Log charuco corner annotations and images in calibration_accumulator

This makes it so that we can readily use foxglove to visualize at least
the charuco detections in a full log being used for extrinsics
calibration. For future changes we will pull in more.

Change-Id: I3416350698d455c7443c2a6f17867614d2e77c0c
Signed-off-by: James Kuszmaul <jabukuszmaul@gmail.com>
diff --git a/frc971/vision/BUILD b/frc971/vision/BUILD
index c52afae..308792c 100644
--- a/frc971/vision/BUILD
+++ b/frc971/vision/BUILD
@@ -102,6 +102,7 @@
         "//y2020/vision/sift:sift_fbs",
         "//y2020/vision/sift:sift_training_fbs",
         "//y2020/vision/tools/python_code:sift_training_data",
+        "@com_github_foxglove_schemas//:schemas",
         "@com_github_google_glog//:glog",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
@@ -121,6 +122,7 @@
     visibility = ["//visibility:public"],
     deps = [
         ":charuco_lib",
+        ":foxglove_image_converter",
         "//aos:init",
         "//aos/events/logging:log_reader",
         "//frc971/analysis:in_process_plotter",
@@ -129,6 +131,8 @@
         "//frc971/wpilib:imu_batch_fbs",
         "//frc971/wpilib:imu_fbs",
         "//third_party:opencv",
+        "@com_github_foxglove_schemas//:CompressedImage_schema",
+        "@com_github_foxglove_schemas//:ImageAnnotations_schema",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_ceres_solver//:ceres",
         "@org_tuxfamily_eigen//:eigen",
diff --git a/frc971/vision/calibration_accumulator.cc b/frc971/vision/calibration_accumulator.cc
index eee22f5..ea9af28 100644
--- a/frc971/vision/calibration_accumulator.cc
+++ b/frc971/vision/calibration_accumulator.cc
@@ -11,6 +11,8 @@
 #include "frc971/control_loops/quaternion_utils.h"
 #include "frc971/vision/charuco_lib.h"
 #include "frc971/wpilib/imu_batch_generated.h"
+#include "external/com_github_foxglove_schemas/ImageAnnotations_schema.h"
+#include "external/com_github_foxglove_schemas/CompressedImage_schema.h"
 
 DEFINE_bool(display_undistorted, false,
             "If true, display the undistorted image.");
@@ -111,6 +113,33 @@
   }
 }
 
+CalibrationFoxgloveVisualizer::CalibrationFoxgloveVisualizer(
+    aos::EventLoop *event_loop)
+    : event_loop_(event_loop),
+      image_converter_(event_loop_, "/camera", "/visualization",
+                       ImageCompression::kJpeg),
+      annotations_sender_(
+          event_loop_->MakeSender<foxglove::ImageAnnotations>("/visualization")) {}
+
+aos::FlatbufferDetachedBuffer<aos::Configuration>
+CalibrationFoxgloveVisualizer::AddVisualizationChannels(
+    const aos::Configuration *config, const aos::Node *node) {
+  constexpr std::string_view channel_name = "/visualization";
+  aos::ChannelT channel_overrides;
+  channel_overrides.max_size = 10000000;
+  aos::FlatbufferDetachedBuffer<aos::Configuration> result =
+      aos::configuration::AddChannelToConfiguration(
+          config, channel_name,
+          aos::FlatbufferSpan<reflection::Schema>(
+              foxglove::ImageAnnotationsSchema()),
+          node, channel_overrides);
+  return aos::configuration::AddChannelToConfiguration(
+      &result.message(), channel_name,
+      aos::FlatbufferSpan<reflection::Schema>(
+          foxglove::CompressedImageSchema()),
+      node, channel_overrides);
+}
+
 Calibration::Calibration(aos::SimulatedEventLoopFactory *event_loop_factory,
                          aos::EventLoop *image_event_loop,
                          aos::EventLoop *imu_event_loop, std::string_view pi,
@@ -140,7 +169,9 @@
           [this](cv::Mat rgb_image, const monotonic_clock::time_point eof) {
             charuco_extractor_.HandleImage(rgb_image, eof);
           }),
