Merge pull request #538 from truher/main

more pose estimation

raspberry_pi/app/config/camera_config.py

@@ -0,0 +1,29 @@
+"""Intrinstic and distortion config of a camera."""
+
+# pylint: disable=E1101,R0903
+
+import gtsam
+import numpy as np
+
+from app.config.identity import Identity
+
+
+class CameraConfig:
+    def __init__(self, identity: Identity) -> None:
+        match identity:
+            case Identity.UNKNOWN:
+                self.camera_offset = gtsam.Pose3(
+                    gtsam.Rot3(),
+                    np.array([0, 0, 1]),
+                )
+                self.calib = gtsam.Cal3DS2(
+                    200.0, 200.0, 0.0, 200.0, 200.0, -0.2, 0.1, 0.0, 0.0
+                )
+            case _:
+                self.camera_offset = gtsam.Pose3(
+                    gtsam.Rot3(),
+                    np.array([0, 0, 1]),
+                )
+                self.calib = gtsam.Cal3DS2(
+                    200.0, 200.0, 0.0, 200.0, 200.0, -0.2, 0.1, 0.0, 0.0
+                )

raspberry_pi/app/localization/note_detector.py

@@ -66,7 +66,6 @@ def analyze(self, req: Request) -> None:
             floodfill = img_range.copy()
             mask = np.zeros((self.height + 2, self.width + 2), np.uint8)
             cv2.floodFill(floodfill, mask, [0, 0], [255])
-
 
             floodfill_inv = cv2.bitwise_not(floodfill)
             img_floodfill = cv2.bitwise_or(img_range, floodfill_inv)
@@ -117,5 +116,5 @@ def analyze(self, req: Request) -> None:
             fps = req.fps()
             self.display.text(img_bgr, f"FPS {fps:2.0f}", (5, 65))
             self.display.text(img_bgr, f"delay (ms) {delay_us/1000:2.0f}", (5, 105))
-            #self.display.put(img_range)
+            # self.display.put(img_range)
             self.display.put(img_bgr)

raspberry_pi/app/network/network_protocol.py

@@ -2,10 +2,10 @@
 
 # pylint: disable=R0902,R0903,W0212,W2301
 
-
 import dataclasses
 from typing import Protocol
 
+import numpy as np
 from wpimath.geometry import Rotation3d, Transform3d
 from wpiutil import wpistruct
 
@@ -27,9 +27,9 @@ class Blip25:
     so i can't use an array here. it's ok because a tag only
     ever has exactly 4 corners; if one is occluded then the
     entire tag is unseen.
