Adding visualization tools for target_mapping and logging playback

Helps to see where the targets are being seen, and the respective localizations
Added ability to draw robot frame on visualizer

Change-Id: I8af7a6d84874fe626d8dc9452f77702741e72fbb
Signed-off-by: Jim Ostrowski <yimmy13@gmail.com>
diff --git a/frc971/vision/BUILD b/frc971/vision/BUILD
index f3ebecb..133281b 100644
--- a/frc971/vision/BUILD
+++ b/frc971/vision/BUILD
@@ -223,6 +223,7 @@
     hdrs = [
         "visualize_robot.h",
     ],
+    visibility = ["//visibility:public"],
     deps = [
         "//aos:init",
         "//third_party:opencv",
diff --git a/frc971/vision/target_mapper.cc b/frc971/vision/target_mapper.cc
index dd332fc..d0dcdc8 100644
--- a/frc971/vision/target_mapper.cc
+++ b/frc971/vision/target_mapper.cc
@@ -232,6 +232,15 @@
   return std::nullopt;
 }
 
+std::optional<TargetMapper::TargetPose> TargetMapper::GetTargetPoseById(
+    TargetId target_id) {
+  if (target_poses_.count(target_id) > 0) {
+    return TargetMapper::TargetPose{target_id, target_poses_[target_id]};
+  }
+
+  return std::nullopt;
+}
+
 // Taken from ceres/examples/slam/pose_graph_3d/pose_graph_3d.cc
 // Constructs the nonlinear least squares optimization problem from the pose
 // graph constraints.
@@ -290,6 +299,9 @@
   // better to properly constrain the gauge freedom. This can be done by
   // setting one of the poses as constant so the optimizer cannot change it.
   ceres::examples::MapOfPoses::iterator pose_start_iter = poses->begin();
+  // TODO<Jim>: This fixes first target, but breaks optimizer if we don't have
+  // an observation on the first target.  We may want to allow other targets as
+  // fixed
   CHECK(pose_start_iter != poses->end()) << "There are no poses.";
   problem->SetParameterBlockConstant(pose_start_iter->second.p.data());
   problem->SetParameterBlockConstant(pose_start_iter->second.q.coeffs().data());
@@ -302,6 +314,7 @@
   ceres::Solver::Options options;
   options.max_num_iterations = FLAGS_max_num_iterations;
   options.linear_solver_type = ceres::SPARSE_NORMAL_CHOLESKY;
+  options.minimizer_progress_to_stdout = true;
 
   ceres::Solver::Summary summary;
   ceres::Solve(options, problem, &summary);
diff --git a/frc971/vision/target_mapper.h b/frc971/vision/target_mapper.h
index 3724832..35c0977 100644
--- a/frc971/vision/target_mapper.h
+++ b/frc971/vision/target_mapper.h
@@ -46,6 +46,9 @@
   static std::optional<TargetPose> GetTargetPoseById(
       std::vector<TargetPose> target_poses, TargetId target_id);
 
+  // Version that gets based on internal target_poses
+  std::optional<TargetPose> GetTargetPoseById(TargetId target_id);
+
   ceres::examples::MapOfPoses target_poses() { return target_poses_; }
 
  private:
