Stick blob results into a flatbuffer

subscibed to this in viewer

Signed-off-by: Henry Speiser <henry@speiser.net>
Change-Id: I5ea9a3b8ab31ea7c8e9351469ba80de9c6cc3d2a
Signed-off-by: Henry Speiser <henry@speiser.net>
diff --git a/y2022/vision/BUILD b/y2022/vision/BUILD
index 823462b..676a4b9 100644
--- a/y2022/vision/BUILD
+++ b/y2022/vision/BUILD
@@ -96,6 +96,7 @@
         ":blob_detector_lib",
         ":calibration_data",
         ":calibration_fbs",
+        ":target_estimate_fbs",
         ":target_estimator_lib",
         "//aos:flatbuffer_merge",
         "//aos/events:event_loop",
diff --git a/y2022/vision/blob_detector.cc b/y2022/vision/blob_detector.cc
index 8ac96cd..aba3550 100644
--- a/y2022/vision/blob_detector.cc
+++ b/y2022/vision/blob_detector.cc
@@ -65,9 +65,10 @@
     const double aspect_ratio =
         static_cast<double>(blob_size.width) / blob_size.height;
     const double area = moments.m00;
-    const size_t points = blob.size();
+    const size_t num_points = blob.size();
 
-    blob_stats.emplace_back(BlobStats{centroid, aspect_ratio, area, points});
+    blob_stats.emplace_back(
+        BlobStats{centroid, aspect_ratio, area, num_points});
   }
   return blob_stats;
 }
@@ -185,7 +186,7 @@
     if ((stats_it->centroid.y <= kMaxY) &&
         (std::abs(kTapeAspectRatio - stats_it->aspect_ratio) <
          kAspectRatioThreshold) &&
-        (stats_it->area >= kMinArea) && (stats_it->points >= kMinPoints)) {
+        (stats_it->area >= kMinArea) && (stats_it->num_points >= kMinPoints)) {
       filtered_blobs.push_back(*blob_it);
       filtered_stats.push_back(*stats_it);
     }
@@ -293,7 +294,7 @@
 }
 
 void BlobDetector::ExtractBlobs(
-    cv::Mat rgb_image, cv::Mat &binarized_image, cv::Mat blob_image,
+    cv::Mat rgb_image, cv::Mat &binarized_image,
     std::vector<std::vector<cv::Point>> &filtered_blobs,
     std::vector<std::vector<cv::Point>> &unfiltered_blobs,
     std::vector<BlobStats> &blob_stats, cv::Point &centroid) {
@@ -308,7 +309,6 @@
   LOG(INFO) << "Blob detection elapsed time: "
             << std::chrono::duration<double, std::milli>(end - start).count()
             << " ms";
-  DrawBlobs(blob_image, unfiltered_blobs, filtered_blobs, blob_stats, centroid);
 }
 
 }  // namespace vision
diff --git a/y2022/vision/blob_detector.h b/y2022/vision/blob_detector.h
index a0663fd..f8a4ab4 100644
--- a/y2022/vision/blob_detector.h
+++ b/y2022/vision/blob_detector.h
@@ -13,7 +13,7 @@
     cv::Point centroid;
     double aspect_ratio;
     double area;
-    size_t points;
+    size_t num_points;
   };
 
   BlobDetector() {}
@@ -45,7 +45,7 @@
       const std::vector<BlobStats> &blob_stats, cv::Point centroid);
 
   static void ExtractBlobs(
-      cv::Mat rgb_image, cv::Mat &binarized_image, cv::Mat blob_image,
+      cv::Mat rgb_image, cv::Mat &binarized_image,
       std::vector<std::vector<cv::Point>> &filtered_blobs,
       std::vector<std::vector<cv::Point>> &unfiltered_blobs,
       std::vector<BlobStats> &blob_stats, cv::Point &centroid);
diff --git a/y2022/vision/camera_reader.cc b/y2022/vision/camera_reader.cc
index ca2535a..69514f9 100644
--- a/y2022/vision/camera_reader.cc
+++ b/y2022/vision/camera_reader.cc
@@ -39,26 +39,80 @@
              << " on " << team_number;
 }
 
