ReefDisplay/main.py at master · wilsonernst/ReefDisplay · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import cv2
import numpy as np
import ntcore
import time
import display


def yaw_pitch_to_pixel(yaw, pitch, img_width = 1920, img_height = 1200, hfov = 104.83/2.0, vfov = 78.13/2.0):
    """
    Converts yaw and pitch angles to pixel coordinates in an image.

    :param yaw: Yaw angle (degrees), provided by PhotonVision
    :param pitch: Pitch angle (degrees), provided by PhotonVision
    :param img_width: Image width in pixels
    :param img_height: Image height in pixels
    :param hfov: Horizontal field of view (degrees)
    :param vfov: Vertical field of view (degrees)
    :return: (x, y) pixel coordinates in the image frame
    """
    # Convert angles to normalized coordinates (-1 to 1 range)
    x_norm = np.tan(np.radians(yaw)) / np.tan(np.radians(hfov / 2))
    y_norm = np.tan(np.radians(pitch)) / np.tan(np.radians(vfov / 2))

    # Convert to pixel coordinates
    x_pixel = (x_norm + 1) * (img_width / 2)
    y_pixel = (1 - y_norm) * (img_height / 2)  # Inverting y since pixels go downwards

    # Clamp values to image bounds
    x_pixel = np.clip(x_pixel, 0, img_width - 1)
    y_pixel = np.clip(y_pixel, 0, img_height - 1)

    return int(x_pixel), int(y_pixel)

# CONFIGURATION
ORANGE_PI_IP = "10.60.36.89"

TEAM_NUMBER = 6036  # Example team number, change this!
ROBORIO_IP = f"10.{TEAM_NUMBER}.2"

PHOTONVISION_STREAM_URL = f"http://{ORANGE_PI_IP}:1183/stream.mjpg"
cap = cv2.VideoCapture(PHOTONVISION_STREAM_URL)

if not cap.isOpened():
    print("Unable to open video stream")
else:
    print("Video stream opened")

inst = ntcore.NetworkTableInstance.getDefault()
inst.startClient4("example client")
inst.setServerTeam(6036)

centers_table = inst.getDoubleArrayTopic("Centers")
centers_sub = centers_table.subscribe([0])

boxes_table = inst.getDoubleArrayTopic("BoundingBoxes")
boxes_sub = boxes_table.subscribe([0])


while True:
    ret, frame = cap.read()
    if not ret:
        print("no ret")
        continue

################CENTERS#####################

    # centers = centers_sub.get()
    # if centers == [0]:
    #     print("no centers")
    #     #continue

    #for i in range(0, len(centers), 2):
       # yaw, pitch = np.degrees(centers[i]), np.degrees(centers[i+1])
       # x, y = yaw_pitch_to_pixel(yaw, pitch, frame.shape[1], frame.shape[0])
      #  cv2.circle(frame, (x, y), 9, (0, 0, 255), -1)  # Draw green dot

##############BOUNDING BOXES##################

    boxes = boxes_sub.get()
    if boxes == [0]:
        print("no boxes")
        continue

    cv2.rectangle(frame, (int(boxes[0]), int(boxes[3])), (int(boxes[2]), int(boxes[6])), (255, 0, 0), 2)
    cv2.line(frame, (int(boxes[1]), int(boxes[3])), (int(boxes[1]), int(boxes[6])),(255, 0, 0), 2)
    cv2.line(frame, (int(boxes[0]), int(boxes[4])), (int(boxes[2]), int(boxes[4])),(255, 0, 0), 2)
    cv2.line(frame, (int(boxes[0]), int(boxes[5])), (int(boxes[2]), int(boxes[5])),(255, 0, 0), 2)

    frame_resized = cv2.resize(frame, (frame.shape[1] // 2, frame.shape[0] // 2))

    cv2.imshow("frame", frame_resized)

    # display.create_rectangles()

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()