diff --git a/frc971/vision/visualize_robot.cc b/frc971/vision/visualize_robot.cc
index 0bbe507..eb2da9a 100644
--- a/frc971/vision/visualize_robot.cc
+++ b/frc971/vision/visualize_robot.cc
@@ -9,6 +9,24 @@
 namespace frc971 {
 namespace vision {
 
+void VisualizeRobot::SetDefaultViewpoint(int image_width, double focal_length) {
+  // 10 meters above the origin, rotated so the camera faces straight down
+  Eigen::Translation3d camera_trans(0, 0, 10.0);
+  Eigen::AngleAxisd camera_rot(M_PI, Eigen::Vector3d::UnitX());
+  Eigen::Affine3d camera_viewpoint = camera_trans * camera_rot;
+  SetViewpoint(camera_viewpoint);
+
+  cv::Mat camera_mat;
+  double half_width = static_cast<double>(image_width) / 2.0;
+  double intr[] = {focal_length, 0.0, half_width, 0.0, focal_length,
+                   half_width,   0.0, 0.0,        1.0};
+  camera_mat = cv::Mat(3, 3, CV_64FC1, intr);
+  SetCameraParameters(camera_mat);
+
+  cv::Mat dist_coeffs = cv::Mat(1, 5, CV_64F, 0.0);
+  SetDistortionCoefficients(dist_coeffs);
+}
+
 cv::Point VisualizeRobot::ProjectPoint(Eigen::Vector3d T_world_point) {
   // Map 3D point in world coordinates to camera frame
   Eigen::Vector3d T_camera_point = H_world_viewpoint_.inverse() * T_world_point;
@@ -28,7 +46,7 @@
   cv::Point start2d = ProjectPoint(start3d);
   cv::Point end2d = ProjectPoint(end3d);
 
-  cv::line(image_, start2d, end2d, cv::Scalar(0, 0, 255));
+  cv::line(image_, start2d, end2d, cv::Scalar(0, 200, 0));
 }
 
 void VisualizeRobot::DrawFrameAxes(Eigen::Affine3d H_world_target,
@@ -53,21 +71,37 @@
     // Grab x axis direction
     cv::Vec3d label_offset = r_mat.col(0);
 
-    // Find 3D coordinate of point at the end of the x-axis, and project it
+    // Find 3D coordinate of point at the end of the x-axis, and put label there
+    // Bump it just a few (5) pixels to the right, to make it read easier
     cv::Mat label_coord_res =
         camera_mat_ * cv::Mat(tvec + label_offset * axis_scale);
     cv::Vec3d label_coord = label_coord_res.col(0);
-    label_coord[0] = label_coord[0] / label_coord[2];
+    label_coord[0] = label_coord[0] / label_coord[2] + 5;
     label_coord[1] = label_coord[1] / label_coord[2];
     cv::putText(image_, label, cv::Point(label_coord[0], label_coord[1]),
                 cv::FONT_HERSHEY_PLAIN, 1.0, cv::Scalar(0, 0, 255));
   }
 }
 
-void VisualizeRobot::DrawBoardOutline(Eigen::Affine3d H_world_board,
+void VisualizeRobot::DrawRobotOutline(Eigen::Affine3d H_world_robot,
                                       std::string label) {
-  LOG(INFO) << "Not yet implemented; drawing axes only";
-  DrawFrameAxes(H_world_board, label);
+  DrawFrameAxes(H_world_robot, label);
+  const double kBotHalfWidth = 0.75 / 2.0;
+  const double kBotHalfLength = 1.0 / 2.0;
+  // Compute coordinates for front/rear and right/left corners
+  Eigen::Vector3d fr_corner =
+      H_world_robot * Eigen::Vector3d(kBotHalfLength, kBotHalfWidth, 0);
+  Eigen::Vector3d fl_corner =
+      H_world_robot * Eigen::Vector3d(kBotHalfLength, -kBotHalfWidth, 0);
+  Eigen::Vector3d rl_corner =
+      H_world_robot * Eigen::Vector3d(-kBotHalfLength, -kBotHalfWidth, 0);
+  Eigen::Vector3d rr_corner =
+      H_world_robot * Eigen::Vector3d(-kBotHalfLength, kBotHalfWidth, 0);
+
+  DrawLine(fr_corner, fl_corner);
+  DrawLine(fl_corner, rl_corner);
+  DrawLine(rl_corner, rr_corner);
+  DrawLine(rr_corner, fr_corner);
 }
 
 }  // namespace vision
diff --git a/frc971/vision/visualize_robot.h b/frc971/vision/visualize_robot.h
index 391a030..f5e75af 100644
--- a/frc971/vision/visualize_robot.h
+++ b/frc971/vision/visualize_robot.h
@@ -31,14 +31,20 @@
 
   // Set camera parameters (for projection into a virtual view)
   void SetCameraParameters(cv::Mat camera_intrinsics) {
-    camera_mat_ = camera_intrinsics;
+    camera_mat_ = camera_intrinsics.clone();
   }
 
   // Set distortion coefficients (defaults to 0)
   void SetDistortionCoefficients(cv::Mat dist_coeffs) {
-    dist_coeffs_ = dist_coeffs;
+    dist_coeffs_ = dist_coeffs.clone();
   }
 
+  // Sets up a default camera view 10 m above the origin, pointing down
+  // Uses image_width and focal_length to define a default camera projection
+  // matrix
+  void SetDefaultViewpoint(int image_width = 1000,
+                           double focal_length = 1000.0);
+
   // Helper function to project a point in 3D to the virtual image coordinates
   cv::Point ProjectPoint(Eigen::Vector3d point3d);
 
@@ -50,8 +56,9 @@
   void DrawFrameAxes(Eigen::Affine3d H_world_target, std::string label = "",
                      double axis_scale = 0.25);
 
-  // TODO<Jim>: Need to implement this, and maybe DrawRobotOutline
-  void DrawBoardOutline(Eigen::Affine3d H_world_board, std::string label = "");
+  // TODO<Jim>: Also implement DrawBoardOutline?  Maybe one function w/
+  // parameters?
+  void DrawRobotOutline(Eigen::Affine3d H_world_robot, std::string label = "");
 
   Eigen::Affine3d H_world_viewpoint_;  // Where to view the world from
   cv::Mat image_;                      // Image to draw on
diff --git a/frc971/vision/visualize_robot_sample.cc b/frc971/vision/visualize_robot_sample.cc
index dc38352..59acba0 100644
--- a/frc971/vision/visualize_robot_sample.cc
+++ b/frc971/vision/visualize_robot_sample.cc
@@ -26,25 +26,8 @@
   cv::Mat image_mat =
       cv::Mat::zeros(cv::Size(image_width, image_width), CV_8UC3);
   vis_robot.SetImage(image_mat);
-
-  // 10 meters above the origin, rotated so the camera faces straight down
-  Eigen::Translation3d camera_trans(0, 0, 10.0);
-  Eigen::AngleAxisd camera_rot(M_PI, Eigen::Vector3d::UnitX());
-  Eigen::Affine3d camera_viewpoint = camera_trans * camera_rot;
-  vis_robot.SetViewpoint(camera_viewpoint);
-
-  cv::Mat camera_mat;
   double focal_length = 1000.0;
-  double intr[] = {focal_length, 0.0,          image_width / 2.0,
-                   0.0,          focal_length, image_width / 2.0,
-                   0.0,          0.0,          1.0};
-  camera_mat = cv::Mat(3, 3, CV_64FC1, intr);
-  vis_robot.SetCameraParameters(camera_mat);
-
-  Eigen::Affine3d offset_rotate_origin(Eigen::Affine3d::Identity());
-
-  cv::Mat dist_coeffs = cv::Mat(1, 5, CV_64F, 0.0);
-  vis_robot.SetDistortionCoefficients(dist_coeffs);
+  vis_robot.SetDefaultViewpoint(image_width, focal_length);
 
   // Go around the clock and plot the coordinate frame at different rotations
   for (int i = 0; i < 12; i++) {
@@ -52,8 +35,9 @@
     Eigen::Vector3d trans;
     trans << 1.0 * cos(angle), 1.0 * sin(angle), 0.0;
 
-    offset_rotate_origin = Eigen::Translation3d(trans) *
-                           Eigen::AngleAxisd(angle, Eigen::Vector3d::UnitX());
+    Eigen::Affine3d offset_rotate_origin =
+        Eigen::Translation3d(trans) *
+        Eigen::AngleAxisd(angle, Eigen::Vector3d::UnitX());
 
     vis_robot.DrawFrameAxes(offset_rotate_origin, std::to_string(i));
   }
diff --git a/y2023/vision/BUILD b/y2023/vision/BUILD
index 1cdbe36..0f75aba 100644
--- a/y2023/vision/BUILD
+++ b/y2023/vision/BUILD
@@ -117,6 +117,7 @@
         "//frc971/vision:target_map_fbs",
         "//frc971/vision:target_mapper",
         "//frc971/vision:vision_fbs",
+        "//frc971/vision:visualize_robot",
         "//third_party:opencv",
         "//third_party/apriltag",
         "//y2023/constants:constants_fbs",
diff --git a/y2023/vision/aprilrobotics.cc b/y2023/vision/aprilrobotics.cc
index 95ad541..cf2244a 100644
--- a/y2023/vision/aprilrobotics.cc
+++ b/y2023/vision/aprilrobotics.cc
@@ -2,6 +2,8 @@
 
 #include "y2023/vision/vision_util.h"
 
+#include <opencv2/highgui.hpp>
+
 DEFINE_bool(
     debug, false,
     "If true, dump a ton of debug and crash on the first valid detection.");
@@ -26,13 +28,12 @@
     : calibration_data_(event_loop),
       image_size_(0, 0),
       ftrace_(),
-      image_callback_(
-          event_loop, channel_name,
-          [&](cv::Mat image_color_mat,
-              const aos::monotonic_clock::time_point eof) {
-            HandleImage(image_color_mat, eof);
-          },
-          chrono::milliseconds(5)),
+      image_callback_(event_loop, channel_name,
+                      [&](cv::Mat image_color_mat,
+                          const aos::monotonic_clock::time_point eof) {
+                        HandleImage(image_color_mat, eof);
+                      },
+                      chrono::milliseconds(5)),
       target_map_sender_(
           event_loop->MakeSender<frc971::vision::TargetMap>("/camera")),
       image_annotations_sender_(
@@ -54,7 +55,6 @@
       calibration_data_.constants(), event_loop->node()->name()->string_view());
 
   extrinsics_ = CameraExtrinsics(calibration_);
-
   intrinsics_ = CameraIntrinsics(calibration_);
   // Create an undistort projection matrix using the intrinsics
   projection_matrix_ = cv::Mat::zeros(3, 4, CV_64F);
@@ -181,6 +181,8 @@
 
 AprilRoboticsDetector::DetectionResult AprilRoboticsDetector::DetectTags(
     cv::Mat image, aos::monotonic_clock::time_point eof) {
+  cv::Mat color_image;
+  cvtColor(image, color_image, cv::COLOR_GRAY2RGB);
   const aos::monotonic_clock::time_point start_time =
       aos::monotonic_clock::now();
 
@@ -275,10 +277,42 @@
                                      .pose = pose,
                                      .pose_error = pose_error,
                                      .distortion_factor = distortion_factor});
