Allowing External Angle Sources in Pose Estimation

photon-lib/src/main/java/org/photonvision/PhotonPoseEstimator.java

@@ -27,13 +27,13 @@
 import edu.wpi.first.apriltag.AprilTagFieldLayout;
 import edu.wpi.first.hal.FRCNetComm.tResourceType;
 import edu.wpi.first.hal.HAL;
+import edu.wpi.first.math.MathUtil;
 import edu.wpi.first.math.Matrix;
 import edu.wpi.first.math.Pair;
-import edu.wpi.first.math.geometry.Pose2d;
-import edu.wpi.first.math.geometry.Pose3d;
-import edu.wpi.first.math.geometry.Rotation3d;
-import edu.wpi.first.math.geometry.Transform3d;
-import edu.wpi.first.math.geometry.Translation3d;
+import edu.wpi.first.math.controller.PIDController;
+import edu.wpi.first.math.geometry.*;
+import edu.wpi.first.math.interpolation.Interpolator;
+import edu.wpi.first.math.interpolation.TimeInterpolatableBuffer;
 import edu.wpi.first.math.numbers.N1;
 import edu.wpi.first.math.numbers.N3;
 import edu.wpi.first.math.numbers.N8;
@@ -43,10 +43,12 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import org.opencv.core.Point;
 import org.photonvision.estimation.TargetModel;
 import org.photonvision.estimation.VisionEstimation;
 import org.photonvision.targeting.PhotonPipelineResult;
 import org.photonvision.targeting.PhotonTrackedTarget;
+import org.photonvision.targeting.TargetCorner;
 
 /**
  * The PhotonPoseEstimator class filters or combines readings from all the AprilTags visible at a
@@ -83,7 +85,14 @@ public enum PoseStrategy {
          * Use all visible tags to compute a single pose estimate. This runs on the RoboRIO, and can
          * take a lot of time.
          */
-        MULTI_TAG_PNP_ON_RIO
+        MULTI_TAG_PNP_ON_RIO,
+
+        /**
+         * Use the best visible tag to compute a single pose estimate. This runs on the RoboRIO in order
+         * to access the robot's yaw heading, and MUST have setReferencePose called every frame so
+         * heading data is up-to-date.
+         */
+        IMU_SOLVE_BEST_TAG
     }
 
     private AprilTagFieldLayout fieldTags;