-      data_(data) {
+      data_(data),
+      visualizer_event_loop_(image_factory_->MakeEventLoop("visualization")),
+      visualizer_(visualizer_event_loop_.get()) {
   imu_factory_->OnShutdown([]() { cv::destroyAllWindows(); });
 
   // Check for IMUValuesBatch topic on both /localizer and /drivetrain channels,
@@ -173,9 +204,10 @@
 void Calibration::HandleCharuco(
     cv::Mat rgb_image, const monotonic_clock::time_point eof,
     std::vector<cv::Vec4i> /*charuco_ids*/,
-    std::vector<std::vector<cv::Point2f>> /*charuco_corners*/, bool valid,
+    std::vector<std::vector<cv::Point2f>> charuco_corners, bool valid,
     std::vector<Eigen::Vector3d> rvecs_eigen,
     std::vector<Eigen::Vector3d> tvecs_eigen) {
+  visualizer_.HandleCharuco(eof, charuco_corners);
   if (valid) {
     CHECK(rvecs_eigen.size() > 0) << "Require at least one target detected";
     // We only use one (the first) target detected for calibration
@@ -237,6 +269,7 @@
                         last_value_.accelerometer_y,
                         last_value_.accelerometer_z);
 
+  // TODO: ToDistributedClock may be too noisy.
   data_->AddImu(imu_factory_->ToDistributedClock(monotonic_clock::time_point(
                     chrono::nanoseconds(imu->monotonic_timestamp_ns()))),
                 gyro, accel * kG);
diff --git a/frc971/vision/calibration_accumulator.h b/frc971/vision/calibration_accumulator.h
index d9f6065..5c435ad 100644
--- a/frc971/vision/calibration_accumulator.h
+++ b/frc971/vision/calibration_accumulator.h
@@ -8,6 +8,7 @@
 #include "aos/time/time.h"
 #include "frc971/control_loops/quaternion_utils.h"
 #include "frc971/vision/charuco_lib.h"
+#include "frc971/vision/foxglove_image_converter.h"
 #include "frc971/wpilib/imu_batch_generated.h"
 
 namespace frc971 {
@@ -76,6 +77,28 @@
       turret_points_;
 };
 
+class CalibrationFoxgloveVisualizer {
+ public:
+  CalibrationFoxgloveVisualizer(aos::EventLoop *event_loop);
+
+  static aos::FlatbufferDetachedBuffer<aos::Configuration>
+  AddVisualizationChannels(const aos::Configuration *config,
+                           const aos::Node *node);
+
+  void HandleCharuco(const aos::monotonic_clock::time_point eof,
+                     std::vector<std::vector<cv::Point2f>> charuco_corners) {
+    auto builder = annotations_sender_.MakeBuilder();
+    builder.CheckOk(
+        builder.Send(BuildAnnotations(eof, charuco_corners, builder.fbb())));
+  }
+
+ private:
+  aos::EventLoop *event_loop_;
+  FoxgloveImageConverter image_converter_;
+
+  aos::Sender<foxglove::ImageAnnotations> annotations_sender_;
+};
+
 // Class to register image and IMU callbacks in AOS and route them to the
 // corresponding CalibrationData class.
 class Calibration {
@@ -110,6 +133,9 @@
 
   CalibrationData *data_;
 
+  std::unique_ptr<aos::EventLoop> visualizer_event_loop_;
+  CalibrationFoxgloveVisualizer visualizer_;
+
   frc971::IMUValuesT last_value_;
 };
 
diff --git a/frc971/vision/charuco_lib.cc b/frc971/vision/charuco_lib.cc
index 80b1fd4..864c5e7 100644
--- a/frc971/vision/charuco_lib.cc
+++ b/frc971/vision/charuco_lib.cc
@@ -487,5 +487,42 @@
                   rvecs_eigen, tvecs_eigen);
 }
 
+flatbuffers::Offset<foxglove::ImageAnnotations> BuildAnnotations(
+    const aos::monotonic_clock::time_point monotonic_now,
+    const std::vector<std::vector<cv::Point2f>> &corners,
+    flatbuffers::FlatBufferBuilder *fbb) {
+  std::vector<flatbuffers::Offset<foxglove::PointsAnnotation>> rectangles;
+  const struct timespec now_t = aos::time::to_timespec(monotonic_now);
+  foxglove::Time time{static_cast<uint32_t>(now_t.tv_sec),
+                      static_cast<uint32_t>(now_t.tv_nsec)};
+  const flatbuffers::Offset<foxglove::Color> color_offset =
+      foxglove::CreateColor(*fbb, 0.0, 1.0, 0.0, 1.0);
+  for (const std::vector<cv::Point2f> &rectangle : corners) {
+    std::vector<flatbuffers::Offset<foxglove::Point2>> points_offsets;
+    for (const cv::Point2f &point : rectangle) {
+      points_offsets.push_back(foxglove::CreatePoint2(*fbb, point.x, point.y));
+    }
+    const flatbuffers::Offset<
+        flatbuffers::Vector<flatbuffers::Offset<foxglove::Point2>>>
+        points_offset = fbb->CreateVector(points_offsets);
+    std::vector<flatbuffers::Offset<foxglove::Color>> color_offsets(
+        points_offsets.size(), color_offset);
+    auto colors_offset = fbb->CreateVector(color_offsets);
+    foxglove::PointsAnnotation::Builder points_builder(*fbb);
+    points_builder.add_timestamp(&time);
+    points_builder.add_type(foxglove::PointsAnnotationType::POINTS);
+    points_builder.add_points(points_offset);
+    points_builder.add_outline_color(color_offset);
+    points_builder.add_outline_colors(colors_offset);
+    points_builder.add_thickness(2.0);
+    rectangles.push_back(points_builder.Finish());
+  }
+
+  const auto rectangles_offset = fbb->CreateVector(rectangles);
+  foxglove::ImageAnnotations::Builder annotation_builder(*fbb);
+  annotation_builder.add_points(rectangles_offset);
+  return annotation_builder.Finish();
+}
+
 }  // namespace vision
 }  // namespace frc971
