Merge pull request #342 from sourcebots/vision-backport

Backport marker object improvements from sr-robot3

sbot/marker.py

@@ -39,34 +39,56 @@ class Coordinates(NamedTuple):
     z: float
 
 
+class Position(NamedTuple):
+    """
+    Position of a marker in space from the camera's perspective.
+
+    :param distance:          Distance from the camera to the marker, in millimetres.
+    :param horizontal_angle:  Horizontal angle from the camera to the marker, in radians.
+                              Ranges from -pi to pi, with positive values indicating
+                              markers to the right of the camera. Directly in front
+                              of the camera is 0 rad.
+    :param vertical_angle:    Vertical angle from the camera to the marker, in radians.
+                              Ranges from -pi to pi, with positive values indicating
+                              markers above the camera. Directly in front of the camera
+                              is 0 rad.
+    """
+
+    distance: float
+    horizontal_angle: float
+    vertical_angle: float
+
+
 PixelCorners = Tuple[PixelCoordinates, PixelCoordinates, PixelCoordinates, PixelCoordinates]
 
 
 class Marker(NamedTuple):
     """
     Wrapper of a marker detection with axis and rotation calculated.
+
+    :param id: The ID of the detected marker
+    :param size: The physical size of the marker in millimeters
+    :param pixel_corners: A tuple of the PixelCoordinates of the marker's corners in the frame
+    :param pixel_centre: The PixelCoordinates of the marker's centre in the frame
+    :param position: Position information of the marker relative to the camera
+    :param orientation: Orientation information of the marker
     """
 
     id: int
     size: int
     pixel_corners: PixelCorners
     pixel_centre: PixelCoordinates
 
-    # The '2D' distance across the floor
-    distance: float = 0
-    # In radians, increasing clockwise
-    azimuth: float = 0
-    # In radians, increasing upwards
-    elevation: float = 0
-
+    position: Position = Position(0, 0, 0)
     orientation: Orientation = Orientation(0, 0, 0)
 
     @classmethod
     def from_april_vision_marker(cls, marker: AprilMarker) -> 'Marker':
+        """Generate a marker object using the data from april_vision's Marker object."""
         if marker.rvec is None or marker.tvec is None:
             raise ValueError("Marker lacks pose information")
 
-        _cartesian = cls._standardise_tvec(marker.tvec)
+        _cartesian = cls._standardise_tvec(marker.tvec.flatten())
 
         return cls(
             id=marker.id,
@@ -76,9 +98,11 @@ def from_april_vision_marker(cls, marker: AprilMarker) -> 'Marker':
                 tuple(PixelCoordinates(*corner) for corner in marker.pixel_corners)),
             pixel_centre=PixelCoordinates(*marker.pixel_centre),
 
-            distance=int(hypot(*_cartesian) * 1000),
-            azimuth=atan2(-_cartesian.y, _cartesian.x),
-            elevation=atan2(_cartesian.z, _cartesian.x),
+            position=Position(
+                distance=int(hypot(*_cartesian) * 1000),
+                horizontal_angle=atan2(-_cartesian.y, _cartesian.x),
+                vertical_angle=atan2(_cartesian.z, _cartesian.x),
+            ),
 
             orientation=Orientation.from_rvec_matrix(marker.rvec),
         )
@@ -97,7 +121,7 @@ def _standardise_tvec(tvec: NDArray) -> Coordinates:
 
     def __repr__(self) -> str:
         return (
-            f"<{self.__class__.__name__} id={self.id} distance={self.distance:.0f}mm "
-            f"bearing={self.azimuth:.2f}rad elevation={self.elevation:.2f}rad "
-            f"size={self.size}mm>"
+            f"<{self.__class__.__name__} id={self.id} distance={self.position.distance:.0f}mm "
+            f"horizontal_angle={self.position.horizontal_angle:.2f}rad "
+            f"vertical_angle={self.position.vertical_angle:.2f}rad size={self.size}mm>"
         )