Merge pull request #1 from illini-robomaster/roger/2023_educational

Roger's 2022 competition season implementation

Aiming/Aim.py

@@ -1,6 +1,75 @@
+import numpy as np
+import config
+import Utils
+
 class Aim:
     def __init__(self):
         pass
 
-    def get_rotation(self, pred):
-        pass
+    def process_one(self, pred_list, enemy_team, depth_map):
+        assert enemy_team in ['blue', 'red']
+        closet_pred, closet_dist = self.get_closet_pred(pred_list, enemy_team, depth_map)
+        if closet_pred is None:
+            return None
+        name, confidence, bbox = closet_pred
+        center_x, center_y, width, height = bbox
+
+        # Get yaw/pitch differences in radians
+        yaw_diff, pitch_diff = self.get_rotation_angle(center_x, center_y)
+
+        calibrated_yaw_diff, calibrated_pitch_diff = self.posterior_calibration(yaw_diff, pitch_diff, closet_dist)
+
+        return {
+            'yaw_diff': calibrated_yaw_diff,
+            'pitch_diff': calibrated_pitch_diff,
+            'center_x': center_x,
+            'center_y': center_y,
+        }
+
+    def posterior_calibration(self, yaw_diff, pitch_diff, distance):
+        """Given a set of naively estimated parameters, return calibrated parameters
+        from a range table?
+
+        Args:
+            yaw_diff (float): yaw difference in radians
+            pitch_diff (float): pitch difference in radians
+            distance (float): distance to target in meters
+
+        Returns:
+            (float, float): calibrated yaw_diff, pitch_diff in radians
+        """
+        if distance >= 2**16:
+            # In this case, distance is a rough estimation from bbox size
+            return (yaw_diff, pitch_diff)
+        else:
+            # In this case, distance comes from D455 stereo estimation
+            # TODO: compute a range table
+            return (yaw_diff, pitch_diff)
+
+    def get_closet_pred(self, pred_list, enemy_team, depth_map):
+        closet_pred = None
+        closet_dist = None # Cloest to camera in z-axis
+        obj_of_interest = [f"armor_{enemy_team}"]
+        for name, conf, bbox in pred_list:
+            # name from C++ string is in bytes; decoding is needed
+            if isinstance(name, bytes):
+                name_str = name.decode('utf-8')
+            else:
+                name_str = name
+            if name_str not in obj_of_interest: continue
+            cur_dist = Utils.estimate_target_depth(bbox, depth_map)
+            if closet_pred is None:
+                closet_pred = (name, conf, bbox)
+                closet_dist = cur_dist
+            else:
+                if cur_dist < closet_dist:
+                    closet_pred = (name, conf, bbox)
+                    closet_dist = cur_dist
+        return closet_pred, closet_dist
+
+    @staticmethod
+    def get_rotation_angle(bbox_center_x, bbox_center_y):
+        yaw_diff = (bbox_center_x - config.IMG_CENTER_X) * (config.RGBD_CAMERA.YAW_FOV_HALF / config.IMG_CENTER_X)
+        pitch_diff = (bbox_center_y - config.IMG_CENTER_Y) * (config.RGBD_CAMERA.PITCH_FOV_HALF / config.IMG_CENTER_Y)
+
+        return yaw_diff, pitch_diff

Camera/d455.py

@@ -0,0 +1,31 @@
+import numpy as np
+import pyrealsense2 as rs
+import Utils
+
+# TODO: make an abstract camera base class
+class D455_camera(object):
+    # Camera parameters from intelrealsense.com/depth-camera/d-455/
+    # TODO: for better accuracy, use individual calibrated intrinsics
+    YAW_FOV_HALF = Utils.deg_to_rad(87) / 2
+    PITCH_FOV_HALF = Utils.deg_to_rad(58) / 2
+
+    def __init__(self, width, height):
+        self.width = width
+        self.height = height
+
+        self.pipeline = rs.pipeline()
+
+        self.config = rs.config()
+        self.config.enable_stream(rs.stream.depth, self.width, self.height, rs.format.z16, 30)
+        self.config.enable_stream(rs.stream.color, self.width, self.height, rs.format.bgr8, 30)
+
+        self.pipeline.start(self.config)
+
+    def get_frame(self):
+        RGBD_frame = self.pipeline.wait_for_frames()
+
+        frame = np.array(RGBD_frame.get_color_frame().get_data()) # HW3
+        depth = np.array(RGBD_frame.get_depth_frame().get_data()) # uint16 0-65535
+        depth[depth == 65535] = 0 # mask ignored label to background
+
+        return (frame, depth)