diff --git a/frc971/vision/charuco_lib.h b/frc971/vision/charuco_lib.h
index c7269f9..e296071 100644
--- a/frc971/vision/charuco_lib.h
+++ b/frc971/vision/charuco_lib.h
@@ -13,6 +13,7 @@
 #include "aos/network/message_bridge_server_generated.h"
 #include "y2020/vision/sift/sift_generated.h"
 #include "y2020/vision/sift/sift_training_generated.h"
+#include "external/com_github_foxglove_schemas/ImageAnnotations_generated.h"
 
 DECLARE_bool(visualize);
 
@@ -175,6 +176,13 @@
       handle_charuco_;
 };
 
+// Puts the provided charuco corners into a foxglove ImageAnnotation type for
+// visualization purposes.
+flatbuffers::Offset<foxglove::ImageAnnotations> BuildAnnotations(
+    const aos::monotonic_clock::time_point monotonic_now,
+    const std::vector<std::vector<cv::Point2f>> &corners,
+    flatbuffers::FlatBufferBuilder *fbb);
+
 }  // namespace vision
 }  // namespace frc971
 
diff --git a/y2022/vision/calibrate_extrinsics.cc b/y2022/vision/calibrate_extrinsics.cc
index 05d3481..9b6c428 100644
--- a/y2022/vision/calibrate_extrinsics.cc
+++ b/y2022/vision/calibrate_extrinsics.cc
@@ -2,6 +2,7 @@
 #include "Eigen/Geometry"
 #include "absl/strings/str_format.h"
 #include "aos/events/logging/log_reader.h"
+#include "aos/events/logging/log_writer.h"
 #include "aos/init.h"
 #include "aos/network/team_number.h"
 #include "aos/time/time.h"
@@ -21,6 +22,8 @@
 DEFINE_string(target_type, "charuco",
               "Type of target: april_tag|aruco|charuco|charuco_diamond");
 DEFINE_string(image_channel, "/camera", "Channel to listen for images on");
+DEFINE_string(output_logs, "/tmp/calibration/",
+              "Output folder for visualization logs.");
 
 namespace frc971 {
 namespace vision {
@@ -33,11 +36,24 @@
 
 void Main(int argc, char **argv) {
   CalibrationData data;
+  std::optional<uint16_t> pi_number = aos::network::ParsePiNumber(FLAGS_pi);
+  CHECK(pi_number);
+  const std::string pi_name = absl::StrCat("pi", *pi_number);
+  LOG(INFO) << "Pi " << *pi_number;
+  aos::FlatbufferDetachedBuffer<aos::Configuration> config = [argc, argv,
+                                                              pi_name]() {
+    aos::logger::LogReader reader(
+        aos::logger::SortParts(aos::logger::FindLogs(argc, argv)));
+    return CalibrationFoxgloveVisualizer::AddVisualizationChannels(
+        reader.logged_configuration(),
+        aos::configuration::GetNode(reader.logged_configuration(), pi_name));
+  }();
 
   {
     // Now, accumulate all the data into the data object.
     aos::logger::LogReader reader(
-        aos::logger::SortParts(aos::logger::FindLogs(argc, argv)));
+        aos::logger::SortParts(aos::logger::FindLogs(argc, argv)),
+        &config.message());
 
     aos::SimulatedEventLoopFactory factory(reader.configuration());
     reader.Register(&factory);
@@ -50,11 +66,8 @@
     const aos::Node *const roborio_node =
         aos::configuration::GetNode(factory.configuration(), "roborio");
 
-    std::optional<uint16_t> pi_number = aos::network::ParsePiNumber(FLAGS_pi);
-    CHECK(pi_number);
-    LOG(INFO) << "Pi " << *pi_number;
-    const aos::Node *const pi_node = aos::configuration::GetNode(
-        factory.configuration(), absl::StrCat("pi", *pi_number));
+    const aos::Node *const pi_node =
+        aos::configuration::GetNode(factory.configuration(), pi_name);
 
     LOG(INFO) << "imu " << aos::FlatbufferToJson(imu_node);
     LOG(INFO) << "roboRIO " << aos::FlatbufferToJson(roborio_node);
@@ -67,6 +80,11 @@
     std::unique_ptr<aos::EventLoop> pi_event_loop =
         factory.MakeEventLoop("calibration", pi_node);
 
+    std::unique_ptr<aos::EventLoop> logger_loop =
+        factory.MakeEventLoop("logger", pi_node);
+    aos::logger::Logger logger(logger_loop.get());
+    logger.StartLoggingOnRun(FLAGS_output_logs);
+
     TargetType target_type = TargetType::kCharuco;
     if (FLAGS_target_type == "april_tag") {
       target_type = TargetType::kAprilTag;