+      if (FLAGS_visualize) {
+        // Draw raw (distorted) corner points in green
+        cv::line(color_image, orig_corner_points[0], orig_corner_points[1],
+                 cv::Scalar(0, 255, 0), 2);
+        cv::line(color_image, orig_corner_points[1], orig_corner_points[2],
+                 cv::Scalar(0, 255, 0), 2);
+        cv::line(color_image, orig_corner_points[2], orig_corner_points[3],
+                 cv::Scalar(0, 255, 0), 2);
+        cv::line(color_image, orig_corner_points[3], orig_corner_points[0],
+                 cv::Scalar(0, 255, 0), 2);
+
+        // Draw undistorted corner points in red
+        cv::line(color_image, corner_points[0], corner_points[1],
+                 cv::Scalar(0, 0, 255), 2);
+        cv::line(color_image, corner_points[2], corner_points[1],
+                 cv::Scalar(0, 0, 255), 2);
+        cv::line(color_image, corner_points[2], corner_points[3],
+                 cv::Scalar(0, 0, 255), 2);
+        cv::line(color_image, corner_points[0], corner_points[3],
+                 cv::Scalar(0, 0, 255), 2);
+      }
+
+      VLOG(1) << "Found tag number " << det->id << " hamming: " << det->hamming
+              << " margin: " << det->decision_margin;
+
     } else {
       rejections_++;
     }
   }
