Setting up Blob detector and viewer code in y2022

Includes viewer with SIFT code stripped out, since we're not using it
in y2022.

Change-Id: Id56e9e5c868a1d730dcfb433d5e4868898a6f011
Signed-off-by: Jim Ostrowski <yimmy13@gmail.com>
diff --git a/y2020/vision/BUILD b/y2020/vision/BUILD
index d94de5a..330848a 100644
--- a/y2020/vision/BUILD
+++ b/y2020/vision/BUILD
@@ -5,7 +5,7 @@
     srcs = ["vision.fbs"],
     gen_reflections = 1,
     target_compatible_with = ["@platforms//os:linux"],
-    visibility = ["//y2020:__subpackages__"],
+    visibility = ["//y2020:__subpackages__"] + ["//y2022:__subpackages__"],
 )
 
 flatbuffer_cc_library(
@@ -61,7 +61,7 @@
         "//y2020:config",
     ],
     target_compatible_with = ["@platforms//os:linux"],
-    visibility = ["//y2020:__subpackages__"],
+    visibility = ["//y2020:__subpackages__"] + ["//y2022:__subpackages__"],
     deps = [
         ":v4l2_reader",
         ":vision_fbs",
diff --git a/y2022/vision/BUILD b/y2022/vision/BUILD
index e69de29..c10bf85 100644
--- a/y2022/vision/BUILD
+++ b/y2022/vision/BUILD
@@ -0,0 +1,49 @@
+cc_binary(
+    name = "camera_reader",
+    srcs = [
+        "camera_reader_main.cc",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//y2022:__subpackages__"],
+    deps = [
+        "//aos:init",
+        "//aos/events:shm_event_loop",
+        "//y2020/vision:camera_reader_lib",
+    ],
+)
+
+cc_library(
+    name = "blob_detector_lib",
+    srcs = [
+        "blob_detector.cc",
+    ],
+    hdrs = [
+        "blob_detector.h",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//y2022:__subpackages__"],
+    deps = [
+        "//aos/network:team_number",
+        "//third_party:opencv",
+    ],
+)
+
+cc_binary(
+    name = "viewer",
+    srcs = [
+        "viewer.cc",
+    ],
+    data = [
+        "//y2020:config",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//y2022:__subpackages__"],
+    deps = [
+        ":blob_detector_lib",
+        "//aos:init",
+        "//aos/events:shm_event_loop",
+        "//third_party:opencv",
+        "//y2020/vision:vision_fbs",
+        "//y2020/vision/sift:sift_fbs",
+    ],
+)
diff --git a/y2022/vision/blob_detector.cc b/y2022/vision/blob_detector.cc
new file mode 100644
index 0000000..bf5afd7
--- /dev/null
+++ b/y2022/vision/blob_detector.cc
@@ -0,0 +1,98 @@
+#include "y2022/vision/blob_detector.h"
+
+#include <opencv2/imgproc.hpp>
+
+#include "aos/network/team_number.h"
+
+DEFINE_uint64(green_delta, 50,
+              "Required difference between green pixels vs. red and blue");
+DEFINE_bool(use_outdoors, false,
+            "If true, change thresholds to handle outdoor illumination");
+
+namespace y2022 {
+namespace vision {
+
+cv::Mat BlobDetector::ThresholdImage(cv::Mat rgb_image) {
+  cv::Mat gray_image(cv::Size(rgb_image.cols, rgb_image.rows), CV_8UC1);
+  for (int row = 0; row < rgb_image.rows; row++) {
+    for (int col = 0; col < rgb_image.cols; col++) {
+      cv::Vec3b pixel = rgb_image.at<cv::Vec3b>(row, col);
+      uint8_t blue = pixel.val[0];
+      uint8_t green = pixel.val[1];
+      uint8_t red = pixel.val[2];
+      // Simple filter that looks for green pixels sufficiently brigher than
+      // red and blue
+      if ((green > blue + 30) && (green > red + 50)) {
+        gray_image.at<uint8_t>(row, col) = 255;
+      } else {
+        gray_image.at<uint8_t>(row, col) = 0;
+      }
+    }
+  }
+
+  return gray_image;
+}
+
+std::vector<std::vector<cv::Point>> BlobDetector::FindBlobs(
+    cv::Mat binarized_image) {
+  // find the contours (blob outlines)
+  std::vector<std::vector<cv::Point>> contours;
+  std::vector<cv::Vec4i> hierarchy;
+  cv::findContours(binarized_image, contours, hierarchy, cv::RETR_CCOMP,
+                   cv::CHAIN_APPROX_SIMPLE);
+
+  return contours;
+}
+
+// Filter blobs to get rid of noise, too large items, etc.
+std::vector<std::vector<cv::Point>> BlobDetector::FilterBlobs(
+    std::vector<std::vector<cv::Point>> blobs) {
+  // TODO: Put in some filters
+
+  std::vector<std::vector<cv::Point>> filtered_blobs;
+  for (auto blob : blobs) {
+    // for now, let's remove all blobs that are at the bottom of the image
+    if (blob[0].y < 400) {
+      filtered_blobs.push_back(blob);
+    } else {
+      // LOG(INFO) << "Found and removed blob";
+    }
+  }
+  return filtered_blobs;
+}
+
+void BlobDetector::DrawBlobs(
+    cv::Mat view_image, std::vector<std::vector<cv::Point>> unfiltered_blobs,
+    std::vector<std::vector<cv::Point>> filtered_blobs) {
+  CHECK_GT(view_image.cols, 0);
+  if (unfiltered_blobs.size() > 0) {
+    // Draw blobs unfilled, with red color border
+    drawContours(view_image, unfiltered_blobs, -1, cv::Scalar(0, 0, 255), 0);
+  }
+
+  drawContours(view_image, filtered_blobs, -1, cv::Scalar(0, 255, 0),
+               cv::FILLED);
+}
+
+std::vector<std::vector<cv::Point>> BlobDetector::ComputeStats(
+    std::vector<std::vector<cv::Point>> blobs) {
+  // Placeholder for now for this
+  // TODO<Jim>: need to compute stats on blobs, like centroid, aspect
+  // ratio, bounding box
+  return blobs;
+}
+
+void BlobDetector::ExtractBlobs(
+    cv::Mat rgb_image, cv::Mat binarized_image, cv::Mat blob_image,
+    std::vector<std::vector<cv::Point>> &filtered_blobs,
+    std::vector<std::vector<cv::Point>> &unfiltered_blobs,
+    std::vector<std::vector<cv::Point>> &blob_stats) {
+  binarized_image = ThresholdImage(rgb_image);
+  unfiltered_blobs = FindBlobs(binarized_image);
+  filtered_blobs = FilterBlobs(unfiltered_blobs);
+  DrawBlobs(blob_image, unfiltered_blobs, filtered_blobs);
+  blob_stats = ComputeStats(filtered_blobs);
+}
+
+}  // namespace vision
+}  // namespace y2022
diff --git a/y2022/vision/blob_detector.h b/y2022/vision/blob_detector.h
new file mode 100644
index 0000000..3da3d4b
--- /dev/null
+++ b/y2022/vision/blob_detector.h
@@ -0,0 +1,43 @@
+#ifndef Y2022_BLOB_DETECTOR_H_
+#define Y2022_BLOB_DETECTOR_H_
+
+#include <opencv2/imgproc.hpp>
+
+namespace y2022 {
+namespace vision {
+
+class BlobDetector {
+ public:
+  BlobDetector() {}
+  // Given an image, threshold it to find "green" pixels
+  // Input: Color image
+  // Output: Grayscale (binarized) image with green pixels set to 255
+  static cv::Mat ThresholdImage(cv::Mat rgb_image);
+
+  // Given binary image, extract blobs
+  static std::vector<std::vector<cv::Point>> FindBlobs(cv::Mat threshold_image);
+
+  // Filter blobs to get rid of noise, too large items, etc.
+  static std::vector<std::vector<cv::Point>> FilterBlobs(
+      std::vector<std::vector<cv::Point>> blobs);
+
+  // Draw Blobs on image
+  // Optionally draw all blobs and filtered blobs
+  static void DrawBlobs(cv::Mat view_image,
+                        std::vector<std::vector<cv::Point>> filtered_blobs,
+                        std::vector<std::vector<cv::Point>> unfiltered_blobs);
+
+  // Extract stats for each blob
+  static std::vector<std::vector<cv::Point>> ComputeStats(
+      std::vector<std::vector<cv::Point>>);
+
+  static void ExtractBlobs(
+      cv::Mat rgb_image, cv::Mat binarized_image, cv::Mat blob_image,
+      std::vector<std::vector<cv::Point>> &filtered_blobs,
+      std::vector<std::vector<cv::Point>> &unfiltered_blobs,
+      std::vector<std::vector<cv::Point>> &blob_stats);
+};
+}  // namespace vision
+}  // namespace y2022
+
+#endif  // Y2022_BLOB_DETECTOR_H_
diff --git a/y2022/vision/camera_reader_main.cc b/y2022/vision/camera_reader_main.cc
new file mode 100644
index 0000000..ee65cfe
--- /dev/null
+++ b/y2022/vision/camera_reader_main.cc
@@ -0,0 +1,58 @@
+#include "aos/events/shm_event_loop.h"
+#include "aos/init.h"
+#include "y2020/vision/camera_reader.h"
+
+// config used to allow running camera_reader independently.  E.g.,
+// bazel run //y2022/vision:camera_reader -- --config y2022/config.json
+//   --override_hostname pi-7971-1  --ignore_timestamps true
+DEFINE_string(config, "config.json", "Path to the config file to use.");
+DEFINE_uint32(exposure, 5, "Exposure time, in 100us increments");
+
+namespace y2022 {
+namespace vision {
+namespace {
+
+using namespace frc971::vision;
+
+void CameraReaderMain() {
+  aos::FlatbufferDetachedBuffer<aos::Configuration> config =
+      aos::configuration::ReadConfig(FLAGS_config);
+
+  const aos::FlatbufferSpan<sift::TrainingData> training_data(
+      SiftTrainingData());
+  CHECK(training_data.Verify());
+
+  const auto index_params = cv::makePtr<cv::flann::IndexParams>();
+  index_params->setAlgorithm(cvflann::FLANN_INDEX_KDTREE);
+  index_params->setInt("trees", 5);
+  const auto search_params =
+      cv::makePtr<cv::flann::SearchParams>(/* checks */ 50);
+  cv::FlannBasedMatcher matcher(index_params, search_params);
+
+  aos::ShmEventLoop event_loop(&config.message());
+
+  // First, log the data for future reference.
+  {
+    aos::Sender<sift::TrainingData> training_data_sender =
+        event_loop.MakeSender<sift::TrainingData>("/camera");
+    CHECK_EQ(training_data_sender.Send(training_data),
+             aos::RawSender::Error::kOk);
+  }
+
+  V4L2Reader v4l2_reader(&event_loop, "/dev/video0");
+  CameraReader camera_reader(&event_loop, &training_data.message(),
+                             &v4l2_reader, index_params, search_params);
+
+  v4l2_reader.SetExposure(FLAGS_exposure);
+
+  event_loop.Run();
+}
+
+}  // namespace
+}  // namespace vision
+}  // namespace y2022
+
+int main(int argc, char **argv) {
+  aos::InitGoogle(&argc, &argv);
+  y2022::vision::CameraReaderMain();
+}
diff --git a/y2022/vision/viewer.cc b/y2022/vision/viewer.cc
new file mode 100644
index 0000000..8b76d21
--- /dev/null
+++ b/y2022/vision/viewer.cc
@@ -0,0 +1,147 @@
+#include <map>
+#include <opencv2/calib3d.hpp>
+#include <opencv2/features2d.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgproc.hpp>
+#include <random>
+
+#include "aos/events/shm_event_loop.h"
+#include "aos/init.h"
+#include "aos/time/time.h"
+#include "y2020/vision/vision_generated.h"
+#include "y2022/vision/blob_detector.h"
+
+DEFINE_string(capture, "",
+              "If set, capture a single image and save it to this filename.");
+DEFINE_string(channel, "/camera", "Channel name for the image.");
+DEFINE_string(config, "config.json", "Path to the config file to use.");
+DEFINE_string(png_dir,
+              "/home/jim/code/FRC/971-Robot-Code/y2020/vision/LED_Ring_exp",
+              "Path to a set of images to display.");
+DEFINE_bool(show_features, true, "Show the blobs.");
+
+namespace y2022 {
+namespace vision {
+namespace {
+
+aos::Fetcher<frc971::vision::CameraImage> image_fetcher;
+
+bool DisplayLoop() {
+  int64_t image_timestamp = 0;
+  const frc971::vision::CameraImage *image;
+  // Read next image
+  if (!image_fetcher.FetchNext()) {
+    VLOG(2) << "Couldn't fetch next image";
+    return true;
+  }
+
+  image = image_fetcher.get();
+  CHECK(image != nullptr) << "Couldn't read image";
+  image_timestamp = image->monotonic_timestamp_ns();
+  VLOG(2) << "Got image at timestamp: " << image_timestamp;
+
+  // Create color image:
+  cv::Mat image_color_mat(cv::Size(image->cols(), image->rows()), CV_8UC2,
+                          (void *)image->data()->data());
+  cv::Mat rgb_image(cv::Size(image->cols(), image->rows()), CV_8UC3);
+  cv::cvtColor(image_color_mat, rgb_image, cv::COLOR_YUV2BGR_YUYV);
+
+  if (!FLAGS_capture.empty()) {
+    cv::imwrite(FLAGS_capture, rgb_image);
+    return false;
+  }
+
+  cv::Mat binarized_image, ret_image;
+  std::vector<std::vector<cv::Point>> unfiltered_blobs, filtered_blobs,
+      blob_stats;
+  BlobDetector::ExtractBlobs(rgb_image, binarized_image, ret_image,
+                             filtered_blobs, unfiltered_blobs, blob_stats);
+
+  LOG(INFO) << image->monotonic_timestamp_ns()
+            << ": # blobs: " << filtered_blobs.size();
+
+  // Downsize for viewing
+  cv::resize(rgb_image, rgb_image,
+             cv::Size(rgb_image.cols / 2, rgb_image.rows / 2),
+             cv::INTER_LINEAR);
+
+  cv::imshow("image", rgb_image);
+  cv::imshow("blobs", ret_image);
+
+  int keystroke = cv::waitKey(1);
+  if ((keystroke & 0xFF) == static_cast<int>('c')) {
+    // Convert again, to get clean image
+    cv::cvtColor(image_color_mat, rgb_image, cv::COLOR_YUV2BGR_YUYV);
+    std::stringstream name;
+    name << "capture-" << aos::realtime_clock::now() << ".png";
+    cv::imwrite(name.str(), rgb_image);
+    LOG(INFO) << "Saved image file: " << name.str();
+  } else if ((keystroke & 0xFF) == static_cast<int>('q')) {
+    return false;
+  }
+  return true;
+}
+
+void ViewerMain() {
+  aos::FlatbufferDetachedBuffer<aos::Configuration> config =
+      aos::configuration::ReadConfig(FLAGS_config);
+
+  aos::ShmEventLoop event_loop(&config.message());
+
+  image_fetcher =
+      event_loop.MakeFetcher<frc971::vision::CameraImage>(FLAGS_channel);
+
+  // Run the display loop
+  event_loop.AddPhasedLoop(
+      [&event_loop](int) {
+        if (!DisplayLoop()) {
+          LOG(INFO) << "Calling event_loop Exit";
+          event_loop.Exit();
+        };
+      },
+      ::std::chrono::milliseconds(100));
+
+  event_loop.Run();
+
+  image_fetcher = aos::Fetcher<frc971::vision::CameraImage>();
+}
+
+void ViewerLocal() {
+  std::vector<cv::String> file_list;
+  cv::glob(FLAGS_png_dir + "/*.png", file_list, false);
+  for (auto file : file_list) {
+    LOG(INFO) << "Reading file " << file;
+    cv::Mat rgb_image = cv::imread(file.c_str());
+    std::vector<std::vector<cv::Point>> filtered_blobs, unfiltered_blobs,
+        blob_stats;
+    cv::Mat binarized_image =
+        cv::Mat::zeros(cv::Size(rgb_image.cols, rgb_image.rows), CV_8UC1);
+    cv::Mat ret_image =
+        cv::Mat::zeros(cv::Size(rgb_image.cols, rgb_image.rows), CV_8UC3);
+
+    BlobDetector::ExtractBlobs(rgb_image, binarized_image, ret_image,
+                               filtered_blobs, unfiltered_blobs, blob_stats);
+
+    LOG(INFO) << ": # blobs: " << filtered_blobs.size() << " (# removed: "
+              << unfiltered_blobs.size() - filtered_blobs.size() << ")";
+    cv::imshow("image", rgb_image);
+    cv::imshow("blobs", ret_image);
+
+    int keystroke = cv::waitKey(1000);
+    if ((keystroke & 0xFF) == static_cast<int>('q')) {
+      return;
+    }
+  }
+}
+}  // namespace
+}  // namespace vision
+}  // namespace y2022
+
+// Quick and lightweight viewer for images
+int main(int argc, char **argv) {
+  aos::InitGoogle(&argc, &argv);
+  if (FLAGS_png_dir != "")
+    y2022::vision::ViewerLocal();
+  else
+    y2022::vision::ViewerMain();
+}