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