+    There's no tag id here because tag id is in the NT key.
     """
 
-    id: int
     llx: float  # lower left
     lly: float
     lrx: float  # lower right
@@ -39,6 +39,21 @@ class Blip25:
     ulx: float  # upper left
     uly: float
 
+    def measurement(self) -> np.ndarray:
+        """Concatenated corners, for GTSAM."""
+        return np.array(
+            [
+                self.llx,
+                self.lly,
+                self.lrx,
+                self.lry,
+                self.urx,
+                self.ury,
+                self.ulx,
+                self.uly,
+            ]
+        )
+
 
 class DoubleSender(Protocol):
     def send(self, val: float, delay_us: int) -> None: ...
@@ -57,7 +72,7 @@ def send(self, val: list[Blip25], delay_us: int) -> None: ...
 
 
 class Blip25Receiver(Protocol):
-    def get(self) -> list[tuple[int, list[Blip25]]]:
+    def get(self) -> list[tuple[int, int, Blip25]]:
         """Returns the list of tuples (timetamp, list[blip]) seen in a single frame"""
         ...
 

raspberry_pi/app/network/real_network.py

@@ -3,21 +3,17 @@
 # pylint: disable=R0902,R0903,W0212
 
 # github workflow crashes in ntcore, bleah
+from typing import cast
+
 import ntcore
 from typing_extensions import override
 from wpimath.geometry import Rotation3d
+from wpiutil import wpistruct
 
 from app.config.identity import Identity
-from app.network.network_protocol import (
-    Blip24,
-    Blip25,
-    Blip25Receiver,
-    Blip25Sender,
-    BlipSender,
-    DoubleSender,
-    Network,
-    NoteSender,
-)
+from app.network.network_protocol import (Blip24, Blip25, Blip25Receiver,
+                                          Blip25Sender, BlipSender,
+                                          DoubleSender, Network, NoteSender)
 
 
 class RealDoubleSender(DoubleSender):
@@ -58,17 +54,39 @@ def send(self, val: list[Blip25], delay_us: int) -> None:
 
 
 class RealBlip25Receiver(Blip25Receiver):
-    def __init__(self, sub: ntcore.StructArraySubscriber) -> None:
-        self.sub = sub
+    def __init__(
+        self,
+        name: str,
+        inst: ntcore.NetworkTableInstance,
+    ) -> None:
+        print(name)
+        self.name = name
+        self.poller = ntcore.NetworkTableListenerPoller(inst)
+        # need to hang on to this reference :-(
+        self.msub = ntcore.MultiSubscriber(inst, [name])
+        self.poller.addListener(self.msub, ntcore.EventFlags.kValueAll)
+        # self.poller.addListener([""], ntcore.EventFlags.kValueAll)
 
     @override
-    def get(self) -> list[tuple[int, list[Blip25]]]:
-        recv = self.sub.readQueue()
-        result: list[tuple[int, list[Blip25]]] = []
-        for item in recv:
-            server_time_us = item.serverTime
-            value_list = item.value
-            result.append((server_time_us, value_list))
+    def get(self) -> list[tuple[int, int, Blip25]]:
+        """(timestamp, tag id, blip)"""
+        result: list[tuple[int, int, Blip25]] = []
+        print("get")
+        # see NotePosition24ArrayListener for example
+        for e in self.poller.readQueue():
+            print("in queue")
+            ve = cast(ntcore.ValueEventData, e.data)
+            v = ve.value
+            name = ve.topic.getName()
+            # TODO: redo the key scheme
+            tag_id = int(name.split("/")[1])
+            print(name)
+            server_time_us = v.server_time()
+            b = v.getRaw()
+            # value = cast(Blip25, v.value)
+            value = wpistruct.unpack(Blip25, b)
+            print(value)
+            result.append((server_time_us, tag_id, value))
         return result
 
 
@@ -113,9 +131,7 @@ def get_blip25_sender(self, name: str) -> RealBlip25Sender:
 
     @override
     def get_blip25_receiver(self, name: str) -> Blip25Receiver:
-        return RealBlip25Receiver(
-            self._inst.getStructArrayTopic(name, Blip25).subscribe([])
-        )
+        return RealBlip25Receiver(name, self._inst)
 
     @override
     def flush(self) -> None:

raspberry_pi/app/pose_estimator/estimate.py

@@ -11,7 +11,7 @@
 from gtsam import noiseModel  # type:ignore
 from gtsam.noiseModel import Base as SharedNoiseModel  # type:ignore
 from gtsam.symbol_shorthand import C, K, X  # type:ignore
-from wpimath.geometry import Pose2d, Rotation2d, Translation2d, Twist2d
+from wpimath.geometry import Rotation2d, Translation2d, Twist2d
 
 import app.pose_estimator.factors.accelerometer as accelerometer
 import app.pose_estimator.factors.apriltag_calibrate as apriltag_calibrate
@@ -22,14 +22,13 @@
 from app.pose_estimator.drive_util import DriveUtil
 from app.pose_estimator.swerve_drive_kinematics import SwerveDriveKinematics100
 from app.pose_estimator.swerve_module_delta import SwerveModuleDelta
-from app.pose_estimator.swerve_module_position import (
-    OptionalRotation2d,
-    SwerveModulePosition100,
-)
+from app.pose_estimator.swerve_module_position import (OptionalRotation2d,
+                                                       SwerveModulePosition100)
 
 # TODO: real noise estimates.
 ODOMETRY_NOISE = noiseModel.Diagonal.Sigmas(np.array([0.01, 0.01, 0.01]))
-PRIOR_NOISE = noiseModel.Diagonal.Sigmas(np.array([0.3, 0.3, 0.1]))
+# prior uncertainty is larger than field, i.e. "no idea"
+PRIOR_NOISE = noiseModel.Diagonal.Sigmas(np.array([16, 8, 6]))
 ACCELEROMETER_NOISE = noiseModel.Diagonal.Sigmas(np.array([0.1, 0.1]))
 GYRO_NOISE = noiseModel.Diagonal.Sigmas(np.array([0.05]))
 
@@ -83,13 +82,8 @@ def __init__(self) -> None:
             SwerveModulePosition100(0, OptionalRotation2d(True, Rotation2d(0))),
         ]
 
-    def init(self, initial_pose: Pose2d) -> None:
-        """Add a state at zero"""
-        prior_mean = gtsam.Pose2(
-            initial_pose.X(), initial_pose.Y(), initial_pose.rotation().radians()
-        )
-        self.add_state(0, prior_mean)
-        self.prior(0, prior_mean, PRIOR_NOISE)
+    def init(self) -> None:
+        """No longer adds a state at zero, if you want that, do it."""
         # there is just one camera factor
         self.new_values.insert(K(0), CAL)
         self.new_timestamps[K(0)] = 0
@@ -241,7 +235,8 @@ def apriltag_for_smoothing_batch(
         calib: gtsam.Cal3DS2,
     ) -> None:
         """landmarks: list of 3d points