+namespace {
+// Converts a vector of cv::Point to PointT for the flatbuffer
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Blob>>>
+CvBlobsToFbs(const std::vector<std::vector<cv::Point>> &blobs,
+             aos::Sender<TargetEstimate>::Builder &builder) {
+  std::vector<flatbuffers::Offset<Blob>> blobs_fbs;
+  for (auto &blob : blobs) {
+    std::vector<Point> points_fbs;
+    for (auto p : blob) {
+      points_fbs.push_back(Point{p.x, p.y});
+    }
+    auto points_offset = builder.fbb()->CreateVectorOfStructs(points_fbs);
+    auto blob_builder = builder.MakeBuilder<Blob>();
+    blob_builder.add_points(points_offset);
+    blobs_fbs.emplace_back(blob_builder.Finish());
+  }
+  return builder.fbb()->CreateVector(blobs_fbs.data(), blobs_fbs.size());
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<BlobStatsFbs>>>
+BlobStatsToFbs(const std::vector<BlobDetector::BlobStats> blob_stats,
+               aos::Sender<TargetEstimate>::Builder &builder) {
+  std::vector<flatbuffers::Offset<BlobStatsFbs>> stats_fbs;
+  for (auto &stats : blob_stats) {
+    // Make BlobStatsFbs builder then fill each field using the BlobStats
+    // struct, then you finish it and add it to stats_fbs.
+    auto stats_builder = builder.MakeBuilder<BlobStatsFbs>();
+    Point centroid_fbs = Point{stats.centroid.x, stats.centroid.y};
+    stats_builder.add_centroid(&centroid_fbs);
+    stats_builder.add_aspect_ratio(stats.aspect_ratio);
+    stats_builder.add_area(stats.area);
+    stats_builder.add_num_points(stats.num_points);
+
+    auto current_stats = stats_builder.Finish();
+    stats_fbs.emplace_back(current_stats);
+  }
+  return builder.fbb()->CreateVector(stats_fbs.data(), stats_fbs.size());
+}
+}  // namespace
+
 void CameraReader::ProcessImage(cv::Mat image_mat) {
   // Remember, we're getting YUYV images, so we start by converting to RGB
 
-  // TOOD: Need to code this up for blob detection
-
   std::vector<std::vector<cv::Point>> filtered_blobs, unfiltered_blobs;
   std::vector<BlobDetector::BlobStats> blob_stats;
   cv::Mat binarized_image =
       cv::Mat::zeros(cv::Size(image_mat.cols, image_mat.rows), CV_8UC1);
-  cv::Mat ret_image =
-      cv::Mat::zeros(cv::Size(image_mat.cols, image_mat.rows), CV_8UC3);
   cv::Point centroid;
-  BlobDetector::ExtractBlobs(image_mat, binarized_image, ret_image,
-                             filtered_blobs, unfiltered_blobs, blob_stats,
-                             centroid);
-  TargetEstimateT target = TargetEstimator::EstimateTargetLocation(
-      centroid, CameraIntrinsics(), CameraExtrinsics());
-
+  BlobDetector::ExtractBlobs(image_mat, binarized_image, filtered_blobs,
+                             unfiltered_blobs, blob_stats, centroid);
   auto builder = target_estimate_sender_.MakeBuilder();
-  builder.CheckOk(builder.Send(TargetEstimate::Pack(*builder.fbb(), &target)));
+  flatbuffers::Offset<BlobResult> blob_result_offset;
+  {
+    const auto filtered_blobs_offset = CvBlobsToFbs(filtered_blobs, builder);
+    const auto unfiltered_blobs_offset =
+        CvBlobsToFbs(unfiltered_blobs, builder);
+    const auto blob_stats_offset = BlobStatsToFbs(blob_stats, builder);
+    const Point centroid_fbs = Point{centroid.x, centroid.y};
+
+    auto blob_result_builder = builder.MakeBuilder<BlobResult>();
+    blob_result_builder.add_filtered_blobs(filtered_blobs_offset);
+    blob_result_builder.add_unfiltered_blobs(unfiltered_blobs_offset);
+    blob_result_builder.add_blob_stats(blob_stats_offset);
+    blob_result_builder.add_centroid(&centroid_fbs);
+    blob_result_offset = blob_result_builder.Finish();
+  }
+
+  auto target_estimate_builder = builder.MakeBuilder<TargetEstimate>();
+  TargetEstimator::EstimateTargetLocation(centroid, CameraIntrinsics(),
+                                          CameraExtrinsics(),
+                                          &target_estimate_builder);
+  target_estimate_builder.add_blob_result(blob_result_offset);
+
+  builder.CheckOk(builder.Send(target_estimate_builder.Finish()));
 }
 
 void CameraReader::ReadImage() {
diff --git a/y2022/vision/target_estimate.fbs b/y2022/vision/target_estimate.fbs
index e22654a..207a37a 100644
--- a/y2022/vision/target_estimate.fbs
+++ b/y2022/vision/target_estimate.fbs
@@ -1,5 +1,34 @@
 namespace y2022.vision;
 
+struct Point {
+  x:int (id: 0);
+  y:int (id: 1);
+}
+
+table Blob  {
+  points:[Point] (id: 0);
+}
+
+// Statistics for each blob used for filtering
+table BlobStatsFbs {
+  centroid:Point (id: 0);
+  aspect_ratio:double (id: 1);
+  area:double (id: 2);
+  num_points:uint64 (id: 3);
+}
+
+// Information for debugging blob detection
+table BlobResult {
+  // Blobs that passed the filtering step
+  filtered_blobs:[Blob] (id: 0);
+  // All detected blobs
+  unfiltered_blobs:[Blob] (id: 1);
+  // Stats on the blobs
+  blob_stats:[BlobStatsFbs] (id: 2);
+  // Average centroid of the filtered blobs
+  centroid:Point (id: 3);
+}
+
 // Contains the information the EKF wants from blobs from a single image.
 table TargetEstimate {
   // Horizontal distance from the camera to the center of the upper hub
@@ -8,7 +37,9 @@
   // Positive means right of center, negative means left.
   angle_to_target:double (id: 1);
 
-  // TODO(milind): add confidence and blob stats
+  blob_result:BlobResult (id: 2);
+
+  // TODO(milind): add confidence
 }
 
 root_type TargetEstimate;
diff --git a/y2022/vision/target_estimator.cc b/y2022/vision/target_estimator.cc
index c51eb08..1ac3d55 100644
--- a/y2022/vision/target_estimator.cc
+++ b/y2022/vision/target_estimator.cc
@@ -2,9 +2,10 @@
 
 namespace y2022::vision {
 
-TargetEstimateT TargetEstimator::EstimateTargetLocation(
-    cv::Point2i blob_point, const cv::Mat &intrinsics,
-    const cv::Mat &extrinsics) {
+void TargetEstimator::EstimateTargetLocation(cv::Point2i centroid,
+                                             const cv::Mat &intrinsics,
+                                             const cv::Mat &extrinsics,
+                                             TargetEstimate::Builder *builder) {
   const cv::Point2d focal_length(intrinsics.at<double>(0, 0),
                                  intrinsics.at<double>(1, 1));
   const cv::Point2d offset(intrinsics.at<double>(0, 2),
@@ -12,20 +13,19 @@
 
   // Blob pitch in camera reference frame
   const double blob_pitch =
-      std::atan(static_cast<double>(-(blob_point.y - offset.y)) /
+      std::atan(static_cast<double>(-(centroid.y - offset.y)) /
                 static_cast<double>(focal_length.y));
   const double camera_height = extrinsics.at<double>(2, 3);
   // Depth from camera to blob
   const double depth = (kTapeHeight - camera_height) / std::tan(blob_pitch);
 
-  TargetEstimateT target;
-  target.angle_to_target =
-      std::atan2(static_cast<double>(blob_point.x - offset.x),
+  double angle_to_target =
+      std::atan2(static_cast<double>(centroid.x - offset.x),
                  static_cast<double>(focal_length.x));
-  target.distance =
-      (depth / std::cos(target.angle_to_target)) + kUpperHubRadius;
+  double distance = (depth / std::cos(angle_to_target)) + kUpperHubRadius;
 
-  return target;
+  builder->add_angle_to_target(angle_to_target);
+  builder->add_distance(distance);
 }
 
 }  // namespace y2022::vision
diff --git a/y2022/vision/target_estimator.h b/y2022/vision/target_estimator.h
index 4988c3c..9db4121 100644
--- a/y2022/vision/target_estimator.h
+++ b/y2022/vision/target_estimator.h
@@ -10,9 +10,11 @@
  public:
   // Computes the location of the target.
   // blob_point is the mean (x, y) of blob pixels.
-  static TargetEstimateT EstimateTargetLocation(cv::Point2i blob_point,
-                                                const cv::Mat &intrinsics,
-                                                const cv::Mat &extrinsics);
+  // Adds angle_to_target and distance to the given builder.
+  static void EstimateTargetLocation(cv::Point2i centroid,
+                                     const cv::Mat &intrinsics,
+                                     const cv::Mat &extrinsics,
+                                     TargetEstimate::Builder *builder);
 
  private:
   // Height of the center of the tape (m)
diff --git a/y2022/vision/viewer.cc b/y2022/vision/viewer.cc
index 2da16e1..7c13d1b 100644
--- a/y2022/vision/viewer.cc
+++ b/y2022/vision/viewer.cc
@@ -10,7 +10,7 @@
 #include "aos/time/time.h"
 #include "frc971/vision/vision_generated.h"
 #include "y2022/vision/blob_detector.h"
-#include "y2022/vision/target_estimator.h"
+#include "y2022/vision/target_estimate_generated.h"
 
 DEFINE_string(capture, "",
               "If set, capture a single image and save it to this filename.");
@@ -24,6 +24,33 @@
 namespace {
 
 aos::Fetcher<frc971::vision::CameraImage> image_fetcher;
+aos::Fetcher<y2022::vision::TargetEstimate> target_estimate_fetcher;
+
+std::vector<std::vector<cv::Point>> FbsToCvBlobs(
+    const flatbuffers::Vector<flatbuffers::Offset<Blob>> &blobs_fbs) {
+  std::vector<std::vector<cv::Point>> blobs;
+  for (const auto blob : blobs_fbs) {
+    std::vector<cv::Point> points;
+    for (const Point *point : *blob->points()) {
+      points.emplace_back(cv::Point{point->x(), point->y()});
+    }
+    blobs.emplace_back(points);
+  }
+  return blobs;
+}
+
+std::vector<BlobDetector::BlobStats> FbsToBlobStats(
+    const flatbuffers::Vector<flatbuffers::Offset<BlobStatsFbs>>
+        &blob_stats_fbs) {
+  std::vector<BlobDetector::BlobStats> blob_stats;
+  for (const auto stats_fbs : blob_stats_fbs) {
+    cv::Point centroid{stats_fbs->centroid()->x(), stats_fbs->centroid()->y()};
+    blob_stats.emplace_back(BlobDetector::BlobStats{
+        centroid, stats_fbs->aspect_ratio(), stats_fbs->area(),
+        static_cast<size_t>(stats_fbs->num_points())});
+  }
+  return blob_stats;
+}
 
 bool DisplayLoop() {
   int64_t image_timestamp = 0;
@@ -39,6 +66,11 @@
   image_timestamp = image->monotonic_timestamp_ns();
   VLOG(2) << "Got image at timestamp: " << image_timestamp;
 
+  // TODO(Milind) Store the target estimates and match them by timestamp to make
+  // sure we're getting the right one.
+  CHECK(target_estimate_fetcher.FetchNext());
+  const TargetEstimate *target = target_estimate_fetcher.get();
+
   // Create color image:
   cv::Mat image_color_mat(cv::Size(image->cols(), image->rows()), CV_8UC2,
                           (void *)image->data()->data());
@@ -50,17 +82,16 @@
     return false;
   }
 
-  cv::Mat binarized_image;
-  cv::Mat ret_image(cv::Size(image->cols(), image->rows()), CV_8UC3);
-  std::vector<std::vector<cv::Point>> unfiltered_blobs, filtered_blobs;
-  std::vector<BlobDetector::BlobStats> blob_stats;
-  cv::Point centroid;
-  BlobDetector::ExtractBlobs(rgb_image, binarized_image, ret_image,
-                             filtered_blobs, unfiltered_blobs, blob_stats,
-                             centroid);
-
   LOG(INFO) << image->monotonic_timestamp_ns()
-            << ": # blobs: " << filtered_blobs.size();
+            << ": # blobs: " << target->blob_result()->filtered_blobs()->size();
+
+  cv::Mat ret_image;
+  BlobDetector::DrawBlobs(
+      ret_image, FbsToCvBlobs(*target->blob_result()->filtered_blobs()),
+      FbsToCvBlobs(*target->blob_result()->unfiltered_blobs()),
+      FbsToBlobStats(*target->blob_result()->blob_stats()),
+      cv::Point{target->blob_result()->centroid()->x(),
+                target->blob_result()->centroid()->y()});
 
   cv::imshow("image", rgb_image);
   cv::imshow("blobs", ret_image);
@@ -102,36 +133,6 @@
 
   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;
-    std::vector<BlobDetector::BlobStats> 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);
-    cv::Point centroid;
-    BlobDetector::ExtractBlobs(rgb_image, binarized_image, ret_image,
-                               filtered_blobs, unfiltered_blobs, blob_stats,
-                               centroid);
-
-    LOG(INFO) << ": # blobs: " << filtered_blobs.size() << " (# removed: "
-              << unfiltered_blobs.size() - filtered_blobs.size() << ")";
-    cv::imshow("image", rgb_image);
-    cv::imshow("mask", binarized_image);
-    cv::imshow("blobs", ret_image);
-
-    int keystroke = cv::waitKey(0);
-    if ((keystroke & 0xFF) == static_cast<int>('q')) {
-      return;
-    }
-  }
-}
 }  // namespace
 }  // namespace vision
 }  // namespace y2022
@@ -139,8 +140,5 @@
 // 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();
+  y2022::vision::ViewerMain();
 }