@@ -97,6 +106,7 @@ public enum PoseStrategy {
     protected double poseCacheTimestampSeconds = -1;
     private final Set<Integer> reportedErrors = new HashSet<>();
 
+
     /**
      * Create a new PhotonPoseEstimator.
      *
@@ -259,6 +269,7 @@ public void setLastPose(Pose2d lastPose) {
         setLastPose(new Pose3d(lastPose));
     }
 
+
     /**
      * @return The current transform from the center of the robot to the camera mount position
      */
@@ -380,6 +391,9 @@ private Optional<EstimatedRobotPose> update(
             case MULTI_TAG_PNP_ON_COPROCESSOR:
                 estimatedPose = multiTagOnCoprocStrategy(cameraResult);
                 break;
+            case IMU_SOLVE_BEST_TAG:
+                estimatedPose = imuTagOnRioStrategy(cameraResult, cameraMatrix);
+                break;
             default:
                 DriverStation.reportError(
                         "[PhotonPoseEstimator] Unknown Position Estimation Strategy!", false);
@@ -413,6 +427,71 @@ private Optional<EstimatedRobotPose> multiTagOnCoprocStrategy(PhotonPipelineResu
         }
     }
 
+    private Optional<EstimatedRobotPose> imuTagOnRioStrategy(
+            PhotonPipelineResult result, Optional<Matrix<N3, N3>> cameraMatrixOpt) {
+        if (result.hasTargets() && cameraMatrixOpt.isPresent()) {
+            PhotonTrackedTarget target = result.getBestTarget();
+
+            List<TargetCorner> targetCorners = target.getDetectedCorners();
+            double sumX = 0.0;
+            double sumY = 0.0;
+            for (TargetCorner t : targetCorners) {
+                sumX += t.x;
+                sumY += t.y;
+            }
+
+            Point tagCenter = new Point(sumX / 4, sumY / 4);
+
+            Rotation3d tagAngle = PhotonUtils.correctPixelRot(tagCenter, cameraMatrixOpt.get());
+
+            int tagID = target.getFiducialId();
+            Optional<Pose3d> tagPose = fieldTags.getTagPose(tagID);
+
+            if (tagPose.isEmpty()) {
+                reportFiducialPoseError(tagID);
+                return Optional.empty();
+            }
+
+            // calculate distance to the tag
+            double hypot =
+                    -PhotonUtils.calculateDistanceToTargetMeters(
+                            robotToCamera.getZ(), tagPose.get().getZ(), robotToCamera.getY(), tagAngle.getY());
+
+            // calculating the translation of the tag from the robot center on the floor plane
+            double robotX = Math.sin(tagAngle.getZ() + robotToCamera.getRotation().getZ()) * hypot;
+            double robotY = -Math.cos(tagAngle.getZ() + robotToCamera.getRotation().getZ()) * hypot;
+
+            Translation2d tagToCamera = new Translation2d(robotX, robotY);
+
+            Translation2d robotToTag = tagToCamera.plus(robotToCamera.getTranslation().toTranslation2d());
+
+            Optional<Rotation2d> headingState = PhotonUtils.getHeadingSample(result.getTimestampSeconds());
+            if(headingState.isEmpty()) {
+                return Optional.empty();
+            }
+
+            Rotation2d robotHeading = headingState.get();
+
+            // rotating the robot-relative tag position by the IMU yaw to make it field-relative
+            Translation2d fixedTrans = robotToTag.rotateBy(robotHeading);
+
+            Translation2d finalTranslation =
+                    tagPose.get().getTranslation().toTranslation2d().plus(fixedTrans);
+
+            Pose3d robotPosition =
+                    new Pose3d(new Pose2d(finalTranslation, robotHeading));
+
+            return Optional.of(
+                    new EstimatedRobotPose(
+                            robotPosition,
+                            result.getTimestampSeconds(),
+                            result.getTargets(),
+                            PoseStrategy.IMU_SOLVE_BEST_TAG));
+        } else {
+            return Optional.empty();
+        }
+    }
+
     private Optional<EstimatedRobotPose> multiTagOnRioStrategy(
             PhotonPipelineResult result,
             Optional<Matrix<N3, N3>> cameraMatrixOpt,

photon-lib/src/main/java/org/photonvision/PhotonUtils.java

@@ -22,9 +22,16 @@
  * SOFTWARE.
  */
 
+
 package org.photonvision;
 
+import edu.wpi.first.math.Matrix;
 import edu.wpi.first.math.geometry.*;
+import edu.wpi.first.math.interpolation.TimeInterpolatableBuffer;
+import edu.wpi.first.math.numbers.N3;
+import org.opencv.core.Point;
+
+import java.util.Optional;
 
 public final class PhotonUtils {
     private PhotonUtils() {
@@ -206,4 +213,50 @@ public static Rotation2d getYawToPose(Pose2d robotPose, Pose2d targetPose) {
     public static double getDistanceToPose(Pose2d robotPose, Pose2d targetPose) {
         return robotPose.getTranslation().getDistance(targetPose.getTranslation());
     }
+
+    /**
+     * Corrects a given pixel for perspective distortion
+     *
+     * @param pixel pixel to be corrected
+     * @param camIntrinsics the distortion matrix of the camera
+     * @return the angles
+     */
+    public static Rotation3d correctPixelRot(Point pixel, Matrix<N3, N3> camIntrinsics) {
+        double fx = camIntrinsics.get(0, 0);
+        double cx = camIntrinsics.get(0, 2);
+        double xOffset = cx - pixel.x;
+
+        double fy = camIntrinsics.get(1, 1);
+        double cy = camIntrinsics.get(1, 2);
+        double yOffset = cy - pixel.y;
+
+        // calculate yaw normally
+        var yaw = new Rotation2d(fx, xOffset);
+        // correct pitch based on yaw
+        var pitch = new Rotation2d(fy / Math.cos(Math.atan(xOffset / fx)), -yOffset);
+
+        return new Rotation3d(0, pitch.getRadians(), yaw.getRadians());
+    }
+
+
+
+    private static final TimeInterpolatableBuffer<Rotation2d> headingBuffer = TimeInterpolatableBuffer.createBuffer(
+            1.0
+    );
+
+    private static final TimeInterpolatableBuffer<Double> headingVelocityBuffer = TimeInterpolatableBuffer.createDoubleBuffer(1.0);
+    public static void registerHeadingState(Rotation2d robotAngle, double velocity, double timestamp) {
+        headingBuffer.addSample(timestamp, robotAngle);
+        headingVelocityBuffer.addSample(timestamp, velocity);
+    }
+
+    public static Optional<Rotation2d> getHeadingSample(double timestamp) {
+        return headingBuffer.getSample(timestamp);
+    }
+
+    private static Optional<Double> getHeadingVelocitySample(double timestamp) {
+        return headingVelocityBuffer.getSample(timestamp);
+    }
+
+
 }