-        measured: concatenated px measurements"""
+        measured: concatenated px measurements
+        TODO: flatten landmarks"""
         noise = noiseModel.Diagonal.Sigmas(
             np.concatenate(
                 [[1, 1] for _ in landmarks])

raspberry_pi/app/pose_estimator/field_map.py

@@ -0,0 +1,34 @@
+"""Tag locations."""
+
+import math
+
+import numpy as np
+
+TAG_SIZE_M = 0.1651
+HALF = TAG_SIZE_M / 2
+
+
+class FieldMap:
+    def __init__(self) -> None:
+        # key = tag number
+        # value = list of arrays, each array is (x, y)
+        # TODO: make this a flat array with 8 numbers instead.
+        self.tags: dict[int, list[np.ndarray]] = {}
+        self.tags[0] = FieldMap.make_tag(4, 0, 1, 0)
+        self.tags[1] = FieldMap.make_tag(2, 2, 1, math.pi / 2)
+
+    def get(self, tag_id: int) -> list[np.ndarray]:
+        """list of corners"""
+        return self.tags[tag_id]
+
+    @staticmethod
+    def make_tag(x, y, z, yaw) -> list[np.ndarray]:
+        """yaw: 'into the page' orientation.
+        pitch and roll are always zero"""
+        s = HALF * math.sin(yaw)
+        c = HALF * math.cos(yaw)
+        ll = np.array([x - s, y + c, z - HALF])
+        lr = np.array([x + s, y - c, z - HALF])
+        ur = np.array([x + s, y - c, z + HALF])
+        ul = np.array([x - s, y + c, z + HALF])
+        return [ll, lr, ur, ul]

raspberry_pi/app/pose_estimator/nt_estimate.py

@@ -1,12 +1,52 @@
 """Read measurements from Network Tables, run the
 smoother, and publish the results on Network Tables."""
 
+# pylint: disable=C0301,E0611,E1101,R0903,R0914
 
+# TODO: remove gtsam
+import gtsam
+import numpy as np
+from gtsam import noiseModel  # type:ignore
+from wpimath.geometry import Pose2d
+
+from app.camera.camera_protocol import Camera
+from app.config.camera_config import CameraConfig
+from app.config.identity import Identity
 from app.network.network_protocol import Network
+from app.pose_estimator.estimate import Estimate
+from app.pose_estimator.field_map import FieldMap
+
+PRIOR_NOISE = noiseModel.Diagonal.Sigmas(np.array([0.3, 0.3, 0.1]))
 
 
 class NTEstimate:
-    def __init__(self, net: Network) -> None:
+    def __init__(self, field_map: FieldMap, net: Network) -> None:
+        self.field_map = field_map
         self.net = net
+        self.blip_receiver = net.get_blip25_receiver("foo")
+        self.est = Estimate()
+        # current estimate, used for initial value for next time
+        # TODO: remove gtsam types
+        self.state = gtsam.Pose2()
+        self.est.init()
+        prior_mean = gtsam.Pose2(0, 0, 0)
+        self.est.add_state(0, prior_mean)
+        self.est.prior(0, prior_mean, PRIOR_NOISE)
+
+    def step(self) -> None:
+        """Collect any pending measurements from
+        the network and add them to the sim."""
+        # TODO: read the camera identity from the blip
+        cam = CameraConfig(Identity.UNKNOWN)
+        sights = self.blip_receiver.get()
+        print("sights ", sights)
+        for sight in sights:
+            timestamp_us = sight[0]
+            tag_id = sight[1]
+            blip = sight[2]
+            pixels = blip.measurement()
+            corners = self.field_map.get(tag_id)
+            self.est.apriltag_for_smoothing_batch(
+                corners, pixels, timestamp_us, cam.camera_offset, cam.calib
+            )
 
-

raspberry_pi/tests/pose_estimator/circle_simulator.py

@@ -22,18 +22,23 @@
 
 import numpy as np
 from gtsam import Cal3DS2  # includes distortion
-from gtsam import (PinholeCameraCal3DS2, Point2, Point3, Pose2,  # type:ignore
-                   Pose3, Rot3)
+from gtsam import Point2  # type:ignore
+from gtsam import Point3  # type:ignore
+from gtsam import Pose2  # type:ignore
+from gtsam import PinholeCameraCal3DS2, Pose3, Rot3
 from wpimath.geometry import Pose2d, Rotation2d, Translation2d
 
+from app.pose_estimator.field_map import FieldMap
 from app.pose_estimator.swerve_drive_kinematics import SwerveDriveKinematics100
-from app.pose_estimator.swerve_module_position import (OptionalRotation2d,
-                                                       SwerveModulePosition100)
+from app.pose_estimator.swerve_module_position import (
+    OptionalRotation2d,
+    SwerveModulePosition100,
+)
 
-TAG_SIZE_M: float = 0.1651
-TAG_X: float = 4
-TAG_Y: float = 0
-TAG_Z: float = 1
+# TAG_SIZE_M: float = 0.1651
+# TAG_X: float = 4
+# TAG_Y: float = 0
+# TAG_Z: float = 1
 PATH_CENTER_X_M = 1
 PATH_CENTER_Y_M = 0
 PATH_RADIUS_M = 1
@@ -48,7 +53,8 @@
 
 
 class CircleSimulator:
-    def __init__(self) -> None:
+    def __init__(self, field_map: FieldMap) -> None:
+        self.field_map = field_map
         # TODO: actual wheelbase etc
         self.kinematics = SwerveDriveKinematics100(
             [
@@ -87,11 +93,12 @@ def __init__(self) -> None:
             SwerveModulePosition100(0, OptionalRotation2d(True, Rotation2d(0))),
         ]
         # constant landmark points
-        self.l0 = np.array([TAG_X, TAG_Y + (TAG_SIZE_M / 2), TAG_Z - (TAG_SIZE_M / 2)])
-        self.l1 = np.array([TAG_X, TAG_Y - (TAG_SIZE_M / 2), TAG_Z - (TAG_SIZE_M / 2)])
-        self.l2 = np.array([TAG_X, TAG_Y - (TAG_SIZE_M / 2), TAG_Z + (TAG_SIZE_M / 2)])
-        self.l3 = np.array([TAG_X, TAG_Y + (TAG_SIZE_M / 2), TAG_Z + (TAG_SIZE_M / 2)])
-        self.landmarks = [self.l0,self.l1,self.l2,self.l3]
+        tag = self.field_map.get(0)
+        self.l0 = tag[0]
+        self.l1 = tag[1]
+        self.l2 = tag[2]
+        self.l3 = tag[3]
+        self.landmarks = [self.l0, self.l1, self.l2, self.l3]
         self.camera_offset = Pose3(Rot3(), np.array([0, 0, 1]))
         self.calib = Cal3DS2(200.0, 200.0, 0.0, 200.0, 200.0, -0.2, 0.1, 0.0, 0.0)