[photonlib-py] Begin implementing PhotonPoseEstimator in Python (#1178)

* [photonlib-py] Initial impl of PhotonPoseEstimator

---------

Co-authored-by: Matt <[email protected]>

photon-lib/py/photonlibpy/estimatedRobotPose.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from wpimath.geometry import Pose3d
+
+from .photonTrackedTarget import PhotonTrackedTarget
+
+if TYPE_CHECKING:
+    from .photonPoseEstimator import PoseStrategy
+
+
+@dataclass
+class EstimatedRobotPose:
+    """An estimated pose based on pipeline result"""
+
+    estimatedPose: Pose3d
+    """The estimated pose"""
+
+    timestampSeconds: float
+    """The estimated time the frame used to derive the robot pose was taken"""
+
+    targetsUsed: [PhotonTrackedTarget]
+    """A list of the targets used to compute this pose"""
+
+    strategy: "PoseStrategy"
+    """The strategy actually used to produce this pose"""

photon-lib/py/photonlibpy/photonPoseEstimator.py

@@ -0,0 +1,321 @@
+import enum
+from typing import Optional
+
+import wpilib
+from robotpy_apriltag import AprilTagFieldLayout
+from wpimath.geometry import Transform3d, Pose3d, Pose2d
+
+from .photonPipelineResult import PhotonPipelineResult
+from .photonCamera import PhotonCamera
+from .estimatedRobotPose import EstimatedRobotPose
+
+
+class PoseStrategy(enum.Enum):
+    """
+    Position estimation strategies that can be used by the PhotonPoseEstimator class.
+    """
+
+    LOWEST_AMBIGUITY = enum.auto()
+    """Choose the Pose with the lowest ambiguity."""
+
+    CLOSEST_TO_CAMERA_HEIGHT = enum.auto()
+    """Choose the Pose which is closest to the camera height."""
+
+    CLOSEST_TO_REFERENCE_POSE = enum.auto()
+    """Choose the Pose which is closest to a set Reference position."""
+
+    CLOSEST_TO_LAST_POSE = enum.auto()
+    """Choose the Pose which is closest to the last pose calculated."""
+
+    AVERAGE_BEST_TARGETS = enum.auto()
+    """Return the average of the best target poses using ambiguity as weight."""
+
+    MULTI_TAG_PNP_ON_COPROCESSOR = enum.auto()
+    """
+    Use all visible tags to compute a single pose estimate on coprocessor.
+    This option needs to be enabled on the PhotonVision web UI as well.
+    """
+
+    MULTI_TAG_PNP_ON_RIO = enum.auto()
+    """
+    Use all visible tags to compute a single pose estimate.
+    This runs on the RoboRIO, and can take a lot of time.
+    """
+
+
+class PhotonPoseEstimator:
+    """
+    The PhotonPoseEstimator class filters or combines readings from all the AprilTags visible at a
+    given timestamp on the field to produce a single robot in field pose, using the strategy set
+    below. Example usage can be found in our apriltagExample example project.
+    """
+
+    def __init__(
+        self,
+        fieldTags: AprilTagFieldLayout,
+        strategy: PoseStrategy,
+        camera: PhotonCamera,
+        robotToCamera: Transform3d,
+    ):
+        """Create a new PhotonPoseEstimator.
+
+        :param fieldTags: A WPILib AprilTagFieldLayout linking AprilTag IDs to Pose3d objects
+                           with respect to the FIRST field using the Field Coordinate System.
+                           Note that setting the origin of this layout object will affect the
+                           results from this class.
+        :param strategy: The strategy it should use to determine the best pose.
+        :param camera: PhotonCamera
+        :param robotToCamera: Transform3d from the center of the robot to the camera mount position (i.e.,
+                                robot ➔ camera) in the Robot Coordinate System.
+        """
+        self._fieldTags = fieldTags
+        self._primaryStrategy = strategy
+        self._camera = camera
+        self.robotToCamera = robotToCamera
+
+        self._multiTagFallbackStrategy = PoseStrategy.LOWEST_AMBIGUITY
+        self._reportedErrors: set[int] = set()
+        self._poseCacheTimestampSeconds = -1
+        self._lastPose: Optional[Pose3d] = None
+        self._referencePose: Optional[Pose3d] = None
+
+        # TODO: Implement HAL reporting
+
+    @property
+    def fieldTags(self) -> AprilTagFieldLayout:
+        """Get the AprilTagFieldLayout being used by the PositionEstimator.
+
+        Note: Setting the origin of this layout will affect the results from this class.
+
+        :returns: the AprilTagFieldLayout
+        """
+        return self._fieldTags
+
+    @fieldTags.setter
+    def fieldTags(self, fieldTags: AprilTagFieldLayout):
+        """Set the AprilTagFieldLayout being used by the PositionEstimator.
+
+        Note: Setting the origin of this layout will affect the results from this class.
+
+        :param fieldTags: the AprilTagFieldLayout
+        """
+        self._checkUpdate(self._fieldTags, fieldTags)
+        self._fieldTags = fieldTags
+
+    @property
+    def primaryStrategy(self) -> PoseStrategy:
+        """Get the Position Estimation Strategy being used by the Position Estimator.
+
+        :returns: the strategy
+        """
+        return self._primaryStrategy
+
+    @primaryStrategy.setter
+    def primaryStrategy(self, strategy: PoseStrategy):
+        """Set the Position Estimation Strategy used by the Position Estimator.
+
+        :param strategy: the strategy to set
+        """
+        self._checkUpdate(self._primaryStrategy, strategy)
+        self._primaryStrategy = strategy
+
+    @property
+    def multiTagFallbackStrategy(self) -> PoseStrategy:
+        return self._multiTagFallbackStrategy
+
+    @multiTagFallbackStrategy.setter
+    def multiTagFallbackStrategy(self, strategy: PoseStrategy):
+        """Set the Position Estimation Strategy used in multi-tag mode when only one tag can be seen. Must
+        NOT be MULTI_TAG_PNP
+
+        :param strategy: the strategy to set
+        """
+        self._checkUpdate(self._multiTagFallbackStrategy, strategy)
+        if (
+            strategy is PoseStrategy.MULTI_TAG_PNP_ON_COPROCESSOR
+            or strategy is PoseStrategy.MULTI_TAG_PNP_ON_RIO
+        ):
+            wpilib.reportWarning(
+                "Fallback cannot be set to MULTI_TAG_PNP! Setting to lowest ambiguity",
+                False,
+            )
+            strategy = PoseStrategy.LOWEST_AMBIGUITY
+        self._multiTagFallbackStrategy = strategy
+
+    @property
+    def referencePose(self) -> Pose3d:
+        """Return the reference position that is being used by the estimator.
+
+        :returns: the referencePose
+        """
+        return self._referencePose
+
+    @referencePose.setter
+    def referencePose(self, referencePose: Pose3d | Pose2d):
+        """Update the stored reference pose for use when using the **CLOSEST_TO_REFERENCE_POSE**
+        strategy.
+
+        :param referencePose: the referencePose to set
+        """
+        if isinstance(referencePose, Pose2d):
+            referencePose = Pose3d(referencePose)
+        self._checkUpdate(self._referencePose, referencePose)
+        self._referencePose = referencePose
+
+    @property
+    def lastPose(self) -> Pose3d:
+        return self._lastPose
+
+    @lastPose.setter
+    def lastPose(self, lastPose: Pose3d | Pose2d):
+        """Update the stored last pose. Useful for setting the initial estimate when using the
+        **CLOSEST_TO_LAST_POSE** strategy.
+
+        :param lastPose: the lastPose to set
+        """
+        if isinstance(lastPose, Pose2d):
+            lastPose = Pose3d(lastPose)
+        self._checkUpdate(self._lastPose, lastPose)
+        self._lastPose = lastPose
+
+    def _invalidatePoseCache(self):
+        self._poseCacheTimestampSeconds = -1
+
+    def _checkUpdate(self, oldObj, newObj):
+        if oldObj != newObj and oldObj is not None and oldObj is not newObj:
+            self._invalidatePoseCache()
+
+    def update(
+        self, cameraResult: Optional[PhotonPipelineResult] = None
+    ) -> Optional[EstimatedRobotPose]:
+        """
+        Updates the estimated position of the robot. Returns empty if:
+
+         - The timestamp of the provided pipeline result is the same as in the previous call to
+         ``update()``.
+
+         - No targets were found in the pipeline results.
+
+        :param cameraResult: The latest pipeline result from the camera
+
+        :returns: an :class:`EstimatedRobotPose` with an estimated pose, timestamp, and targets used to
+                   create the estimate.
+        """
+        if not cameraResult:
+            if not self._camera:
+                wpilib.reportError("[PhotonPoseEstimator] Missing camera!", False)
+                return
+            cameraResult = self._camera.getLatestResult()
+
+        if cameraResult.timestampSec < 0:
+            return
+
+        # If the pose cache timestamp was set, and the result is from the same
+        # timestamp, return an
+        # empty result
+        if (
+            self._poseCacheTimestampSeconds > 0
+            and abs(self._poseCacheTimestampSeconds - cameraResult.timestampSec) < 1e-6
+        ):
+            return
+
+        # Remember the timestamp of the current result used
+        self._poseCacheTimestampSeconds = cameraResult.timestampSec
+
+        # If no targets seen, trivial case -- return empty result
+        if not cameraResult.targets:
+            return
+
+        return self._update(cameraResult, self._primaryStrategy)
+
+    def _update(
+        self, cameraResult: PhotonPipelineResult, strat: PoseStrategy
+    ) -> Optional[EstimatedRobotPose]:
+        if strat is PoseStrategy.LOWEST_AMBIGUITY:
+            estimatedPose = self._lowestAmbiguityStrategy(cameraResult)
+        elif strat is PoseStrategy.MULTI_TAG_PNP_ON_COPROCESSOR:
+            estimatedPose = self._multiTagOnCoprocStrategy(cameraResult)
+        else:
+            wpilib.reportError(
+                "[PhotonPoseEstimator] Unknown Position Estimation Strategy!", False
+            )
+            return
+
+        if not estimatedPose:
+            self._lastPose = None
+
+        return estimatedPose
+
+    def _multiTagOnCoprocStrategy(
+        self, result: PhotonPipelineResult
+    ) -> Optional[EstimatedRobotPose]:
+        if result.multiTagResult.estimatedPose.isPresent:
+            best_tf = result.multiTagResult.estimatedPose.best
+            best = (
+                Pose3d()
+                .transformBy(best_tf)  # field-to-camera
+                .relativeTo(self._fieldTags.getOrigin())
+                .transformBy(self.robotToCamera.inverse())  # field-to-robot
+            )
+            return EstimatedRobotPose(
+                best,
+                result.timestampSec,
+                result.targets,
+                PoseStrategy.MULTI_TAG_PNP_ON_COPROCESSOR,
+            )
+        else:
+            return self._update(result, self._multiTagFallbackStrategy)
+
+    def _lowestAmbiguityStrategy(
+        self, result: PhotonPipelineResult
+    ) -> Optional[EstimatedRobotPose]:
+        """
+        Return the estimated position of the robot with the lowest position ambiguity from a List of
+        pipeline results.
+
+        :param result: pipeline result
+
+        :returns: the estimated position of the robot in the FCS and the estimated timestamp of this
+                  estimation.
+        """
+        lowestAmbiguityTarget = None
+
+        lowestAmbiguityScore = 10
+
+        for target in result.targets:
+            targetPoseAmbiguity = target.poseAmbiguity
+
+            # Make sure the target is a Fiducial target.
+            if targetPoseAmbiguity != -1 and targetPoseAmbiguity < lowestAmbiguityScore:
+                lowestAmbiguityScore = targetPoseAmbiguity
+                lowestAmbiguityTarget = target
+
+        # Although there are confirmed to be targets, none of them may be fiducial
+        # targets.
+        if not lowestAmbiguityTarget:
+            return
+
+        targetFiducialId = lowestAmbiguityTarget.fiducialId
+
+        targetPosition = self._fieldTags.getTagPose(targetFiducialId)
+
+        if not targetPosition:
+            self._reportFiducialPoseError(targetFiducialId)
+            return
+
+        return EstimatedRobotPose(
+            targetPosition.transformBy(
+                lowestAmbiguityTarget.getBestCameraToTarget().inverse()
+            ).transformBy(self.robotToCamera.inverse()),
+            result.timestampSec,
+            result.targets,
+            PoseStrategy.LOWEST_AMBIGUITY,
+        )
+
+    def _reportFiducialPoseError(self, fiducialId: int) -> None:
+        if fiducialId not in self._reportedErrors:
+            wpilib.reportError(
+                f"[PhotonPoseEstimator] Tried to get pose of unknown AprilTag: {fiducialId}",
+                False,
+            )
+            self._reportedErrors.add(fiducialId)

photon-lib/py/setup.py

@@ -60,6 +60,7 @@
     install_requires=[
         "wpilib<2025,>=2024.0.0b2",
         "robotpy-wpimath<2025,>=2024.0.0b2",
+        "robotpy-apriltag<2025,>=2024.0.0b2",
         "pyntcore<2025,>=2024.0.0b2",
     ],
     description=descriptionStr,