Add constrained solvePNP strategy (#1682)

Signed-off-by: Jade Turner <spacey-sooty@proton.me>
Co-authored-by: Jade Turner <spacey-sooty@proton.me>

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

@@ -97,9 +97,34 @@ public class PhotonPoseEstimator {
          *
          * <p>https://www.chiefdelphi.com/t/frc-6328-mechanical-advantage-2025-build-thread/477314/98
          */
-        PNP_DISTANCE_TRIG_SOLVE
+        PNP_DISTANCE_TRIG_SOLVE,
+
+        /**
+         * Solve a constrained version of the Perspective-n-Point problem with the robot's drivebase
+         * flat on the floor. This computation takes place on the RoboRIO, and typically takes not more
+         * than 2ms. See {@link PhotonPoseEstimator.ConstrainedSolvepnpParams} and {@link
+         * org.photonvision.jni.ConstrainedSolvepnpJni} for details and tuning handles this strategy
+         * exposes. This strategy needs addHeadingData called every frame so heading data is up-to-date.
+         * If Multi-Tag PNP is enabled on the coprocessor, it will be used to provide an initial seed to
+         * the optimization algorithm -- otherwise, the multi-tag fallback strategy will be used as the
+         * seed.
+         */
+        CONSTRAINED_SOLVEPNP
     }
 
+    /**
+     * Tuning handles we have over the CONSTRAINED_SOLVEPNP {@link PhotonPoseEstimator.PoseStrategy}.
+     * Internally, the cost function is a sum-squared of pixel reprojection error + (optionally)
+     * heading error * heading scale factor.
+     *
+     * @param headingFree If true, heading is completely free to vary. If false, heading excursions
+     *     from the provided heading measurement will be penalized
+     * @param headingScaleFactor If headingFree is false, this weights the cost of changing our robot
+     *     heading estimate against the tag corner reprojection error const.
+     */
+    public static final record ConstrainedSolvepnpParams(
+            boolean headingFree, double headingScaleFactor) {}
+
     private AprilTagFieldLayout fieldTags;
     private TargetModel tagModel = TargetModel.kAprilTag36h11;
     private PoseStrategy primaryStrategy;
@@ -345,6 +370,8 @@ public class PhotonPoseEstimator {
      *   <li>The timestamp of the provided pipeline result is the same as in the previous call to
      *       {@code update()}.
      *   <li>No targets were found in the pipeline results.
+     *   <li>The strategy is CONSTRAINED_SOLVEPNP, but no constrainedPnpParams were provided (use the
+     *       other function overload).
      * </ul>
      *
      * @param cameraMatrix Camera calibration data for multi-tag-on-rio strategy - can be empty
@@ -358,6 +385,32 @@ public class PhotonPoseEstimator {
             PhotonPipelineResult cameraResult,
             Optional<Matrix<N3, N3>> cameraMatrix,
             Optional<Matrix<N8, N1>> distCoeffs) {
+        return update(cameraResult, cameraMatrix, distCoeffs, Optional.empty());
+    }
+
+    /**
+     * Updates the estimated position of the robot. Returns empty if:
+     *
+     * <ul>
+     *   <li>The timestamp of the provided pipeline result is the same as in the previous call to
+     *       {@code update()}.
+     *   <li>No targets were found in the pipeline results.
+     *   <li>The strategy is CONSTRAINED_SOLVEPNP, but the provided constrainedPnpParams are empty.
+     * </ul>
+     *
+     * @param cameraMatrix Camera calibration data for multi-tag-on-rio strategy - can be empty
+     *     otherwise
+     * @param distCoeffs Camera calibration data for multi-tag-on-rio strategy - can be empty
+     *     otherwise
+     * @param constrainedPnpParams Constrained SolvePNP params, if needed.
+     * @return an {@link EstimatedRobotPose} with an estimated pose, timestamp, and targets used to
+     *     create the estimate.
+     */
+    public Optional<EstimatedRobotPose> update(
+            PhotonPipelineResult cameraResult,
+            Optional<Matrix<N3, N3>> cameraMatrix,
+            Optional<Matrix<N8, N1>> distCoeffs,
+            Optional<ConstrainedSolvepnpParams> constrainedPnpParams) {
         // Time in the past -- give up, since the following if expects times > 0
         if (cameraResult.getTimestampSeconds() < 0) {
             return Optional.empty();
@@ -379,7 +432,8 @@ public class PhotonPoseEstimator {
             return Optional.empty();
         }
 
-        return update(cameraResult, cameraMatrix, distCoeffs, this.primaryStrategy);
+        return update(
+                cameraResult, cameraMatrix, distCoeffs, constrainedPnpParams, this.primaryStrategy);
     }
 
     /**
@@ -387,19 +441,20 @@ public class PhotonPoseEstimator {
      * called after timestamp checks have been done by another update() overload.
      *
      * @param cameraResult The latest pipeline result from the camera
-     * @param strategy The pose strategy to use
+     * @param strategy The pose strategy to use. Can't be CONSTRAINED_SOLVEPNP.
      * @return an {@link EstimatedRobotPose} with an estimated pose, timestamp, and targets used to
      *     create the estimate.
      */
     private Optional<EstimatedRobotPose> update(
             PhotonPipelineResult cameraResult, PoseStrategy strategy) {
-        return update(cameraResult, Optional.empty(), Optional.empty(), strategy);
+        return update(cameraResult, Optional.empty(), Optional.empty(), Optional.empty(), strategy);
     }
 
     private Optional<EstimatedRobotPose> update(
             PhotonPipelineResult cameraResult,
             Optional<Matrix<N3, N3>> cameraMatrix,
             Optional<Matrix<N8, N1>> distCoeffs,
+            Optional<ConstrainedSolvepnpParams> constrainedPnpParams,
             PoseStrategy strategy) {
         Optional<EstimatedRobotPose> estimatedPose =
                 switch (strategy) {
@@ -416,6 +471,8 @@ public class PhotonPoseEstimator {
                             multiTagOnRioStrategy(cameraResult, cameraMatrix, distCoeffs);
                     case MULTI_TAG_PNP_ON_COPROCESSOR -> multiTagOnCoprocStrategy(cameraResult);
                     case PNP_DISTANCE_TRIG_SOLVE -> pnpDistanceTrigSolveStrategy(cameraResult);
+                    case CONSTRAINED_SOLVEPNP ->
+                            constrainedPnpStrategy(cameraResult, cameraMatrix, distCoeffs, constrainedPnpParams);
                 };
 
         if (estimatedPose.isPresent()) {
@@ -470,6 +527,87 @@ public class PhotonPoseEstimator {
                         PoseStrategy.PNP_DISTANCE_TRIG_SOLVE));
     }
 
+    private Optional<EstimatedRobotPose> constrainedPnpStrategy(
+            PhotonPipelineResult result,
+            Optional<Matrix<N3, N3>> cameraMatrixOpt,
+            Optional<Matrix<N8, N1>> distCoeffsOpt,
+            Optional<ConstrainedSolvepnpParams> constrainedPnpParams) {
+        boolean hasCalibData = cameraMatrixOpt.isPresent() && distCoeffsOpt.isPresent();
+        // cannot run multitagPNP, use fallback strategy
+        if (!hasCalibData) {
+            return update(
+                    result, cameraMatrixOpt, distCoeffsOpt, Optional.empty(), this.multiTagFallbackStrategy);
+        }
+
+        if (constrainedPnpParams.isEmpty()) {
+            return Optional.empty();
+        }
+
+        // Need heading if heading fixed
+        if (!constrainedPnpParams.get().headingFree
+                && headingBuffer.getSample(result.getTimestampSeconds()).isEmpty()) {
+            return update(
+                    result, cameraMatrixOpt, distCoeffsOpt, Optional.empty(), this.multiTagFallbackStrategy);
+        }
+
+        Pose3d fieldToRobotSeed;
+
+        // Attempt to use multi-tag to get a pose estimate seed
+        if (result.getMultiTagResult().isPresent()) {
+            fieldToRobotSeed =
+                    Pose3d.kZero.plus(
+                            result.getMultiTagResult().get().estimatedPose.best.plus(robotToCamera.inverse()));
+        } else {
+            // HACK - use fallback strategy to gimme a seed pose
+            // TODO - make sure nested update doesn't break state
+            var nestedUpdate =
+                    update(
+                            result,
+                            cameraMatrixOpt,
+                            distCoeffsOpt,
+                            Optional.empty(),
+                            this.multiTagFallbackStrategy);
+            if (nestedUpdate.isEmpty()) {
+                // best i can do is bail
+                return Optional.empty();
+            }
+            fieldToRobotSeed = nestedUpdate.get().estimatedPose;
+        }
+
+        if (!constrainedPnpParams.get().headingFree) {
+            // If heading fixed, force rotation component
+            fieldToRobotSeed =
+                    new Pose3d(
+                            fieldToRobotSeed.getTranslation(),
+                            new Rotation3d(headingBuffer.getSample(result.getTimestampSeconds()).get()));
+        }
+
+        var pnpResult =
+                VisionEstimation.estimateRobotPoseConstrainedSolvepnp(
+                        cameraMatrixOpt.get(),
+                        distCoeffsOpt.get(),
+                        result.getTargets(),
+                        robotToCamera,
+                        fieldToRobotSeed,
+                        fieldTags,
+                        tagModel,
+                        constrainedPnpParams.get().headingFree,
+                        headingBuffer.getSample(result.getTimestampSeconds()).get(),
+                        constrainedPnpParams.get().headingScaleFactor);
+        // try fallback strategy if solvePNP fails for some reason
+        if (!pnpResult.isPresent())
+            return update(
+                    result, cameraMatrixOpt, distCoeffsOpt, Optional.empty(), this.multiTagFallbackStrategy);
+        var best = Pose3d.kZero.plus(pnpResult.get().best); // field-to-robot
+
+        return Optional.of(
+                new EstimatedRobotPose(
+                        best,
+                        result.getTimestampSeconds(),
+                        result.getTargets(),
+                        PoseStrategy.CONSTRAINED_SOLVEPNP));
+    }
+
     private Optional<EstimatedRobotPose> multiTagOnCoprocStrategy(PhotonPipelineResult result) {
         if (result.getMultiTagResult().isEmpty()) {
             return update(result, this.multiTagFallbackStrategy);
@@ -477,7 +615,7 @@ public class PhotonPoseEstimator {
 
         var best_tf = result.getMultiTagResult().get().estimatedPose.best;
         var best =
-                new Pose3d()
+                Pose3d.kZero
                         .plus(best_tf) // field-to-camera
                         .relativeTo(fieldTags.getOrigin())
                         .plus(robotToCamera.inverse()); // field-to-robot
@@ -508,9 +646,12 @@ public class PhotonPoseEstimator {
                 VisionEstimation.estimateCamPosePNP(
                         cameraMatrixOpt.get(), distCoeffsOpt.get(), result.getTargets(), fieldTags, tagModel);
         // try fallback strategy if solvePNP fails for some reason
-        if (!pnpResult.isPresent()) return update(result, this.multiTagFallbackStrategy);
+        if (!pnpResult.isPresent())
+            return update(
+                    result, cameraMatrixOpt, distCoeffsOpt, Optional.empty(), this.multiTagFallbackStrategy);
+
         var best =
-                new Pose3d()
+                Pose3d.kZero
                         .plus(pnpResult.get().best) // field-to-camera
                         .plus(robotToCamera.inverse()); // field-to-robot