Add Constrained PNP Pose Strategies to C++ photonlib (#1908)

This adds the two missing pose strategies from the java version of
photonlib (Constrained PNP and the Trig solve), to C++ photonlib

---------

Co-authored-by: Matthew Morley <matthew.morley.ca@gmail.com>

photon-lib/src/main/native/cpp/photon/PhotonPoseEstimator.cpp

@@ -49,6 +49,8 @@
 #include <units/time.h>
 
 #include "photon/PhotonCamera.h"
+#include "photon/estimation/TargetModel.h"
+#include "photon/estimation/VisionEstimation.h"
 #include "photon/targeting/PhotonPipelineResult.h"
 #include "photon/targeting/PhotonTrackedTarget.h"
 
@@ -99,7 +101,8 @@ void PhotonPoseEstimator::SetMultiTagFallbackStrategy(PoseStrategy strategy) {
 std::optional<EstimatedRobotPose> PhotonPoseEstimator::Update(
     const PhotonPipelineResult& result,
     std::optional<PhotonCamera::CameraMatrix> cameraMatrixData,
-    std::optional<PhotonCamera::DistortionMatrix> cameraDistCoeffs) {
+    std::optional<PhotonCamera::DistortionMatrix> cameraDistCoeffs,
+    std::optional<ConstrainedSolvepnpParams> constrainedPnpParams) {
   // Time in the past -- give up, since the following if expects times > 0
   if (result.GetTimestamp() < 0_s) {
     FRC_ReportError(frc::warn::Warning,
@@ -122,13 +125,15 @@ std::optional<EstimatedRobotPose> PhotonPoseEstimator::Update(
     return std::nullopt;
   }
 
-  return Update(result, cameraMatrixData, cameraDistCoeffs, this->strategy);
+  return Update(result, cameraMatrixData, cameraDistCoeffs,
+                constrainedPnpParams, this->strategy);
 }
 
 std::optional<EstimatedRobotPose> PhotonPoseEstimator::Update(
     const PhotonPipelineResult& result,
     std::optional<PhotonCamera::CameraMatrix> cameraMatrixData,
     std::optional<PhotonCamera::DistortionMatrix> cameraDistCoeffs,
+    std::optional<ConstrainedSolvepnpParams> constrainedPnpParams,
     PoseStrategy strategy) {
   std::optional<EstimatedRobotPose> ret = std::nullopt;
 
@@ -161,6 +166,10 @@ std::optional<EstimatedRobotPose> PhotonPoseEstimator::Update(
                         "");
       }
       break;
+    case CONSTRAINED_SOLVEPNP:
+      ret = ConstrainedPnpStrategy(result, cameraMatrixData, cameraDistCoeffs,
+                                   constrainedPnpParams);
+      break;
     case PNP_DISTANCE_TRIG_SOLVE:
       ret = PnpDistanceTrigSolveStrategy(result);
       break;
@@ -527,4 +536,74 @@ PhotonPoseEstimator::AverageBestTargetsStrategy(PhotonPipelineResult result) {
                             result.GetTimestamp(), result.GetTargets(),
                             AVERAGE_BEST_TARGETS};
 }
+
+std::optional<EstimatedRobotPose> PhotonPoseEstimator::ConstrainedPnpStrategy(
+    photon::PhotonPipelineResult result,
+    std::optional<photon::PhotonCamera::CameraMatrix> camMat,
+    std::optional<photon::PhotonCamera::DistortionMatrix> distCoeffs,
+    std::optional<ConstrainedSolvepnpParams> constrainedPnpParams) {
+  using namespace frc;
+
+  if (!camMat || !distCoeffs) {
+    FRC_ReportError(frc::warn::Warning,
+                    "No camera calibration data provided to "
+                    "StrPoseEstimator::MultiTagOnRioStrategy!",
+                    "");
+    return Update(result, this->multiTagFallbackStrategy);
+  }
+
+  if (!constrainedPnpParams) {
+    return {};
+  }
+
+  if (!constrainedPnpParams->headingFree &&
+      !headingBuffer.Sample(result.GetTimestamp()).has_value()) {
+    return Update(result, camMat, distCoeffs, {},
+                  this->multiTagFallbackStrategy);
+  }
+
+  frc::Pose3d fieldToRobotSeed;
+
+  if (result.MultiTagResult().has_value()) {
+    fieldToRobotSeed =
+        frc::Pose3d{} + (result.MultiTagResult()->estimatedPose.best +
+                         m_robotToCamera.Inverse());
+  } else {
+    std::optional<EstimatedRobotPose> nestedUpdate =
+        Update(result, camMat, distCoeffs, {}, this->multiTagFallbackStrategy);
+
+    if (!nestedUpdate.has_value()) {
+      return {};
+    }
+
+    fieldToRobotSeed = nestedUpdate->estimatedPose;
+  }
+
+  if (!constrainedPnpParams.value().headingFree) {
+    fieldToRobotSeed = frc::Pose3d{
+        fieldToRobotSeed.Translation(),
+        frc::Rotation3d{headingBuffer.Sample(result.GetTimestamp()).value()}};
+  }
+
+  std::vector<photon::PhotonTrackedTarget> targets{result.GetTargets().begin(),
+                                                   result.GetTargets().end()};
+
+  std::optional<photon::PnpResult> pnpResult =
+      VisionEstimation::EstimateRobotPoseConstrainedSolvePNP(
+          camMat.value(), distCoeffs.value(), targets, m_robotToCamera,
+          fieldToRobotSeed, aprilTags, photon::kAprilTag36h11,
+          constrainedPnpParams->headingFree,
+          frc::Rotation2d{headingBuffer.Sample(result.GetTimestamp()).value()},
+          constrainedPnpParams->headingScalingFactor);
+
+  if (!pnpResult) {
+    return Update(result, camMat, distCoeffs, {},
+                  this->multiTagFallbackStrategy);
+  }
+
+  frc::Pose3d best = frc::Pose3d{} + pnpResult->best;
+
+  return EstimatedRobotPose{best, result.GetTimestamp(), result.GetTargets(),
+                            PoseStrategy::CONSTRAINED_SOLVEPNP};
+}
 }  // namespace photon

photon-lib/src/main/native/include/photon/PhotonPoseEstimator.h

@@ -28,6 +28,7 @@
 
 #include <frc/apriltag/AprilTagFieldLayout.h>
 #include <frc/geometry/Pose3d.h>
+#include <frc/geometry/Rotation3d.h>
 #include <frc/geometry/Transform3d.h>
 #include <frc/interpolation/TimeInterpolatableBuffer.h>
 #include <opencv2/core/mat.hpp>
@@ -48,7 +49,13 @@ enum PoseStrategy {
   AVERAGE_BEST_TARGETS,
   MULTI_TAG_PNP_ON_COPROCESSOR,
   MULTI_TAG_PNP_ON_RIO,
-  PNP_DISTANCE_TRIG_SOLVE,
+  CONSTRAINED_SOLVEPNP,
+  PNP_DISTANCE_TRIG_SOLVE
+};
+
+struct ConstrainedSolvepnpParams {
+  bool headingFree{false};
+  double headingScalingFactor{0.0};
 };
 
 struct EstimatedRobotPose {
@@ -239,11 +246,16 @@ class PhotonPoseEstimator {
    * Only required if doing multitag-on-rio, and may be nullopt otherwise.
    * @param distCoeffsData The camera calibration distortion coefficients. Only
    * required if doing multitag-on-rio, and may be nullopt otherwise.
+   * @param constrainedPnpParams Constrained SolvePNP params, if needed.
    */
   std::optional<EstimatedRobotPose> Update(
-      const PhotonPipelineResult& result,
-      std::optional<PhotonCamera::CameraMatrix> cameraMatrixData = std::nullopt,
-      std::optional<PhotonCamera::DistortionMatrix> coeffsData = std::nullopt);
+      const photon::PhotonPipelineResult& result,
+      std::optional<photon::PhotonCamera::CameraMatrix> cameraMatrixData =
+          std::nullopt,
+      std::optional<photon::PhotonCamera::DistortionMatrix> coeffsData =
+          std::nullopt,
+      std::optional<ConstrainedSolvepnpParams> constrainedPnpParams =
+          std::nullopt);
 
  private:
   frc::AprilTagFieldLayout aprilTags;
@@ -275,13 +287,14 @@ class PhotonPoseEstimator {
    */
   std::optional<EstimatedRobotPose> Update(const PhotonPipelineResult& result,
                                            PoseStrategy strategy) {
-    return Update(result, std::nullopt, std::nullopt, strategy);
+    return Update(result, std::nullopt, std::nullopt, std::nullopt, strategy);
   }
 
   std::optional<EstimatedRobotPose> Update(
       const PhotonPipelineResult& result,
       std::optional<PhotonCamera::CameraMatrix> cameraMatrixData,
       std::optional<PhotonCamera::DistortionMatrix> coeffsData,
+      std::optional<ConstrainedSolvepnpParams> constrainedPnpParams,
       PoseStrategy strategy);
 
   /**
@@ -355,6 +368,12 @@ class PhotonPoseEstimator {
    */
   std::optional<EstimatedRobotPose> AverageBestTargetsStrategy(
       PhotonPipelineResult result);
+
+  std::optional<EstimatedRobotPose> ConstrainedPnpStrategy(
+      photon::PhotonPipelineResult result,
+      std::optional<photon::PhotonCamera::CameraMatrix> camMat,
+      std::optional<photon::PhotonCamera::DistortionMatrix> distCoeffs,
+      std::optional<ConstrainedSolvepnpParams> constrainedPnpParams);
 };
 
 }  // namespace photon

photon-lib/src/test/native/cpp/PhotonPoseEstimatorTest.cpp

@@ -40,6 +40,7 @@
 #include "photon/dataflow/structures/Packet.h"
 #include "photon/simulation/PhotonCameraSim.h"
 #include "photon/simulation/SimCameraProperties.h"
+#include "photon/simulation/VisionTargetSim.h"
 #include "photon/targeting/MultiTargetPNPResult.h"
 #include "photon/targeting/PhotonPipelineResult.h"
 #include "photon/targeting/PhotonTrackedTarget.h"
@@ -583,3 +584,76 @@ TEST(PhotonPoseEstimatorTest, CopyResult) {
   EXPECT_NEAR(testResult.GetTimestamp().to<double>(),
               test2.GetTimestamp().to<double>(), 0.001);
 }
+
+TEST(PhotonPoseEstimatorTest, ConstrainedPnpEmptyCase) {
+  photon::PhotonPoseEstimator estimator(
+      frc::AprilTagFieldLayout::LoadField(frc::AprilTagField::k2024Crescendo),
+      photon::CONSTRAINED_SOLVEPNP, frc::Transform3d());
+
+  photon::PhotonPipelineResult result;
+  auto estimate = estimator.Update(result);
+  EXPECT_FALSE(estimate.has_value());
+}
+
+TEST(PhotonPoseEstimatorTest, ConstrainedPnpOneTag) {
+  photon::PhotonCamera cameraOne = photon::PhotonCamera("test");
+  auto distortion = Eigen::VectorXd::Zero(8);
+  auto cameraMat = Eigen::Matrix3d{{399.37500000000006, 0, 319.5},
+                                   {0, 399.16666666666674, 239.5},
+                                   {0, 0, 1}};
+
+  // Create corners data matching the Java test
+  std::vector<photon::TargetCorner> corners8{
+      photon::TargetCorner{98.09875447066685, 331.0093220119495},
+      photon::TargetCorner{122.20226758624413, 335.50083894738486},
+      photon::TargetCorner{127.17118732489361, 313.81406314178633},
+      photon::TargetCorner{104.28543773760417, 309.6516557438994}};
+
+  frc::Transform3d poseTransform(
+      frc::Translation3d(3.1665557336121353_m, 4.430673446050584_m,
+                         0.48678786477534686_m),
+      frc::Rotation3d(frc::Quaternion(0.3132532247418243, 0.24722671090692333,
+                                      -0.08413452932300695,
+                                      0.9130568172784148)));
+
+  std::vector<photon::PhotonTrackedTarget> targets{
+      photon::PhotonTrackedTarget{0.0, 0.0, 0.0, 0.0, 8, 0, 0.0f, poseTransform,
+                                  poseTransform, 0.0, corners8, corners8}};
+
+  auto multiTagResult = std::make_optional<photon::MultiTargetPNPResult>(
+      photon::PnpResult{poseTransform, poseTransform, 0.1, 0.1, 0.0},
+      std::vector<int16_t>{8});
+
+  photon::PhotonPipelineResult result{
+      photon::PhotonPipelineMetadata{1, 10000, 2000, 100}, targets,
+      multiTagResult};
+
+  cameraOne.test = true;
+  cameraOne.testResult = {result};
+  cameraOne.testResult[0].SetReceiveTimestamp(units::second_t(15));
+
+  const units::radian_t camPitch = 30_deg;
+  const frc::Transform3d kRobotToCam{frc::Translation3d(0.5_m, 0.0_m, 0.5_m),
+                                     frc::Rotation3d(0_rad, -camPitch, 0_rad)};
+
+  photon::PhotonPoseEstimator estimator(
+      frc::AprilTagFieldLayout::LoadField(frc::AprilTagField::k2024Crescendo),
+      photon::CONSTRAINED_SOLVEPNP, kRobotToCam);
+
+  estimator.AddHeadingData(cameraOne.testResult[0].GetTimestamp(),
+                           frc::Rotation2d());
+
+  auto estimatedPose =
+      estimator.Update(cameraOne.testResult[0], cameraMat, distortion,
+                       photon::ConstrainedSolvepnpParams{true, 0});
+
+  ASSERT_TRUE(estimatedPose.has_value());
+
+  frc::Pose3d pose = estimatedPose.value().estimatedPose;
+
+  EXPECT_NEAR(3.58, units::unit_cast<double>(pose.X()), 0.01);
+  EXPECT_NEAR(4.13, units::unit_cast<double>(pose.Y()), 0.01);
+  EXPECT_NEAR(0.0, units::unit_cast<double>(pose.Z()), 0.01);
+
+  EXPECT_EQ(photon::CONSTRAINED_SOLVEPNP, estimatedPose.value().strategy);
+}