+  if (FLAGS_visualize) {
+    // Display the result
+    // Rotate by 180 degrees to make it upright
+    // TDOO<Jim>: Make this a flag, since we don't want it for box of pis
+    cv::rotate(color_image, color_image, 1);
+    cv::imshow("AprilRoboticsDetector Image", color_image);
+  }
 
   const auto corners_offset = builder.fbb()->CreateVector(foxglove_corners);
   foxglove::ImageAnnotations::Builder annotation_builder(*builder.fbb());
diff --git a/y2023/vision/aprilrobotics.h b/y2023/vision/aprilrobotics.h
index fab2d30..7caa848 100644
--- a/y2023/vision/aprilrobotics.h
+++ b/y2023/vision/aprilrobotics.h
@@ -11,6 +11,7 @@
 #include "frc971/vision/target_map_generated.h"
 #include "frc971/vision/target_mapper.h"
 #include "frc971/vision/vision_generated.h"
+#include "frc971/vision/visualize_robot.h"
 #include "opencv2/core/eigen.hpp"
 #include "opencv2/imgproc.hpp"
 #include "third_party/apriltag/apriltag.h"
@@ -83,8 +84,8 @@
   frc971::vision::ImageCallback image_callback_;
   aos::Sender<frc971::vision::TargetMap> target_map_sender_;
   aos::Sender<foxglove::ImageAnnotations> image_annotations_sender_;
