Add secondary solver to target mapper

We first already solved for the relative constraints between targets.
Now fit this whole map onto where it should be in the world, by solving
for the transformation between the fixed target pose and where it
actually should be on the field.

Signed-off-by: milind-u <milind.upadhyay@gmail.com>
Change-Id: Ib93bc2d5f38a1e282826461a7cd33389f8566021
diff --git a/frc971/vision/BUILD b/frc971/vision/BUILD
index 133281b..84f3be7 100644
--- a/frc971/vision/BUILD
+++ b/frc971/vision/BUILD
@@ -166,6 +166,7 @@
         ":target_map_fbs",
         "//aos/events:simulated_event_loop",
         "//frc971/control_loops:control_loop",
+        "//frc971/vision:visualize_robot",
         "//frc971/vision/ceres:pose_graph_3d_lib",
         "//third_party:opencv",
         "@com_google_ceres_solver//:ceres",
diff --git a/frc971/vision/target_mapper.cc b/frc971/vision/target_mapper.cc
index 3eba919..0d6b9e9 100644
--- a/frc971/vision/target_mapper.cc
+++ b/frc971/vision/target_mapper.cc
@@ -10,9 +10,12 @@
 DEFINE_double(distortion_noise_scalar, 1.0,
               "Scale the target pose distortion factor by this when computing "
               "the noise.");
