Add support for different image formats in ImageCallback

The april tag code is faster in halide than opencv.

Change-Id: I6531c5a45158e8bd5e50ee0b0ff69cc1978bf588
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/frc971/vision/charuco_lib.cc b/frc971/vision/charuco_lib.cc
index 01e18e8..1fb1795 100644
--- a/frc971/vision/charuco_lib.cc
+++ b/frc971/vision/charuco_lib.cc
@@ -168,9 +168,22 @@
     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::COLOR_YUV2BGR_YUYV);
-    handle_image_(rgb_image, eof);
+    switch (format_) {
+      case Format::GRAYSCALE: {
+        ftrace_.FormatMessage("Starting yuyv->greyscale\n");
+        cv::Mat gray_image(image_size, CV_8UC3);
+        cv::cvtColor(image_color_mat, gray_image, cv::COLOR_YUV2GRAY_YUYV);
+        handle_image_(gray_image, eof);
+      } break;
+      case Format::BGR: {
+        cv::Mat rgb_image(image_size, CV_8UC3);
+        cv::cvtColor(image_color_mat, rgb_image, cv::COLOR_YUV2BGR_YUYV);
+        handle_image_(rgb_image, eof);
+      } break;
+      case Format::YUYV2: {
+        handle_image_(image_color_mat, eof);
+      };
+    }
   });
 }
 
diff --git a/frc971/vision/charuco_lib.h b/frc971/vision/charuco_lib.h
index fdf25e1..f4f1233 100644
--- a/frc971/vision/charuco_lib.h
+++ b/frc971/vision/charuco_lib.h
@@ -51,10 +51,19 @@
 // full-service callback functionality
 class ImageCallback {
  public:
+  enum class Format {
+    YUYV2 = 0,
+    BGR = 1,
+    GRAYSCALE = 2,
+  };
   ImageCallback(aos::EventLoop *event_loop, std::string_view channel,
                 std::function<void(cv::Mat, aos::monotonic_clock::time_point)>
                     &&handle_image_fn);
 
+  void set_format(Format format) {
+    format_ = format;
+  }
+
  private:
   void DisableTracing();
 
@@ -67,6 +76,8 @@
   bool disabling_ = false;
 
   aos::Ftrace ftrace_;
+
+  Format format_ = Format::BGR;
 };
 
 // Class which calls a callback each time an image arrives with the information