-
   size_t rejections_;
+  frc971::vision::VisualizeRobot vis_robot_;
 };
 
 }  // namespace vision
diff --git a/y2023/vision/target_mapping.cc b/y2023/vision/target_mapping.cc
index 9352bdd..a110794 100644
--- a/y2023/vision/target_mapping.cc
+++ b/y2023/vision/target_mapping.cc
@@ -7,6 +7,7 @@
 #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"
 #include "opencv2/calib3d.hpp"
 #include "opencv2/core/eigen.hpp"
@@ -28,8 +29,8 @@
               "Directory to write solved target map to");
 DEFINE_string(field_name, "charged_up",
               "Field name, for the output json filename and flatbuffer field");
-DEFINE_int32(team_number, 7971,
-             "Use the calibration for a node with this team number");
+DEFINE_int32(team_number, 0,
+             "Required: Use the calibration for a node with this team number");
 DEFINE_string(mcap_output_path, "/tmp/log.mcap", "Log to output.");
 DEFINE_string(pi, "pi1", "Pi name to generate mcap log for; defaults to pi1.");
 DEFINE_double(max_pose_error, 1e-6,
@@ -52,16 +53,22 @@
   return H_robot_target;
 }
 
+frc971::vision::VisualizeRobot vis_robot_;
+const int kImageWidth = 1000;
+
 // Add detected apriltag poses relative to the robot to
 // timestamped_target_detections
 void HandleAprilTags(const TargetMap &map,
                      aos::distributed_clock::time_point pi_distributed_time,
                      std::vector<DataAdapter::TimestampedDetection>
                          *timestamped_target_detections,
-                     Eigen::Affine3d extrinsics) {
+                     Eigen::Affine3d extrinsics, std::string node_name,
+                     frc971::vision::TargetMapper mapper) {
   for (const auto *target_pose_fbs : *map.target_poses()) {
     // Skip detections with high pose errors
     if (target_pose_fbs->pose_error() > FLAGS_max_pose_error) {
+      LOG(INFO) << " Skipping tag " << target_pose_fbs->id()
+                << " due to pose error of " << target_pose_fbs->pose_error();
       continue;
     }
 
@@ -88,6 +95,27 @@
             .distance_from_camera = distance_from_camera,
             .distortion_factor = distortion_factor,
             .id = static_cast<TargetMapper::TargetId>(target_pose.id)});