-DEFINE_uint64(
+DEFINE_int32(
     frozen_target_id, 1,
     "Freeze the pose of this target so the map can have one fixed point.");
+DEFINE_int32(min_target_id, 1, "Minimum target id to solve for");
+DEFINE_int32(max_target_id, 8, "Maximum target id to solve for");
+DEFINE_bool(visualize_solver, false, "If true, visualize the solving process.");
 
 namespace frc971::vision {
 Eigen::Affine3d PoseUtils::Pose3dToAffine3d(
@@ -211,19 +214,28 @@
 TargetMapper::TargetMapper(
     std::string_view target_poses_path,
     const ceres::examples::VectorOfConstraints &target_constraints)
-    : target_constraints_(target_constraints) {
+    : target_constraints_(target_constraints),
+      T_frozen_actual_(Eigen::Vector3d::Zero()),
+      R_frozen_actual_(Eigen::Quaterniond::Identity()),
+      vis_robot_(cv::Size(1280, 1000)) {
   aos::FlatbufferDetachedBuffer<TargetMap> target_map =
       aos::JsonFileToFlatbuffer<TargetMap>(target_poses_path);
   for (const auto *target_pose_fbs : *target_map.message().target_poses()) {
-    target_poses_[target_pose_fbs->id()] =
+    ideal_target_poses_[target_pose_fbs->id()] =
         PoseUtils::TargetPoseFromFbs(*target_pose_fbs).pose;
   }
+  target_poses_ = ideal_target_poses_;
 }
 
 TargetMapper::TargetMapper(
     const ceres::examples::MapOfPoses &target_poses,
     const ceres::examples::VectorOfConstraints &target_constraints)
-    : target_poses_(target_poses), target_constraints_(target_constraints) {}
+    : ideal_target_poses_(target_poses),
+      target_poses_(ideal_target_poses_),
+      target_constraints_(target_constraints),
+      T_frozen_actual_(Eigen::Vector3d::Zero()),
+      R_frozen_actual_(Eigen::Quaterniond::Identity()),
+      vis_robot_(cv::Size(1280, 1000)) {}
 
 std::optional<TargetMapper::TargetPose> TargetMapper::GetTargetPoseById(
     std::vector<TargetMapper::TargetPose> target_poses, TargetId target_id) {
@@ -248,7 +260,7 @@
 // Taken from ceres/examples/slam/pose_graph_3d/pose_graph_3d.cc
 // Constructs the nonlinear least squares optimization problem from the pose
 // graph constraints.
-void TargetMapper::BuildOptimizationProblem(
+void TargetMapper::BuildTargetPoseOptimizationProblem(
     const ceres::examples::VectorOfConstraints &constraints,
     ceres::examples::MapOfPoses *poses, ceres::Problem *problem) {
   CHECK(poses != nullptr);
@@ -311,6 +323,33 @@
   problem->SetParameterBlockConstant(pose_start_iter->second.q.coeffs().data());
 }
 
+void TargetMapper::BuildMapFittingOptimizationProblem(ceres::Problem *problem) {
+  // Setup robot visualization
+  vis_robot_.ClearImage();
+  constexpr int kImageWidth = 1280;
+  constexpr double kFocalLength = 500.0;
+  vis_robot_.SetDefaultViewpoint(kImageWidth, kFocalLength);
+
+  const size_t num_targets = FLAGS_max_target_id - FLAGS_min_target_id;
+  // Translation and rotation error for each target
+  const size_t num_residuals = num_targets * 6;
+  // Set up the only cost function (also known as residual). This uses
+  // auto-differentiation to obtain the derivative (jacobian).
+  ceres::CostFunction *cost_function =
+      new ceres::AutoDiffCostFunction<TargetMapper, ceres::DYNAMIC, 3, 4>(
+          this, num_residuals, ceres::DO_NOT_TAKE_OWNERSHIP);
+
+  ceres::LossFunction *loss_function = new ceres::HuberLoss(2.0);
+  ceres::LocalParameterization *quaternion_local_parameterization =
+      new ceres::EigenQuaternionParameterization;
+
+  problem->AddResidualBlock(cost_function, loss_function,
+                            T_frozen_actual_.vector().data(),
+                            R_frozen_actual_.coeffs().data());
+  problem->SetParameterization(R_frozen_actual_.coeffs().data(),
+                               quaternion_local_parameterization);
+}
+
 // Taken from ceres/examples/slam/pose_graph_3d/pose_graph_3d.cc
 bool TargetMapper::SolveOptimizationProblem(ceres::Problem *problem) {
   CHECK_NOTNULL(problem);
@@ -330,11 +369,39 @@
 
 void TargetMapper::Solve(std::string_view field_name,
                          std::optional<std::string_view> output_dir) {
-  ceres::Problem problem;
-  BuildOptimizationProblem(target_constraints_, &target_poses_, &problem);
+  ceres::Problem target_pose_problem;
+  BuildTargetPoseOptimizationProblem(target_constraints_, &target_poses_,
+                                     &target_pose_problem);
+  CHECK(SolveOptimizationProblem(&target_pose_problem))
+      << "The target pose solve was not successful, exiting.";
 
-  CHECK(SolveOptimizationProblem(&problem))
-      << "The solve was not successful, exiting.";
+  ceres::Problem map_fitting_problem;
+  BuildMapFittingOptimizationProblem(&map_fitting_problem);
+  CHECK(SolveOptimizationProblem(&map_fitting_problem))
+      << "The map fitting solve was not successful, exiting.";
+
+  Eigen::Affine3d H_frozen_actual = T_frozen_actual_ * R_frozen_actual_;
+  LOG(INFO) << "H_frozen_actual: "
+            << PoseUtils::Affine3dToPose3d(H_frozen_actual);
+
+  auto H_world_frozen =
+      PoseUtils::Pose3dToAffine3d(target_poses_[FLAGS_frozen_target_id]);
+  auto H_world_frozenactual = H_world_frozen * H_frozen_actual;
+
+  // Offset the solved poses to become the actual ones
+  for (auto &[id, pose] : target_poses_) {
+    // Don't offset targets we didn't solve for
+    if (id < FLAGS_min_target_id || id > FLAGS_max_target_id) {
+      continue;
+    }
+
+    // Take the delta between the frozen target and the solved target, and put
+    // that on top of the actual pose of the frozen target
+    auto H_world_solved = PoseUtils::Pose3dToAffine3d(pose);
+    auto H_frozen_solved = H_world_frozen.inverse() * H_world_solved;
+    auto H_world_actual = H_world_frozenactual * H_frozen_solved;
+    pose = PoseUtils::Affine3dToPose3d(H_world_actual);
+  }
 
   auto map_json = MapToJson(field_name);
   VLOG(1) << "Solved target poses: " << map_json;
@@ -366,6 +433,110 @@
       {.multi_line = true});
 }
 
+namespace {
+
+// Hacks to extract a double from a scalar, which is either a ceres jet or a
+// double. Only used for debugging and displaying.
+template <typename S>
+double ScalarToDouble(S s) {
+  const double *ptr = reinterpret_cast<double *>(&s);
+  return *ptr;
+}
+
+template <typename S>
+Eigen::Affine3d ScalarAffineToDouble(Eigen::Transform<S, 3, Eigen::Affine> H) {
+  Eigen::Affine3d H_double;
+  for (size_t i = 0; i < H.rows(); i++) {
+    for (size_t j = 0; j < H.cols(); j++) {
+      H_double(i, j) = ScalarToDouble(H(i, j));
+    }
+  }
+  return H_double;
+}
+
+}  // namespace
+
+template <typename S>
+bool TargetMapper::operator()(const S *const translation,
+                              const S *const rotation, S *residual) const {
+  using Affine3s = Eigen::Transform<S, 3, Eigen::Affine>;
+  Eigen::Quaternion<S> R_frozen_actual(rotation[3], rotation[1], rotation[2],
+                                       rotation[0]);
+  Eigen::Translation<S, 3> T_frozen_actual(translation[0], translation[1],
+                                           translation[2]);
+  // Actual target pose in the frame of the fixed pose.
+  Affine3s H_frozen_actual = T_frozen_actual * R_frozen_actual;
+  VLOG(2) << "H_frozen_actual: "
+          << PoseUtils::Affine3dToPose3d(ScalarAffineToDouble(H_frozen_actual));
+
+  Affine3s H_world_frozen =
+      PoseUtils::Pose3dToAffine3d(target_poses_.at(FLAGS_frozen_target_id))
+          .cast<S>();
+  Affine3s H_world_frozenactual = H_world_frozen * H_frozen_actual;
+
+  size_t residual_index = 0;
+  if (FLAGS_visualize_solver) {
+    vis_robot_.ClearImage();
+  }
+
+  for (const auto &[id, solved_pose] : target_poses_) {
+    if (id < FLAGS_min_target_id || id > FLAGS_max_target_id) {
+      continue;
+    }
+
+    Affine3s H_world_ideal =
+        PoseUtils::Pose3dToAffine3d(ideal_target_poses_.at(id)).cast<S>();
+    Affine3s H_world_solved =
+        PoseUtils::Pose3dToAffine3d(solved_pose).cast<S>();
+    // Take the delta between the frozen target and the solved target, and put
+    // that on top of the actual pose of the frozen target
+    auto H_frozen_solved = H_world_frozen.inverse() * H_world_solved;
+    auto H_world_actual = H_world_frozenactual * H_frozen_solved;
+    VLOG(2) << id << ": " << H_world_actual.translation();
+    Affine3s H_ideal_actual = H_world_ideal.inverse() * H_world_actual;
+    auto T_ideal_actual = H_ideal_actual.translation();
+    VLOG(2) << "T_ideal_actual: " << T_ideal_actual;
+    VLOG(2);
+    auto R_ideal_actual = Eigen::AngleAxis<S>(H_ideal_actual.rotation());
+
+    constexpr double kTranslationScalar = 100.0;
+    constexpr double kRotationScalar = 1000.0;
+
+    // Penalize based on how much our actual poses matches the ideal
+    // ones. We've already solved for the relative poses, now figure out
+    // where all of them fit in the world.
+    residual[residual_index++] = kTranslationScalar * T_ideal_actual(0);
+    residual[residual_index++] = kTranslationScalar * T_ideal_actual(1);
+    residual[residual_index++] = kTranslationScalar * T_ideal_actual(2);
+    residual[residual_index++] =
+        kRotationScalar * R_ideal_actual.angle() * R_ideal_actual.axis().x();
+    residual[residual_index++] =
+        kRotationScalar * R_ideal_actual.angle() * R_ideal_actual.axis().y();
+    residual[residual_index++] =
+        kRotationScalar * R_ideal_actual.angle() * R_ideal_actual.axis().z();
+
+    if (FLAGS_visualize_solver) {
+      vis_robot_.DrawFrameAxes(ScalarAffineToDouble(H_world_actual),
+                               std::to_string(id), cv::Scalar(0, 255, 0));
+      vis_robot_.DrawFrameAxes(ScalarAffineToDouble(H_world_ideal),
+                               std::to_string(id), cv::Scalar(255, 255, 255));
+    }
+  }
+  if (FLAGS_visualize_solver) {
+    cv::imshow("Target maps", vis_robot_.image_);
+    cv::waitKey(0);
+  }
+
+  // Ceres can't handle residual values of exactly zero
+  for (size_t i = 0; i < residual_index; i++) {
+    if (residual[i] == S(0)) {
+      residual[i] = S(1e-9);
+    }
+  }
+
+  return true;
+}
+
 }  // namespace frc971::vision
 
 std::ostream &operator<<(std::ostream &os, ceres::examples::Pose3d pose) {
diff --git a/frc971/vision/target_mapper.h b/frc971/vision/target_mapper.h
index c782992..ca36866 100644
--- a/frc971/vision/target_mapper.h
+++ b/frc971/vision/target_mapper.h
@@ -7,6 +7,7 @@
 #include "ceres/ceres.h"
 #include "frc971/vision/ceres/types.h"
 #include "frc971/vision/target_map_generated.h"
+#include "frc971/vision/visualize_robot.h"
 
 namespace frc971::vision {
 
@@ -51,18 +52,36 @@
 
   ceres::examples::MapOfPoses target_poses() { return target_poses_; }
 
+  // Cost function for the secondary solver finding out where the whole map fits
+  // in the world
+  template <typename S>
+  bool operator()(const S *const translation, const S *const rotation,
+                  S *residual) const;
+
  private:
   // Constructs the nonlinear least squares optimization problem from the
   // pose graph constraints.
-  void BuildOptimizationProblem(
+  void BuildTargetPoseOptimizationProblem(
       const ceres::examples::VectorOfConstraints &constraints,
       ceres::examples::MapOfPoses *poses, ceres::Problem *problem);
 
+  // Constructs the nonlinear least squares optimization problem for the solved
+  // -> actual pose solver.
+  void BuildMapFittingOptimizationProblem(ceres::Problem *problem);
+
   // Returns true if the solve was successful.
   bool SolveOptimizationProblem(ceres::Problem *problem);
 
+  ceres::examples::MapOfPoses ideal_target_poses_;
   ceres::examples::MapOfPoses target_poses_;
   ceres::examples::VectorOfConstraints target_constraints_;
+
+  // Transformation moving the target map we solved for to where it actually
+  // should be in the world
+  Eigen::Translation3d T_frozen_actual_;
+  Eigen::Quaterniond R_frozen_actual_;
+
+  mutable VisualizeRobot vis_robot_;
 };
 
 // Utility functions for dealing with ceres::examples::Pose3d structs
diff --git a/frc971/vision/target_mapper_test.cc b/frc971/vision/target_mapper_test.cc
index cd7d18b..1371f89 100644
--- a/frc971/vision/target_mapper_test.cc
+++ b/frc971/vision/target_mapper_test.cc
@@ -9,6 +9,9 @@
 #include "glog/logging.h"
 #include "gtest/gtest.h"
 
+DECLARE_int32(min_target_id);
+DECLARE_int32(max_target_id);
+
 namespace frc971::vision {
 
 namespace {
@@ -338,6 +341,9 @@
 }
 
 TEST(TargetMapperTest, TwoTargetsOneConstraint) {
+  FLAGS_min_target_id = 0;
+  FLAGS_max_target_id = 1;
+
   ceres::examples::MapOfPoses target_poses;
   target_poses[0] = Make2dPose(5.0, 0.0, M_PI);
   target_poses[1] = Make2dPose(-5.0, 0.0, 0.0);
@@ -361,6 +367,9 @@
 }
 
 TEST(TargetMapperTest, TwoTargetsTwoConstraints) {
+  FLAGS_min_target_id = 0;
+  FLAGS_max_target_id = 1;
+
   ceres::examples::MapOfPoses target_poses;
   target_poses[0] = Make2dPose(5.0, 0.0, M_PI);
   target_poses[1] = Make2dPose(-5.0, 0.0, -M_PI_2);
@@ -390,6 +399,9 @@
 }
 
 TEST(TargetMapperTest, TwoTargetsOneNoisyConstraint) {
+  FLAGS_min_target_id = 0;
+  FLAGS_max_target_id = 1;
+
   ceres::examples::MapOfPoses target_poses;
   target_poses[0] = Make2dPose(5.0, 0.0, M_PI);
   target_poses[1] = Make2dPose(-5.0, 0.0, 0.0);
@@ -549,33 +561,6 @@
             .value();
     EXPECT_POSE_NEAR(mapper_target_pose, actual_target_pose.pose);
   }
-
-  //
-  // See what happens when we don't start with the
-  // correct values
-  //
-  for (auto [target_id, target_pose] : target_poses) {
-    // Skip first pose, since that needs to be correct
-    // and is fixed in the solver
-    if (target_id != 1) {
-      ceres::examples::Pose3d bad_pose{
-          Eigen::Vector3d::Zero(),
-          PoseUtils::EulerAnglesToQuaternion(Eigen::Vector3d::Zero())};
-      target_poses[target_id] = bad_pose;
-    }
-  }
-
-  frc971::vision::TargetMapper mapper_bad_poses(target_poses,
-                                                target_constraints);
-  mapper_bad_poses.Solve(kFieldName);
-
-  for (auto [target_pose_id, mapper_target_pose] :
-       mapper_bad_poses.target_poses()) {
-    TargetMapper::TargetPose actual_target_pose =
-        TargetMapper::GetTargetPoseById(actual_target_poses, target_pose_id)
-            .value();
-    EXPECT_POSE_NEAR(mapper_target_pose, actual_target_pose.pose);
-  }
 }
 
 }  // namespace frc971::vision
diff --git a/y2023/vision/target_mapping.cc b/y2023/vision/target_mapping.cc
index ba5f5d9..daba90c 100644
--- a/y2023/vision/target_mapping.cc
+++ b/y2023/vision/target_mapping.cc
@@ -5,7 +5,6 @@
 #include "aos/util/mcap_logger.h"
 #include "frc971/control_loops/pose.h"
 #include "frc971/vision/calibration_generated.h"
-#include "frc971/vision/charuco_lib.h"
 #include "frc971/vision/target_mapper.h"
 #include "frc971/vision/visualize_robot.h"
 #include "opencv2/aruco.hpp"
@@ -49,7 +48,10 @@
 DEFINE_bool(solve, true, "Whether to solve for the field's target map.");
 DEFINE_string(dump_constraints_to, "/tmp/constraints.txt",
               "Write the target constraints to this path");
-DECLARE_uint64(frozen_target_id);
+DECLARE_int32(frozen_target_id);
+DECLARE_int32(min_target_id);
+DECLARE_int32(max_target_id);
+DECLARE_bool(visualize_solver);
 
 namespace y2023 {
 namespace vision {
@@ -61,9 +63,6 @@
 using frc971::vision::VisualizeRobot;
 namespace calibration = frc971::vision::calibration;
 
-constexpr TargetMapper::TargetId kMinTargetId = 1;
-constexpr TargetMapper::TargetId kMaxTargetId = 8;
-
 // Class to handle reading target poses from a replayed log,
 // displaying various debug info, and passing the poses to
 // frc971::vision::TargetMapper for field mapping.
@@ -203,7 +202,7 @@
         });
   }
 
-  if (FLAGS_visualize) {
+  if (FLAGS_visualize_solver) {
     vis_robot_.ClearImage();
     const double kFocalLength = 500.0;
     vis_robot_.SetDefaultViewpoint(kImageWidth, kFocalLength);
@@ -229,22 +228,23 @@
 
   for (const auto *target_pose_fbs : *map.target_poses()) {
     // Skip detections with invalid ids
-    if (target_pose_fbs->id() < kMinTargetId ||
-        target_pose_fbs->id() > kMaxTargetId) {
-      LOG(WARNING) << "Skipping tag with invalid id of "
-                   << target_pose_fbs->id();
+    if (static_cast<TargetMapper::TargetId>(target_pose_fbs->id()) <
+            FLAGS_min_target_id ||
+        static_cast<TargetMapper::TargetId>(target_pose_fbs->id()) >
+            FLAGS_max_target_id) {
+      VLOG(1) << "Skipping tag with invalid id of " << target_pose_fbs->id();
       continue;
     }
 
     // Skip detections with high pose errors
     if (target_pose_fbs->pose_error() > FLAGS_max_pose_error) {
-      VLOG(1) << " Skipping tag " << target_pose_fbs->id()
+      VLOG(1) << "Skipping tag " << target_pose_fbs->id()
               << " due to pose error of " << target_pose_fbs->pose_error();
       continue;
     }
     // Skip detections with high pose error ratios
     if (target_pose_fbs->pose_error_ratio() > FLAGS_max_pose_error_ratio) {
-      VLOG(1) << " Skipping tag " << target_pose_fbs->id()
+      VLOG(1) << "Skipping tag " << target_pose_fbs->id()
               << " due to pose error ratio of "
               << target_pose_fbs->pose_error_ratio();
       continue;
@@ -274,7 +274,7 @@
             .distortion_factor = distortion_factor,
             .id = static_cast<TargetMapper::TargetId>(target_pose.id)});
 
-    if (FLAGS_visualize) {
+    if (FLAGS_visualize_solver) {
       // If we've already drawn in the current image,
       // display it before clearing and adding the new poses
       if (drawn_nodes_.count(node_name) != 0) {
@@ -419,10 +419,13 @@
                 .pose),
         .distance_from_camera = kSeedDistanceFromCamera,
         .distortion_factor = kSeedDistortionFactor,
-        .id = kMinTargetId};
+        .id = FLAGS_frozen_target_id};
 
-    for (TargetMapper::TargetId id = kMinTargetId; id <= kMaxTargetId; id++) {
-      if (id == static_cast<TargetMapper::TargetId>(FLAGS_frozen_target_id)) {
+    constexpr TargetMapper::TargetId kAbsMinTargetId = 1;
+    constexpr TargetMapper::TargetId kAbsMaxTargetId = 8;
+    for (TargetMapper::TargetId id = kAbsMinTargetId; id <= kAbsMaxTargetId;
+         id++) {
+      if (id == FLAGS_frozen_target_id) {
         continue;
       }