Extract common charuco detection code into charuco_lib

This lets us reuse all the charuco extraction code when we do extrinsics
calibration soon.  We really just want a callback with the detected
board.

Change-Id: If806b026cc67bf9e7b61935f93b1902acdcb2783
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/y2020/vision/BUILD b/y2020/vision/BUILD
index 187dedc..1b5fce7 100644
--- a/y2020/vision/BUILD
+++ b/y2020/vision/BUILD
@@ -87,6 +87,33 @@
     ],
 )
 
+cc_library(
+    name = "charuco_lib",
+    srcs = [
+        "charuco_lib.cc",
+    ],
+    hdrs = [
+        "charuco_lib.h",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//y2020:__subpackages__"],
+    deps = [
+        ":vision_fbs",
+        "//aos:flatbuffers",
+        "//aos/events:event_loop",
+        "//aos/network:team_number",
+        "//frc971/control_loops:quaternion_utils",
+        "//third_party:opencv",
+        "//y2020/vision/sift:sift_fbs",
+        "//y2020/vision/sift:sift_training_fbs",
+        "//y2020/vision/tools/python_code:sift_training_data",
+        "@com_github_google_glog//:glog",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/types:span",
+        "@org_tuxfamily_eigen//:eigen",
+    ],
+)
+
 cc_binary(
     name = "calibration",
     srcs = [
@@ -98,10 +125,13 @@
     target_compatible_with = ["@platforms//os:linux"],
     visibility = ["//y2020:__subpackages__"],
     deps = [
+        ":charuco_lib",
         ":vision_fbs",
         "//aos:init",
         "//aos/events:shm_event_loop",
         "//frc971/control_loops/drivetrain:improved_down_estimator",
+        "//frc971/wpilib:imu_batch_fbs",
+        "//frc971/wpilib:imu_fbs",
         "//third_party:opencv",
         "//y2020/vision/sift:sift_fbs",
         "//y2020/vision/sift:sift_training_fbs",
diff --git a/y2020/vision/calibration.cc b/y2020/vision/calibration.cc
index 68002db..517e65c 100644
--- a/y2020/vision/calibration.cc
+++ b/y2020/vision/calibration.cc
@@ -1,332 +1,184 @@
-// These need to come before opencv, or they don't compile. Presumably opencv
-// #defines something annoying.
-// clang-format off
 #include "Eigen/Dense"
 #include "Eigen/Geometry"
-// clang-format on
 
-#include <opencv2/aruco/charuco.hpp>
 #include <opencv2/calib3d.hpp>
-#include <opencv2/core/eigen.hpp>
-#include <opencv2/features2d.hpp>
 #include <opencv2/highgui/highgui.hpp>
 #include <opencv2/imgproc.hpp>
 
 #include "absl/strings/str_format.h"
+#include "y2020/vision/charuco_lib.h"
 #include "aos/events/shm_event_loop.h"
 #include "aos/init.h"
 #include "aos/network/team_number.h"
 #include "aos/time/time.h"
 #include "aos/util/file.h"
-#include "frc971/control_loops/drivetrain/improved_down_estimator.h"
-#include "y2020/vision/sift/sift_generated.h"
-#include "y2020/vision/sift/sift_training_generated.h"
-#include "y2020/vision/tools/python_code/sift_training_data.h"
-#include "y2020/vision/vision_generated.h"
 
 DEFINE_string(config, "config.json", "Path to the config file to use.");
 DEFINE_string(pi, "pi-7971-1", "Pi name to calibrate.");
 DEFINE_string(calibration_folder, "y2020/vision/tools/python_code/calib_files",
               "Folder to place calibration files.");
-DEFINE_bool(large_board, true, "If true, use the large calibration board.");
-DEFINE_bool(coarse_pattern, true, "If true, use coarse arucos; else, use fine");
 DEFINE_bool(display_undistorted, false,
             "If true, display the undistorted image.");
-DEFINE_string(board_template_path, "",
-              "If specified, write an image to the specified path for the "
-              "charuco board pattern.");
-DEFINE_uint32(min_targets, 10,
-              "The mininum number of targets required to match.");
 
 namespace frc971 {
 namespace vision {
 
-class CameraCalibration {
+class Calibration {
  public:
-  CameraCalibration(const absl::Span<const uint8_t> training_data_bfbs) {
-    const aos::FlatbufferSpan<sift::TrainingData> training_data(
-        training_data_bfbs);
-    CHECK(training_data.Verify());
-    camera_calibration_ = FindCameraCalibration(&training_data.message());
+  Calibration(aos::ShmEventLoop *event_loop, std::string_view pi)
+      : event_loop_(event_loop),
+        pi_(pi),
+        pi_number_(aos::network::ParsePiNumber(pi)),
+        charuco_extractor_(
+            event_loop, pi,
+            [this](cv::Mat rgb_image, const double age_double,
+                   std::vector<int> charuco_ids,
+                   std::vector<cv::Point2f> charuco_corners, bool valid,
+                   Eigen::Vector3d rvec_eigen, Eigen::Vector3d tvec_eigen) {
+              HandleCharuco(rgb_image, age_double, charuco_ids, charuco_corners,
+                            valid, rvec_eigen, tvec_eigen);
+            }) {
+    CHECK(pi_number_) << ": Invalid pi number " << pi
+                      << ", failed to parse pi number";
   }
 
-  cv::Mat CameraIntrinsics() const {
-    const cv::Mat result(3, 3, CV_32F,
-                         const_cast<void *>(static_cast<const void *>(
-                             camera_calibration_->intrinsics()->data())));
-    CHECK_EQ(result.total(), camera_calibration_->intrinsics()->size());
-    return result;
-  }
-
-  cv::Mat CameraDistCoeffs() const {
-    const cv::Mat result(5, 1, CV_32F,
-                         const_cast<void *>(static_cast<const void *>(
-                             camera_calibration_->dist_coeffs()->data())));
-    CHECK_EQ(result.total(), camera_calibration_->dist_coeffs()->size());
-    return result;
-  }
-
- private:
-  const sift::CameraCalibration *FindCameraCalibration(
-      const sift::TrainingData *const training_data) const {
-    std::optional<uint16_t> pi_number = aos::network::ParsePiNumber(FLAGS_pi);
-    std::optional<uint16_t> team_number =
-        aos::network::team_number_internal::ParsePiTeamNumber(FLAGS_pi);
-    CHECK(pi_number);
-    CHECK(team_number);
-    const std::string node_name = absl::StrFormat("pi%d", pi_number.value());
-    LOG(INFO) << "Looking for node name " << node_name << " team number "
-              << team_number.value();
-    for (const sift::CameraCalibration *candidate :
-         *training_data->camera_calibrations()) {
-      if (candidate->node_name()->string_view() != node_name) {
-        continue;
-      }
-      if (candidate->team_number() != team_number.value()) {
-        continue;
-      }
-      return candidate;
-    }
-    LOG(FATAL) << ": Failed to find camera calibration for " << node_name
-               << " on " << team_number.value();
-  }
-
-  const sift::CameraCalibration *camera_calibration_;
-};
-
-namespace {
-
-void ViewerMain() {
-  aos::FlatbufferDetachedBuffer<aos::Configuration> config =
-      aos::configuration::ReadConfig(FLAGS_config);
-
-  aos::ShmEventLoop event_loop(&config.message());
-
-  const CameraCalibration calibration(SiftTrainingData());
-
-  cv::Ptr<cv::aruco::Dictionary> dictionary =
-      cv::aruco::getPredefinedDictionary(FLAGS_large_board
-                                             ? cv::aruco::DICT_5X5_250
-                                             : cv::aruco::DICT_6X6_250);
-
-  LOG(INFO) << "Using " << (FLAGS_large_board ? "large" : "small")
-            << " board with " << (FLAGS_coarse_pattern ? "coarse" : "fine")
-            << " pattern";
-  cv::Ptr<cv::aruco::CharucoBoard> board =
-      FLAGS_large_board
-          ? (FLAGS_coarse_pattern ? cv::aruco::CharucoBoard::create(
-                                        12, 9, 0.06, 0.04666, dictionary)
-                                  : cv::aruco::CharucoBoard::create(
-                                        25, 18, 0.03, 0.0233, dictionary))
-          : (FLAGS_coarse_pattern ? cv::aruco::CharucoBoard::create(
-                                        7, 5, 0.04, 0.025, dictionary)
-                                  // TODO: Need to figure out what size is
-                                  // for small board, fine pattern
-                                  : cv::aruco::CharucoBoard::create(
-                                        7, 5, 0.03, 0.0233, dictionary));
-
-  if (!FLAGS_board_template_path.empty()) {
-    cv::Mat board_image;
-    board->draw(cv::Size(600, 500), board_image, 10, 1);
-    cv::imwrite(FLAGS_board_template_path, board_image);
-  }
-
-  std::vector<std::vector<int>> all_charuco_ids;
-  std::vector<std::vector<cv::Point2f>> all_charuco_corners;
-
-  const cv::Mat camera_matrix = calibration.CameraIntrinsics();
-  Eigen::Matrix3d eigen_camera_matrix;
-  cv::cv2eigen(camera_matrix, eigen_camera_matrix);
-
-  const cv::Mat dist_coeffs = calibration.CameraDistCoeffs();
-  LOG(INFO) << "Camera matrix " << camera_matrix;
-  LOG(INFO) << "Distortion Coefficients " << dist_coeffs;
-
-  std::optional<uint16_t> pi_number = aos::network::ParsePiNumber(FLAGS_pi);
-  CHECK(pi_number) << ": Invalid pi number " << FLAGS_pi
-                   << ", failed to parse pi number";
-
-  const std::string channel =
-      absl::StrCat("/pi", std::to_string(pi_number.value()), "/camera");
-  LOG(INFO) << "Connecting to channel " << channel;
-
-  event_loop.MakeWatcher(channel, [&event_loop, &board, &all_charuco_ids,
-                                   &all_charuco_corners, camera_matrix,
-                                   dist_coeffs, eigen_camera_matrix](
-                                      const CameraImage &image) {
-    const aos::monotonic_clock::duration age =
-        event_loop.monotonic_now() - event_loop.context().monotonic_event_time;
-    const double age_double =
-        std::chrono::duration_cast<std::chrono::duration<double>>(age).count();
-    if (age > std::chrono::milliseconds(100)) {
-      LOG(INFO) << "Age: " << age_double << ", getting behind, skipping";
-      return;
-    }
-    // Create color image:
-    cv::Mat image_color_mat(cv::Size(image.cols(), image.rows()), CV_8UC2,
-                            (void *)image.data()->data());
-    const cv::Size image_size(image.cols(), image.rows());
-    cv::Mat rgb_image(image_size, CV_8UC3);
-    cv::cvtColor(image_color_mat, rgb_image, CV_YUV2BGR_YUYV);
-
-    std::vector<int> marker_ids;
-    std::vector<std::vector<cv::Point2f>> marker_corners;
-
-    cv::aruco::detectMarkers(rgb_image, board->dictionary, marker_corners,
-                             marker_ids);
-
-    std::vector<cv::Point2f> charuco_corners;
-    std::vector<int> charuco_ids;
-    // If at least one marker detected
-    if (marker_ids.size() >= FLAGS_min_targets) {
-      // Run everything twice, once with the calibration, and once without.
-      // This lets us both calibrate, and also print out the pose real time
-      // with the previous calibration.
-      cv::aruco::interpolateCornersCharuco(marker_corners, marker_ids,
-                                           rgb_image, board, charuco_corners,
-                                           charuco_ids);
-
-      std::vector<cv::Point2f> charuco_corners_with_calibration;
-      std::vector<int> charuco_ids_with_calibration;
-
-      cv::aruco::interpolateCornersCharuco(
-          marker_corners, marker_ids, rgb_image, board,
-          charuco_corners_with_calibration, charuco_ids_with_calibration,
-          camera_matrix, dist_coeffs);
-
-      cv::aruco::drawDetectedMarkers(rgb_image, marker_corners, marker_ids);
-
-      if (charuco_ids.size() >= FLAGS_min_targets) {
-        cv::aruco::drawDetectedCornersCharuco(
-            rgb_image, charuco_corners, charuco_ids, cv::Scalar(255, 0, 0));
-
-        cv::Vec3d rvec, tvec;
-        bool valid = cv::aruco::estimatePoseCharucoBoard(
-            charuco_corners_with_calibration, charuco_ids_with_calibration,
-            board, camera_matrix, dist_coeffs, rvec, tvec);
-
-        // if charuco pose is valid
-        if (valid) {
-          LOG(INFO) << std::fixed << std::setprecision(6)
-                    << "Age: " << age_double << ", Pose is R:" << rvec
-                    << " T:" << tvec;
-          cv::aruco::drawAxis(rgb_image, camera_matrix, dist_coeffs, rvec, tvec,
-                              0.1);
-        } else {
-          LOG(INFO) << "Age: " << age_double << ", invalid pose";
-        }
-      } else {
-        LOG(INFO) << "Age: " << age_double << ", no charuco IDs.";
-      }
-    } else {
-      LOG(INFO) << "Age: " << age_double << ", no marker IDs";
-    }
-
+  void HandleCharuco(cv::Mat rgb_image, const double /*age_double*/,
+                     std::vector<int> charuco_ids,
+                     std::vector<cv::Point2f> charuco_corners, bool valid,
+                     Eigen::Vector3d /*rvec_eigen*/,
+                     Eigen::Vector3d /*tvec_eigen*/) {
     cv::imshow("Display", rgb_image);
 
     if (FLAGS_display_undistorted) {
+      const cv::Size image_size(rgb_image.cols, rgb_image.rows);
       cv::Mat undistorted_rgb_image(image_size, CV_8UC3);
-      cv::undistort(rgb_image, undistorted_rgb_image, camera_matrix,
-                    dist_coeffs);
+      cv::undistort(rgb_image, undistorted_rgb_image,
+                    charuco_extractor_.camera_matrix(),
+                    charuco_extractor_.dist_coeffs());
 
       cv::imshow("Display undist", undistorted_rgb_image);
     }
 
     int keystroke = cv::waitKey(1);
     if ((keystroke & 0xFF) == static_cast<int>('c')) {
-      if (charuco_ids.size() >= FLAGS_min_targets) {
-        all_charuco_ids.emplace_back(std::move(charuco_ids));
-        all_charuco_corners.emplace_back(std::move(charuco_corners));
-        LOG(INFO) << "Image " << all_charuco_corners.size();
+      if (valid) {
+        all_charuco_ids_.emplace_back(std::move(charuco_ids));
+        all_charuco_corners_.emplace_back(std::move(charuco_corners));
+        LOG(INFO) << "Image " << all_charuco_corners_.size();
       }
 
-      if (all_charuco_ids.size() >= 50) {
+      if (all_charuco_ids_.size() >= 50) {
         LOG(INFO) << "Got enough images to calibrate";
-        event_loop.Exit();
+        event_loop_->Exit();
       }
     } else if ((keystroke & 0xFF) == static_cast<int>('q')) {
-      event_loop.Exit();
+      event_loop_->Exit();
     }
-  });
+  }
+
+  void MaybeCalibrate() {
+    if (all_charuco_ids_.size() >= 50) {
+      cv::Mat cameraMatrix, distCoeffs;
+      std::vector<cv::Mat> rvecs, tvecs;
+      cv::Mat stdDeviationsIntrinsics, stdDeviationsExtrinsics, perViewErrors;
+
+      // Set calibration flags (same as in calibrateCamera() function)
+      int calibration_flags = 0;
+      cv::Size img_size(640, 480);
+      const double reprojection_error = cv::aruco::calibrateCameraCharuco(
+          all_charuco_corners_, all_charuco_ids_, charuco_extractor_.board(),
+          img_size, cameraMatrix, distCoeffs, rvecs, tvecs,
+          stdDeviationsIntrinsics, stdDeviationsExtrinsics, perViewErrors,
+          calibration_flags);
+      CHECK_LE(reprojection_error, 1.0) << ": Reproduction error is bad.";
+      LOG(INFO) << "Reprojection Error is " << reprojection_error;
+
+      flatbuffers::FlatBufferBuilder fbb;
+      flatbuffers::Offset<flatbuffers::String> name_offset =
+          fbb.CreateString(absl::StrFormat("pi%d", pi_number_.value()));
+      flatbuffers::Offset<flatbuffers::Vector<float>> intrinsics_offset =
+          fbb.CreateVector<float>(9u, [&cameraMatrix](size_t i) {
+            return static_cast<float>(
+                reinterpret_cast<double *>(cameraMatrix.data)[i]);
+          });
+
+      flatbuffers::Offset<flatbuffers::Vector<float>>
+          distortion_coefficients_offset =
+              fbb.CreateVector<float>(5u, [&distCoeffs](size_t i) {
+                return static_cast<float>(
+                    reinterpret_cast<double *>(distCoeffs.data)[i]);
+              });
+
+      sift::CameraCalibration::Builder camera_calibration_builder(fbb);
+      std::optional<uint16_t> team_number =
+          aos::network::team_number_internal::ParsePiTeamNumber(pi_);
+      CHECK(team_number) << ": Invalid pi hostname " << pi_
+                         << ", failed to parse team number";
+
+      const aos::realtime_clock::time_point realtime_now =
+          aos::realtime_clock::now();
+      camera_calibration_builder.add_node_name(name_offset);
+      camera_calibration_builder.add_team_number(team_number.value());
+      camera_calibration_builder.add_calibration_timestamp(
+          realtime_now.time_since_epoch().count());
+      camera_calibration_builder.add_intrinsics(intrinsics_offset);
+      camera_calibration_builder.add_dist_coeffs(
+          distortion_coefficients_offset);
+      fbb.Finish(camera_calibration_builder.Finish());
+
+      aos::FlatbufferDetachedBuffer<sift::CameraCalibration> camera_calibration(
+          fbb.Release());
+      std::stringstream time_ss;
+      time_ss << realtime_now;
+
+      const std::string calibration_filename =
+          FLAGS_calibration_folder +
+          absl::StrFormat("/calibration_pi-%d-%d_%s.json", team_number.value(),
+                          pi_number_.value(), time_ss.str());
+
+      LOG(INFO) << calibration_filename << " -> "
+                << aos::FlatbufferToJson(camera_calibration,
+                                         {.multi_line = true});
+
+      aos::util::WriteStringToFileOrDie(
+          calibration_filename,
+          aos::FlatbufferToJson(camera_calibration, {.multi_line = true}));
+    } else {
+      LOG(INFO) << "Skipping calibration due to not enough images.";
+    }
+  }
+
+ private:
+  aos::ShmEventLoop *event_loop_;
+  std::string pi_;
+  const std::optional<uint16_t> pi_number_;
+
+  std::vector<std::vector<int>> all_charuco_ids_;
+  std::vector<std::vector<cv::Point2f>> all_charuco_corners_;
+
+  CharucoExtractor charuco_extractor_;
+};
+
+namespace {
+
+void Main() {
+  aos::FlatbufferDetachedBuffer<aos::Configuration> config =
+      aos::configuration::ReadConfig(FLAGS_config);
+
+  aos::ShmEventLoop event_loop(&config.message());
+
+  Calibration extractor(&event_loop, FLAGS_pi);
 
   event_loop.Run();
 
-  if (all_charuco_ids.size() >= 50) {
-    cv::Mat cameraMatrix, distCoeffs;
-    std::vector<cv::Mat> rvecs, tvecs;
-    cv::Mat stdDeviationsIntrinsics, stdDeviationsExtrinsics, perViewErrors;
-
-    // Set calibration flags (same than in calibrateCamera() function)
-    int calibration_flags = 0;
-    cv::Size img_size(640, 480);
-    const double reprojection_error = cv::aruco::calibrateCameraCharuco(
-        all_charuco_corners, all_charuco_ids, board, img_size, cameraMatrix,
-        distCoeffs, rvecs, tvecs, stdDeviationsIntrinsics,
-        stdDeviationsExtrinsics, perViewErrors, calibration_flags);
-    CHECK_LE(reprojection_error, 1.0) << ": Reproduction error is bad.";
-    LOG(INFO) << "Reprojection Error is " << reprojection_error;
-
-    flatbuffers::FlatBufferBuilder fbb;
-    flatbuffers::Offset<flatbuffers::String> name_offset =
-        fbb.CreateString(absl::StrFormat("pi%d", pi_number.value()));
-    flatbuffers::Offset<flatbuffers::Vector<float>> intrinsics_offset =
-        fbb.CreateVector<float>(9u, [&cameraMatrix](size_t i) {
-          return static_cast<float>(
-              reinterpret_cast<double *>(cameraMatrix.data)[i]);
-        });
-
-    flatbuffers::Offset<flatbuffers::Vector<float>>
-        distortion_coefficients_offset =
-            fbb.CreateVector<float>(5u, [&distCoeffs](size_t i) {
-              return static_cast<float>(
-                  reinterpret_cast<double *>(distCoeffs.data)[i]);
-            });
-
-    sift::CameraCalibration::Builder camera_calibration_builder(fbb);
-    std::optional<uint16_t> team_number =
-        aos::network::team_number_internal::ParsePiTeamNumber(FLAGS_pi);
-    CHECK(team_number) << ": Invalid pi hostname " << FLAGS_pi
-                       << ", failed to parse team number";
-
-    const aos::realtime_clock::time_point realtime_now =
-        aos::realtime_clock::now();
-    camera_calibration_builder.add_node_name(name_offset);
-    camera_calibration_builder.add_team_number(team_number.value());
-    camera_calibration_builder.add_calibration_timestamp(
-        realtime_now.time_since_epoch().count());
-    camera_calibration_builder.add_intrinsics(intrinsics_offset);
-    camera_calibration_builder.add_dist_coeffs(distortion_coefficients_offset);
-    camera_calibration_builder.add_dist_coeffs(distortion_coefficients_offset);
-    fbb.Finish(camera_calibration_builder.Finish());
-
-    aos::FlatbufferDetachedBuffer<sift::CameraCalibration> camera_calibration(
-        fbb.Release());
-    std::stringstream time_ss;
-    time_ss << realtime_now;
-
-    const std::string calibration_filename =
-        FLAGS_calibration_folder +
-        absl::StrFormat("/calibration_pi-%d-%d_%s.json", team_number.value(),
-                        pi_number.value(), time_ss.str());
-
-    LOG(INFO) << calibration_filename << " -> "
-              << aos::FlatbufferToJson(camera_calibration,
-                                       {.multi_line = true});
-
-    aos::util::WriteStringToFileOrDie(
-        calibration_filename,
-        aos::FlatbufferToJson(camera_calibration, {.multi_line = true}));
-  } else {
-    LOG(INFO) << "Skipping calibration due to not enough images.";
-  }
+  extractor.MaybeCalibrate();
 }
 
 }  // namespace
 }  // namespace vision
 }  // namespace frc971
 
-// Quick and lightweight grayscale viewer for images
 int main(int argc, char **argv) {
   aos::InitGoogle(&argc, &argv);
-  frc971::vision::ViewerMain();
+  frc971::vision::Main();
 }
diff --git a/y2020/vision/charuco_lib.cc b/y2020/vision/charuco_lib.cc
new file mode 100644
index 0000000..c86663a
--- /dev/null
+++ b/y2020/vision/charuco_lib.cc
@@ -0,0 +1,248 @@
+#include "y2020/vision/charuco_lib.h"
+
+#include <chrono>
+#include <functional>
+#include <string_view>
+
+#include <opencv2/core/eigen.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgproc.hpp>
+#include "aos/events/event_loop.h"
+#include "aos/flatbuffers.h"
+#include "aos/network/team_number.h"
+#include "frc971/control_loops/quaternion_utils.h"
+#include "glog/logging.h"
+#include "y2020/vision/sift/sift_generated.h"
+#include "y2020/vision/sift/sift_training_generated.h"
+#include "y2020/vision/tools/python_code/sift_training_data.h"
+#include "y2020/vision/vision_generated.h"
+
+DEFINE_uint32(min_targets, 10,
+              "The mininum number of targets required to match.");
+DEFINE_bool(large_board, true, "If true, use the large calibration board.");
+DEFINE_bool(coarse_pattern, true, "If true, use coarse arucos; else, use fine");
+DEFINE_string(board_template_path, "",
+              "If specified, write an image to the specified path for the "
+              "charuco board pattern.");
+
+namespace frc971 {
+namespace vision {
+namespace chrono = std::chrono;
+
+CameraCalibration::CameraCalibration(
+    const absl::Span<const uint8_t> training_data_bfbs, std::string_view pi) {
+  const aos::FlatbufferSpan<sift::TrainingData> training_data(
+      training_data_bfbs);
+  CHECK(training_data.Verify());
+  camera_calibration_ = FindCameraCalibration(&training_data.message(), pi);
+}
+
+cv::Mat CameraCalibration::CameraIntrinsics() const {
+  const cv::Mat result(3, 3, CV_32F,
+                       const_cast<void *>(static_cast<const void *>(
+                           camera_calibration_->intrinsics()->data())));
+  CHECK_EQ(result.total(), camera_calibration_->intrinsics()->size());
+  return result;
+}
+
+Eigen::Matrix3d CameraCalibration::CameraIntrinsicsEigen() const {
+  cv::Mat camera_intrinsics = CameraIntrinsics();
+  Eigen::Matrix3d result;
+  cv::cv2eigen(camera_intrinsics, result);
+  return result;
+}
+
+cv::Mat CameraCalibration::CameraDistCoeffs() const {
+  const cv::Mat result(5, 1, CV_32F,
+                       const_cast<void *>(static_cast<const void *>(
+                           camera_calibration_->dist_coeffs()->data())));
+  CHECK_EQ(result.total(), camera_calibration_->dist_coeffs()->size());
+  return result;
+}
+
+const sift::CameraCalibration *CameraCalibration::FindCameraCalibration(
+    const sift::TrainingData *const training_data, std::string_view pi) const {
+  std::optional<uint16_t> pi_number = aos::network::ParsePiNumber(pi);
+  std::optional<uint16_t> team_number =
+      aos::network::team_number_internal::ParsePiTeamNumber(pi);
+  CHECK(pi_number);
+  CHECK(team_number);
+  const std::string node_name = absl::StrFormat("pi%d", pi_number.value());
+  LOG(INFO) << "Looking for node name " << node_name << " team number "
+            << team_number.value();
+  for (const sift::CameraCalibration *candidate :
+       *training_data->camera_calibrations()) {
+    if (candidate->node_name()->string_view() != node_name) {
+      continue;
+    }
+    if (candidate->team_number() != team_number.value()) {
+      continue;
+    }
+    return candidate;
+  }
+  LOG(FATAL) << ": Failed to find camera calibration for " << node_name
+             << " on " << team_number.value();
+}
+
+ImageCallback::ImageCallback(aos::EventLoop *event_loop,
+                             std::string_view channel,
+                             std::function<void(cv::Mat, double)> &&fn)
+    : event_loop_(event_loop), handle_image_(std::move(fn)) {
+  event_loop_->MakeWatcher(channel, [this](const CameraImage &image) {
+    const aos::monotonic_clock::duration age =
+        event_loop_->monotonic_now() -
+        event_loop_->context().monotonic_event_time;
+    const double age_double =
+        std::chrono::duration_cast<std::chrono::duration<double>>(age).count();
+    if (age > std::chrono::milliseconds(100)) {
+      LOG(INFO) << "Age: " << age_double << ", getting behind, skipping";
+      return;
+    }
+    // Create color image:
+    cv::Mat image_color_mat(cv::Size(image.cols(), image.rows()), CV_8UC2,
+                            (void *)image.data()->data());
+    const cv::Size image_size(image.cols(), image.rows());
+    cv::Mat rgb_image(image_size, CV_8UC3);
+    cv::cvtColor(image_color_mat, rgb_image, CV_YUV2BGR_YUYV);
+    handle_image_(rgb_image, age_double);
+  });
+}
+
+CharucoExtractor::CharucoExtractor(
+    aos::EventLoop *event_loop, std::string_view pi,
+    std::function<void(cv::Mat, const double, std::vector<int>,
+                       std::vector<cv::Point2f>, bool, Eigen::Vector3d,
+                       Eigen::Vector3d)> &&fn)
+    : calibration_(SiftTrainingData(), pi),
+      dictionary_(cv::aruco::getPredefinedDictionary(
+          FLAGS_large_board ? cv::aruco::DICT_5X5_250
+                            : cv::aruco::DICT_6X6_250)),
+      board_(
+          FLAGS_large_board
+              ? (FLAGS_coarse_pattern ? cv::aruco::CharucoBoard::create(
+                                            12, 9, 0.06, 0.04666, dictionary_)
+                                      : cv::aruco::CharucoBoard::create(
+                                            25, 18, 0.03, 0.0233, dictionary_))
+              : (FLAGS_coarse_pattern ? cv::aruco::CharucoBoard::create(
+                                            7, 5, 0.04, 0.025, dictionary_)
+                                      // TODO(jim): Need to figure out what size
+                                      // is for small board, fine pattern
+                                      : cv::aruco::CharucoBoard::create(
+                                            7, 5, 0.03, 0.0233, dictionary_))),
+      camera_matrix_(calibration_.CameraIntrinsics()),
+      eigen_camera_matrix_(calibration_.CameraIntrinsicsEigen()),
+      dist_coeffs_(calibration_.CameraDistCoeffs()),
+      pi_number_(aos::network::ParsePiNumber(pi)),
+      image_callback_(
+          event_loop,
+          absl::StrCat("/pi", std::to_string(pi_number_.value()), "/camera"),
+          [this](cv::Mat rgb_image, const double age_double) {
+            HandleImage(rgb_image, age_double);
+          }),
+      handle_charuco_(std::move(fn)) {
+  LOG(INFO) << "Using " << (FLAGS_large_board ? "large" : "small")
+            << " board with " << (FLAGS_coarse_pattern ? "coarse" : "fine")
+            << " pattern";
+  if (!FLAGS_board_template_path.empty()) {
+    cv::Mat board_image;
+    board_->draw(cv::Size(600, 500), board_image, 10, 1);
+    cv::imwrite(FLAGS_board_template_path, board_image);
+  }
+
+  LOG(INFO) << "Camera matrix " << camera_matrix_;
+  LOG(INFO) << "Distortion Coefficients " << dist_coeffs_;
+
+  CHECK(pi_number_) << ": Invalid pi number " << pi
+                    << ", failed to parse pi number";
+
+  LOG(INFO) << "Connecting to channel /pi" << pi_number_.value() << "/camera";
+}
+
+void CharucoExtractor::HandleImage(cv::Mat rgb_image, const double age_double) {
+  std::vector<int> marker_ids;
+  std::vector<std::vector<cv::Point2f>> marker_corners;
+
+  cv::aruco::detectMarkers(rgb_image, board_->dictionary, marker_corners,
+                           marker_ids);
+
+  std::vector<cv::Point2f> charuco_corners;
+  std::vector<int> charuco_ids;
+  // If at least one marker detected
+  if (marker_ids.size() >= FLAGS_min_targets) {
+    // Run everything twice, once with the calibration, and once
+    // without. This lets us both calibrate, and also print out the pose
+    // real time with the previous calibration.
+    cv::aruco::interpolateCornersCharuco(marker_corners, marker_ids, rgb_image,
+                                         board_, charuco_corners, charuco_ids);
+
+    std::vector<cv::Point2f> charuco_corners_with_calibration;
+    std::vector<int> charuco_ids_with_calibration;
+
+    cv::aruco::interpolateCornersCharuco(
+        marker_corners, marker_ids, rgb_image, board_,
+        charuco_corners_with_calibration, charuco_ids_with_calibration,
+        camera_matrix_, dist_coeffs_);
+
+    cv::aruco::drawDetectedMarkers(rgb_image, marker_corners, marker_ids);
+
+    if (charuco_ids.size() >= FLAGS_min_targets) {
+      cv::aruco::drawDetectedCornersCharuco(rgb_image, charuco_corners,
+                                            charuco_ids, cv::Scalar(255, 0, 0));
+
+      cv::Vec3d rvec, tvec;
+      const bool valid = cv::aruco::estimatePoseCharucoBoard(
+          charuco_corners_with_calibration, charuco_ids_with_calibration,
+          board_, camera_matrix_, dist_coeffs_, rvec, tvec);
+
+      // if charuco pose is valid
+      if (valid) {
+        Eigen::Vector3d rvec_eigen;
+        Eigen::Vector3d tvec_eigen;
+        cv::cv2eigen(rvec, rvec_eigen);
+        cv::cv2eigen(tvec, tvec_eigen);
+
+        Eigen::Quaternion<double> rotation(
+            frc971::controls::ToQuaternionFromRotationVector(rvec_eigen));
+        Eigen::Translation3d translation(tvec_eigen);
+
+        const Eigen::Affine3d board_to_camera = translation * rotation;
+
+        Eigen::Matrix<double, 3, 4> camera_projection =
+            Eigen::Matrix<double, 3, 4>::Identity();
+        Eigen::Vector3d result = eigen_camera_matrix_ * camera_projection *
+                                 board_to_camera * Eigen::Vector3d::Zero();
+
+        result /= result.z();
+        cv::circle(rgb_image, cv::Point(result.x(), result.y()), 4,
+                   cv::Scalar(255, 255, 255), 0, cv::LINE_8);
+
+        cv::aruco::drawAxis(rgb_image, camera_matrix_, dist_coeffs_, rvec, tvec,
+                            0.1);
+
+        handle_charuco_(rgb_image, age_double, charuco_ids, charuco_corners,
+
+                        true, rvec_eigen, tvec_eigen);
+      } else {
+        LOG(INFO) << "Age: " << age_double << ", invalid pose";
+        handle_charuco_(rgb_image, age_double, charuco_ids, charuco_corners,
+
+                        false, Eigen::Vector3d::Zero(),
+                        Eigen::Vector3d::Zero());
+      }
+    } else {
+      LOG(INFO) << "Age: " << age_double << ", no charuco IDs.";
+      handle_charuco_(rgb_image, age_double, charuco_ids, charuco_corners,
+
+                      false, Eigen::Vector3d::Zero(), Eigen::Vector3d::Zero());
+    }
+  } else {
+    LOG(INFO) << "Age: " << age_double << ", no marker IDs";
+    cv::aruco::drawDetectedMarkers(rgb_image, marker_corners, marker_ids);
+
+    handle_charuco_(rgb_image, age_double, charuco_ids, charuco_corners, false,
+                    Eigen::Vector3d::Zero(), Eigen::Vector3d::Zero());
+  }
+}
+
+}  // namespace vision
+}  // namespace frc971
diff --git a/y2020/vision/charuco_lib.h b/y2020/vision/charuco_lib.h
new file mode 100644
index 0000000..e8dc3eb
--- /dev/null
+++ b/y2020/vision/charuco_lib.h
@@ -0,0 +1,112 @@
+#ifndef Y2020_VISION_CHARUCO_LIB_H_
+#define Y2020_VISION_CHARUCO_LIB_H_
+
+#include <functional>
+#include <string_view>
+
+#include <opencv2/aruco/charuco.hpp>
+#include <opencv2/calib3d.hpp>
+#include "Eigen/Dense"
+#include "Eigen/Geometry"
+
+#include "absl/types/span.h"
+#include "aos/events/event_loop.h"
+#include "y2020/vision/sift/sift_generated.h"
+#include "y2020/vision/sift/sift_training_generated.h"
+
+namespace frc971 {
+namespace vision {
+
+// Class to find extrinsics for a specified pi's camera using the provided
+// training data.
+class CameraCalibration {
+ public:
+  CameraCalibration(const absl::Span<const uint8_t> training_data_bfbs,
+                    std::string_view pi);
+
+  // Intrinsics for the located camera.
+  cv::Mat CameraIntrinsics() const;
+  Eigen::Matrix3d CameraIntrinsicsEigen() const;
+
+  // Distortion coefficients for the located camera.
+  cv::Mat CameraDistCoeffs() const;
+
+ private:
+  // Finds the camera specific calibration flatbuffer.
+  const sift::CameraCalibration *FindCameraCalibration(
+      const sift::TrainingData *const training_data, std::string_view pi) const;
+
+  // Pointer to this camera's calibration parameters.
+  const sift::CameraCalibration *camera_calibration_;
+};
+
+// Class to call a function with a cv::Mat and age when an image shows up on the
+// provided channel.  This hides all the conversions and wrangling needed to
+// view the image.
+class ImageCallback {
+ public:
+  ImageCallback(aos::EventLoop *event_loop, std::string_view channel,
+                std::function<void(cv::Mat, double)> &&fn);
+
+ private:
+  aos::EventLoop *event_loop_;
+  std::function<void(cv::Mat, double)> handle_image_;
+};
+
+// Class which calls a callback each time an image arrives with the information
+// extracted from it.
+class CharucoExtractor {
+ public:
+  // The callback takes the following arguments:
+  //   cv::Mat -> image with overlays drawn on it.
+  //   const double -> Duration between when the image was captured and this
+  //                   callback was called.
+  //   std::vector<int> -> charuco_ids
+  //   std::vector<cv::Point2f> -> charuco_corners
+  //   bool -> true if rvec/tvec is valid.
+  //   Eigen::Vector3d -> rvec
+  //   Eigen::Vector3d -> tvec
+  CharucoExtractor(aos::EventLoop *event_loop, std::string_view pi,
+                   std::function<void(cv::Mat, const double, std::vector<int>,
+                                      std::vector<cv::Point2f>, bool,
+                                      Eigen::Vector3d, Eigen::Vector3d)> &&fn);
+
+  // Returns the aruco dictionary in use.
+  cv::Ptr<cv::aruco::Dictionary> dictionary() const { return dictionary_; }
+  // Returns the aruco board in use.
+  cv::Ptr<cv::aruco::CharucoBoard> board() const { return board_; }
+
+  // Returns the camera matrix for this camera.
+  const cv::Mat camera_matrix() const { return camera_matrix_; }
+  // Returns the distortion coefficients for this camera.
+  const cv::Mat dist_coeffs() const { return dist_coeffs_; }
+
+ private:
+  // Handles the image by detecting the charuco board in it and calling the
+  // callback with the extracted board corners, corresponding IDs, and pose.
+  void HandleImage(cv::Mat rgb_image, const double age_double);
+
+  CameraCalibration calibration_;
+
+  cv::Ptr<cv::aruco::Dictionary> dictionary_;
+  cv::Ptr<cv::aruco::CharucoBoard> board_;
+
+  const cv::Mat camera_matrix_;
+  const Eigen::Matrix3d eigen_camera_matrix_;
+  const cv::Mat dist_coeffs_;
+
+  const std::optional<uint16_t> pi_number_;
+
+  ImageCallback image_callback_;
+
+  // Function to call.
+  std::function<void(cv::Mat, const double, std::vector<int>,
+                     std::vector<cv::Point2f>, bool, Eigen::Vector3d,
+                     Eigen::Vector3d)>
+      handle_charuco_;
+};
+
+}  // namespace vision
+}  // namespace frc971
+
+#endif  // Y2020_VISION_CHARUCO_LIB_H_