+
+    if (FLAGS_visualize) {
+      Eigen::Affine3d H_world_target = PoseUtils::Pose3dToAffine3d(
+          mapper.GetTargetPoseById(target_pose_fbs->id())->pose);
+      Eigen::Affine3d H_world_robot = H_world_target * H_robot_target.inverse();
+      Eigen::Affine3d H_world_camera =
+          H_world_target * H_camera_target.inverse();
+      LOG(INFO) << node_name << ", " << target_pose_fbs->id()
+                << ", t = " << pi_distributed_time
+                << ", pose_error = " << target_pose_fbs->pose_error()
+                << ", robot_pos (x,y,z) + "
+                << H_world_robot.translation().transpose();
+      vis_robot_.DrawRobotOutline(
+          H_world_robot, node_name + "-" +
+                             std::to_string(target_pose_fbs->id()) + "  " +
+                             std::to_string(target_pose_fbs->pose_error() /
+                                            FLAGS_max_pose_error));
+      vis_robot_.DrawFrameAxes(H_world_camera, node_name);
+      vis_robot_.DrawFrameAxes(H_world_target,
+                               std::to_string(target_pose_fbs->id()));
+    }
   }
 }
 
@@ -109,6 +137,10 @@
 
   detectors->emplace_back(std::move(detector_ptr));
 
+  ceres::examples::VectorOfConstraints target_constraints;
+  frc971::vision::TargetMapper target_loc_mapper(FLAGS_json_path,
+                                                 target_constraints);
+
   mapping_event_loop->MakeWatcher("/camera", [=](const TargetMap &map) {
     aos::distributed_clock::time_point pi_distributed_time =
         reader->event_loop_factory()
@@ -116,8 +148,16 @@
             ->ToDistributedClock(aos::monotonic_clock::time_point(
                 aos::monotonic_clock::duration(map.monotonic_timestamp_ns())));
 
+    std::string node_name = detection_event_loop->node()->name()->str();
     HandleAprilTags(map, pi_distributed_time, timestamped_target_detections,
-                    extrinsics);
+                    extrinsics, node_name, target_loc_mapper);
+    if (FLAGS_visualize) {
+      cv::imshow("View", vis_robot_.image_);
+      cv::waitKey();
+      cv::Mat image_mat =
+          cv::Mat::zeros(cv::Size(kImageWidth, kImageWidth), CV_8UC3);
+      vis_robot_.SetImage(image_mat);
+    }
   });
 }
 
@@ -197,7 +237,6 @@
   std::unique_ptr<aos::McapLogger> relogger;
   if (!FLAGS_mcap_output_path.empty()) {
     LOG(INFO) << "Writing out mcap file to " << FLAGS_mcap_output_path;
-    // TODO: Should make this work for any pi
     const aos::Node *node =
         aos::configuration::GetNode(reader.configuration(), FLAGS_pi);
     reader.event_loop_factory()->GetNodeEventLoopFactory(node)->OnStartup(
@@ -212,6 +251,14 @@
         });
   }
 
+  if (FLAGS_visualize) {
+    cv::Mat image_mat =
+        cv::Mat::zeros(cv::Size(kImageWidth, kImageWidth), CV_8UC3);
+    vis_robot_.SetImage(image_mat);
+    const double kFocalLength = 500.0;
+    vis_robot_.SetDefaultViewpoint(kImageWidth, kFocalLength);
+  }
+
   reader.event_loop_factory()->Run();
 
   auto target_constraints =