Communication/communicator.py

@@ -2,19 +2,53 @@
 import serial
 import crc8
 import crc16
-from variables import SERIAL_PORT
-
-serial_port = serial.Serial(
-    port=SERIAL_PORT,
-    baudrate=115200,
-    bytesize=serial.EIGHTBITS,
-    parity=serial.PARITY_NONE,
-    stopbits=serial.STOPBITS_ONE,
-)
+
+import config
+
+# USB is for TTL-only device
+USB_PREFIX = "/dev/ttyUSB"
+# ACM is for st-link/TTL device
+ACM_PREFIX = "/dev/ttyACM"
+
+success_flag = False
+
+for prefix in [USB_PREFIX, ACM_PREFIX]:
+    if success_flag: break
+    for i in range(5):
+        if success_flag: break
+        try:
+            serial_port = serial.Serial(
+                port=prefix + str(i),
+                baudrate=115200,
+                bytesize=serial.EIGHTBITS,
+                parity=serial.PARITY_NONE,
+                stopbits=serial.STOPBITS_ONE,
+            )
+            success_flag = True
+            break # if succeed, break
+        except serial.serialutil.SerialException:
+            serial_port = None
+
 # Wait a second to let the port initialize
 time.sleep(1)
 
-def create_packet(cmd_id, data, seq):
+def create_packet(header, seq_num, yaw_offset, pitch_offset):
+    assert header in [config.SEARCH_TARGET, config.MOVE_YOKE]
+    packet = header
+    assert seq_num >= 0 and seq_num < 2**32 - 1 # uint32
+    packet += seq_num.to_bytes(4, 'big')
+    # YAW/PITCH offset should not be too high
+    assert yaw_offset >= -config.RGBD_CAMERA.YAW_FOV_HALF and yaw_offset <= config.RGBD_CAMERA.YAW_FOV_HALF
+    assert pitch_offset >= -config.RGBD_CAMERA.PITCH_FOV_HALF and pitch_offset <= config.RGBD_CAMERA.PITCH_FOV_HALF
+    discrete_yaw_offset = int(yaw_offset * 100000)
+    discrete_pitch_offset = int(pitch_offset * 100000)
+    packet += (discrete_yaw_offset & 0xFFFFFFFF).to_bytes(4, 'big')
+    packet += (discrete_pitch_offset & 0xFFFFFFFF).to_bytes(4, 'big')
+    # ENDING
+    packet += config.PACK_END
+    return packet
+
+def create_packet_w_crc(cmd_id, data, seq):
     '''
     Args:
         cmd_id: bytes, ID for command to send, see "variables" for different cmd
@@ -38,10 +72,9 @@ def create_packet(cmd_id, data, seq):
     pkt = SOF+data_len+seq+crc_header+cmd_id+data+crc_data
     return pkt
 
-def send_packet(pkt, serial_port):
-    serial_port.write(pkt)
-
+# Testing code
 if __name__ =='__main__':
+    assert serial_port is not None, "No serial device found; check root priviledge and USB devices"
     try:
         cmd_id = b'\xde\xad'
         data = 0xffff*b'\xAA'
@@ -53,11 +86,3 @@ def send_packet(pkt, serial_port):
                 print(data)
     except:
         print("Falied to write")
-
-
-class Communicator:
-    def __init__(self):
-        pass
-
-    def move(self, yaw_offset, pitch_offset):
-        pass