Add python simulation (#1532)

photon-lib/py/photonlibpy/estimation/__init__.py

@@ -0,0 +1,5 @@
+from .cameraTargetRelation import CameraTargetRelation
+from .openCVHelp import OpenCVHelp
+from .rotTrlTransform3d import RotTrlTransform3d
+from .targetModel import TargetModel
+from .visionEstimation import VisionEstimation

photon-lib/py/photonlibpy/estimation/cameraTargetRelation.py

@@ -0,0 +1,25 @@
+import math
+
+from wpimath.geometry import Pose3d, Rotation2d, Transform3d
+from wpimath.units import meters
+
+
+class CameraTargetRelation:
+    def __init__(self, cameraPose: Pose3d, targetPose: Pose3d):
+        self.camPose = cameraPose
+        self.camToTarg = Transform3d(cameraPose, targetPose)
+        self.camToTargDist = self.camToTarg.translation().norm()
+        self.camToTargDistXY: meters = math.hypot(
+            self.camToTarg.translation().X(), self.camToTarg.translation().Y()
+        )
+        self.camToTargYaw = Rotation2d(self.camToTarg.X(), self.camToTarg.Y())
+        self.camToTargPitch = Rotation2d(self.camToTargDistXY, -self.camToTarg.Z())
+        self.camToTargAngle = Rotation2d(
+            math.hypot(self.camToTargYaw.radians(), self.camToTargPitch.radians())
+        )
+        self.targToCam = Transform3d(targetPose, cameraPose)
+        self.targToCamYaw = Rotation2d(self.targToCam.X(), self.targToCam.Y())
+        self.targToCamPitch = Rotation2d(self.camToTargDistXY, -self.targToCam.Z())
+        self.targtoCamAngle = Rotation2d(
+            math.hypot(self.targToCamYaw.radians(), self.targToCamPitch.radians())
+        )

photon-lib/py/photonlibpy/estimation/openCVHelp.py

@@ -0,0 +1,200 @@
+import math
+from typing import Any, Tuple
+
+import cv2 as cv
+import numpy as np
+from wpimath.geometry import Rotation3d, Transform3d, Translation3d
+
+from ..targeting import PnpResult, TargetCorner
+from .rotTrlTransform3d import RotTrlTransform3d
+
+NWU_TO_EDN = Rotation3d(np.array([[0, -1, 0], [0, 0, -1], [1, 0, 0]]))
+EDN_TO_NWU = Rotation3d(np.array([[0, 0, 1], [-1, 0, 0], [0, -1, 0]]))
+
+
+class OpenCVHelp:
+    @staticmethod
+    def getMinAreaRect(points: np.ndarray) -> cv.RotatedRect:
+        return cv.RotatedRect(*cv.minAreaRect(points))
+
+    @staticmethod
+    def translationNWUtoEDN(trl: Translation3d) -> Translation3d:
+        return trl.rotateBy(NWU_TO_EDN)
+
+    @staticmethod
+    def rotationNWUtoEDN(rot: Rotation3d) -> Rotation3d:
+        return -NWU_TO_EDN + (rot + NWU_TO_EDN)
+
+    @staticmethod
+    def translationToTVec(translations: list[Translation3d]) -> np.ndarray:
+        retVal: list[list] = []
+        for translation in translations:
+            trl = OpenCVHelp.translationNWUtoEDN(translation)
+            retVal.append([trl.X(), trl.Y(), trl.Z()])
+        return np.array(
+            retVal,
+            dtype=np.float32,
+        )
+
+    @staticmethod
+    def rotationToRVec(rotation: Rotation3d) -> np.ndarray:
+        retVal: list[np.ndarray] = []
+        rot = OpenCVHelp.rotationNWUtoEDN(rotation)
+        rotVec = rot.getQuaternion().toRotationVector()
+        retVal.append(rotVec)
+        return np.array(
+            retVal,
+            dtype=np.float32,
+        )
+
+    @staticmethod
+    def avgPoint(points: list[Tuple[float, float]]) -> Tuple[float, float]:
+        x = 0.0
+        y = 0.0
+        for p in points:
+            x += p[0]
+            y += p[1]
+        return (x / len(points), y / len(points))
+
+    @staticmethod
+    def pointsToTargetCorners(points: np.ndarray) -> list[TargetCorner]:
+        corners = [TargetCorner(p[0, 0], p[0, 1]) for p in points]
+        return corners
+
+    @staticmethod
+    def cornersToPoints(corners: list[TargetCorner]) -> np.ndarray:
+        points = [[[c.x, c.y]] for c in corners]
+        return np.array(points)
+
+    @staticmethod
+    def projectPoints(
+        cameraMatrix: np.ndarray,
+        distCoeffs: np.ndarray,
+        camRt: RotTrlTransform3d,
+        objectTranslations: list[Translation3d],
+    ) -> np.ndarray:
+        objectPoints = OpenCVHelp.translationToTVec(objectTranslations)
+        rvec = OpenCVHelp.rotationToRVec(camRt.getRotation())
+        tvec = OpenCVHelp.translationToTVec(
+            [
+                camRt.getTranslation(),
+            ]
+        )
+
+        pts, _ = cv.projectPoints(objectPoints, rvec, tvec, cameraMatrix, distCoeffs)
+        return pts
+
+    @staticmethod
+    def reorderCircular(
+        elements: list[Any] | np.ndarray, backwards: bool, shiftStart: int
+    ) -> list[Any]:
+        size = len(elements)
+        reordered = []
+        dir = -1 if backwards else 1
+        for i in range(size):
+            index = (i * dir + shiftStart * dir) % size
+            if index < 0:
+                index += size
+            reordered.append(elements[index])
+        return reordered
+
+    @staticmethod
+    def translationEDNToNWU(trl: Translation3d) -> Translation3d:
+        return trl.rotateBy(EDN_TO_NWU)
+
+    @staticmethod
+    def rotationEDNToNWU(rot: Rotation3d) -> Rotation3d:
+        return -EDN_TO_NWU + (rot + EDN_TO_NWU)
+
+    @staticmethod
+    def tVecToTranslation(tvecInput: np.ndarray) -> Translation3d:
+        return OpenCVHelp.translationEDNToNWU(Translation3d(tvecInput))
+
+    @staticmethod
+    def rVecToRotation(rvecInput: np.ndarray) -> Rotation3d:
+        return OpenCVHelp.rotationEDNToNWU(Rotation3d(rvecInput))
+
+    @staticmethod
+    def solvePNP_Square(
+        cameraMatrix: np.ndarray,
+        distCoeffs: np.ndarray,
+        modelTrls: list[Translation3d],
+        imagePoints: np.ndarray,
+    ) -> PnpResult | None:
+        modelTrls = OpenCVHelp.reorderCircular(modelTrls, True, -1)
+        imagePoints = np.array(OpenCVHelp.reorderCircular(imagePoints, True, -1))
+        objectMat = np.array(OpenCVHelp.translationToTVec(modelTrls))
+
+        alt: Transform3d | None = None
+        for tries in range(2):
+            retval, rvecs, tvecs, reprojectionError = cv.solvePnPGeneric(
+                objectMat,
+                imagePoints,
+                cameraMatrix,
+                distCoeffs,
+                flags=cv.SOLVEPNP_IPPE_SQUARE,
+            )
+
+            best = Transform3d(
+                OpenCVHelp.tVecToTranslation(tvecs[0]),
+                OpenCVHelp.rVecToRotation(rvecs[0]),
+            )
+            if len(tvecs) > 1:
+                alt = Transform3d(
+                    OpenCVHelp.tVecToTranslation(tvecs[1]),
+                    OpenCVHelp.rVecToRotation(rvecs[1]),
+                )
+
+            if not math.isnan(reprojectionError[0, 0]):
+                break
+            else:
+                pt = imagePoints[0]
+                pt[0, 0] -= 0.001
+                pt[0, 1] -= 0.001
+                imagePoints[0] = pt
+
+        if math.isnan(reprojectionError[0, 0]):
+            print("SolvePNP_Square failed!")
+            return None
+
+        if alt:
+            return PnpResult(
+                best=best,
+                bestReprojErr=reprojectionError[0, 0],
+                alt=alt,
+                altReprojErr=reprojectionError[1, 0],
+                ambiguity=reprojectionError[0, 0] / reprojectionError[1, 0],
+            )
+        else:
+            # We have no alternative so set it to best as well
+            return PnpResult(
+                best=best,
+                bestReprojErr=reprojectionError[0],
+                alt=best,
+                altReprojErr=reprojectionError[0],
+            )
+
+    @staticmethod
+    def solvePNP_SQPNP(
+        cameraMatrix: np.ndarray,
+        distCoeffs: np.ndarray,
+        modelTrls: list[Translation3d],
+        imagePoints: np.ndarray,
+    ) -> PnpResult | None:
+        objectMat = np.array(OpenCVHelp.translationToTVec(modelTrls))
+
+        retval, rvecs, tvecs, reprojectionError = cv.solvePnPGeneric(
+            objectMat, imagePoints, cameraMatrix, distCoeffs, flags=cv.SOLVEPNP_SQPNP
+        )
+
+        error = reprojectionError[0, 0]
+        best = Transform3d(
+            OpenCVHelp.tVecToTranslation(tvecs[0]), OpenCVHelp.rVecToRotation(rvecs[0])
+        )
+
+        if math.isnan(error):
+            return None
+
+        # We have no alternative so set it to best as well
+        result = PnpResult(best=best, bestReprojErr=error, alt=best, altReprojErr=error)
+        return result

photon-lib/py/photonlibpy/estimation/rotTrlTransform3d.py

@@ -0,0 +1,32 @@
+from typing import Self
+
+from wpimath.geometry import Pose3d, Rotation3d, Transform3d, Translation3d
+
+
+class RotTrlTransform3d:
+    def __init__(
+        self, rot: Rotation3d = Rotation3d(), trl: Translation3d = Translation3d()
+    ):
+        self.rot = rot
+        self.trl = trl
+
+    def inverse(self) -> Self:
+        invRot = -self.rot
+        invTrl = -(self.trl.rotateBy(invRot))
+        return type(self)(invRot, invTrl)
+
+    def getTransform(self) -> Transform3d:
+        return Transform3d(self.trl, self.rot)
+
+    def getTranslation(self) -> Translation3d:
+        return self.trl
+
+    def getRotation(self) -> Rotation3d:
+        return self.rot
+
+    def apply(self, trlToApply: Translation3d) -> Translation3d:
+        return trlToApply.rotateBy(self.rot) + self.trl
+
+    @classmethod
+    def makeRelativeTo(cls, pose: Pose3d) -> Self:
+        return cls(pose.rotation(), pose.translation()).inverse()

photon-lib/py/photonlibpy/estimation/targetModel.py

@@ -0,0 +1,137 @@
+import math
+from typing import List, Self
+
+from wpimath.geometry import Pose3d, Rotation2d, Rotation3d, Translation3d
+from wpimath.units import meters
+
+from . import RotTrlTransform3d
+
+
+class TargetModel:
+    def __init__(
+        self,
+        *,
+        width: meters | None = None,
+        height: meters | None = None,
+        length: meters | None = None,
+        diameter: meters | None = None,
+        verts: List[Translation3d] | None = None
+    ):
+
+        if (
+            width is not None
+            and height is not None
+            and length is None
+            and diameter is None
+            and verts is None
+        ):
+            self.isPlanar = True
+            self.isSpherical = False
+            self.vertices = [
+                Translation3d(0.0, -width / 2.0, -height / 2.0),
+                Translation3d(0.0, width / 2.0, -height / 2.0),
+                Translation3d(0.0, width / 2.0, height / 2.0),
+                Translation3d(0.0, -width / 2.0, height / 2.0),
+            ]
+
+            return
+
+        elif (
+            length is not None
+            and width is not None
+            and height is not None
+            and diameter is None
+            and verts is None
+        ):
+            verts = [
+                Translation3d(length / 2.0, -width / 2.0, -height / 2.0),
+                Translation3d(length / 2.0, width / 2.0, -height / 2.0),
+                Translation3d(length / 2.0, width / 2.0, height / 2.0),
+                Translation3d(length / 2.0, -width / 2.0, height / 2.0),
+                Translation3d(-length / 2.0, -width / 2.0, height / 2.0),
+                Translation3d(-length / 2.0, width / 2.0, height / 2.0),
+                Translation3d(-length / 2.0, width / 2.0, -height / 2.0),
+                Translation3d(-length / 2.0, -width / 2.0, -height / 2.0),
+            ]
+            # Handle the rest of this in the "default" case
+        elif (
+            diameter is not None
+            and width is None
+            and height is None
+            and length is None
+            and verts is None
+        ):
+            self.isPlanar = False
+            self.isSpherical = True
+            self.vertices = [
+                Translation3d(0.0, -diameter / 2.0, 0.0),
+                Translation3d(0.0, 0.0, -diameter / 2.0),
+                Translation3d(0.0, diameter / 2.0, 0.0),
+                Translation3d(0.0, 0.0, diameter / 2.0),
+            ]
+            return
+        elif (
+            verts is not None
+            and width is None
+            and height is None
+            and length is None
+            and diameter is None
+        ):
+            # Handle this in the "default" case
+            pass
+        else:
+            raise Exception("Not a valid overload")
+
+        # TODO maybe remove this if there is a better/preferred way
+        # make the python type checking gods happy
+        assert verts is not None
+
+        self.isSpherical = False
+        if len(verts) <= 2:
+            self.vertices: List[Translation3d] = []
+            self.isPlanar = False
+        else:
+            cornersPlaner = True
+            for corner in verts:
+                if abs(corner.X() < 1e-4):
+                    cornersPlaner = False
+            self.isPlanar = cornersPlaner
+
+        self.vertices = verts
+
+    def getFieldVertices(self, targetPose: Pose3d) -> List[Translation3d]:
+        basisChange = RotTrlTransform3d(targetPose.rotation(), targetPose.translation())
+
+        retVal = []
+
+        for vert in self.vertices:
+            retVal.append(basisChange.apply(vert))
+
+        return retVal
+
+    @classmethod
+    def getOrientedPose(cls, tgtTrl: Translation3d, cameraTrl: Translation3d):
+        relCam = cameraTrl - tgtTrl
+        orientToCam = Rotation3d(
+            0.0,
+            Rotation2d(math.hypot(relCam.X(), relCam.Y()), relCam.Z()).radians(),
+            Rotation2d(relCam.X(), relCam.Y()).radians(),
+        )
+        return Pose3d(tgtTrl, orientToCam)
+
+    def getVertices(self) -> List[Translation3d]:
+        return self.vertices
+
+    def getIsPlanar(self) -> bool:
+        return self.isPlanar
+
+    def getIsSpherical(self) -> bool:
+        return self.isSpherical
+
+    @classmethod
+    def AprilTag36h11(cls) -> Self:
+        return cls(width=6.5 * 0.0254, height=6.5 * 0.0254)
+
+    @classmethod
+    def AprilTag16h5(cls) -> Self:
+        return cls(width=6.0 * 0.0254, height=6.0 * 0.0254)

photon-lib/py/photonlibpy/estimation/visionEstimation.py

@@ -0,0 +1,91 @@
+import numpy as np
+from robotpy_apriltag import AprilTag, AprilTagFieldLayout
+from wpimath.geometry import Pose3d, Transform3d, Translation3d
+
+from ..targeting import PhotonTrackedTarget, PnpResult, TargetCorner
+from . import OpenCVHelp, TargetModel
+
+
+class VisionEstimation:
+    @staticmethod
+    def getVisibleLayoutTags(
+        visTags: list[PhotonTrackedTarget], layout: AprilTagFieldLayout
+    ) -> list[AprilTag]:
+        retVal: list[AprilTag] = []
+        for tag in visTags:
+            id = tag.getFiducialId()
+            maybePose = layout.getTagPose(id)
+            if maybePose:
+                tag = AprilTag()
+                tag.ID = id
+                tag.pose = maybePose
+                retVal.append(tag)
+        return retVal
+
+    @staticmethod
+    def estimateCamPosePNP(
+        cameraMatrix: np.ndarray,
+        distCoeffs: np.ndarray,
+        visTags: list[PhotonTrackedTarget],
+        layout: AprilTagFieldLayout,
+        tagModel: TargetModel,
+    ) -> PnpResult | None:
+        if len(visTags) == 0:
+            return None
+
+        corners: list[TargetCorner] = []
+        knownTags: list[AprilTag] = []
+
+        for tgt in visTags:
+            id = tgt.getFiducialId()
+            maybePose = layout.getTagPose(id)
+            if maybePose:
+                tag = AprilTag()
+                tag.ID = id
+                tag.pose = maybePose
+                knownTags.append(tag)
+                currentCorners = tgt.getDetectedCorners()
+                if currentCorners:
+                    corners += currentCorners
+
+        if len(knownTags) == 0 or len(corners) == 0 or len(corners) % 4 != 0:
+            return None
+
+        points = OpenCVHelp.cornersToPoints(corners)
+
+        if len(knownTags) == 1:
+            camToTag = OpenCVHelp.solvePNP_Square(
+                cameraMatrix, distCoeffs, tagModel.getVertices(), points
+            )
+            if not camToTag:
+                return None
+
+            bestPose = knownTags[0].pose.transformBy(camToTag.best.inverse())
+            altPose = Pose3d()
+            if camToTag.ambiguity != 0:
+                altPose = knownTags[0].pose.transformBy(camToTag.alt.inverse())
+
+            o = Pose3d()
+            result = PnpResult(
+                best=Transform3d(o, bestPose),
+                alt=Transform3d(o, altPose),
+                ambiguity=camToTag.ambiguity,
+                bestReprojErr=camToTag.bestReprojErr,
+                altReprojErr=camToTag.altReprojErr,
+            )
+            return result
+        else:
+            objectTrls: list[Translation3d] = []
+            for tag in knownTags:
+                verts = tagModel.getFieldVertices(tag.pose)
+                objectTrls += verts
+
+            ret = OpenCVHelp.solvePNP_SQPNP(
+                cameraMatrix, distCoeffs, objectTrls, points
+            )
+            if ret:
+                # Invert best/alt transforms
+                ret.best = ret.best.inverse()
+                ret.alt = ret.alt.inverse()
+
+            return ret