Filter blobs based on aspect ratio and size

Signed-off-by: milind-u <milind.upadhyay@gmail.com>
Change-Id: Idad861a777a9f09de7da10792ce5ab73c5daa6d1
diff --git a/y2022/vision/blob_detector.cc b/y2022/vision/blob_detector.cc
index bf5afd7..737363d 100644
--- a/y2022/vision/blob_detector.cc
+++ b/y2022/vision/blob_detector.cc
@@ -1,7 +1,5 @@
 #include "y2022/vision/blob_detector.h"
 
-#include <opencv2/imgproc.hpp>
-
 #include "aos/network/team_number.h"
 
 DEFINE_uint64(green_delta, 50,
@@ -13,7 +11,7 @@
 namespace vision {
 
 cv::Mat BlobDetector::ThresholdImage(cv::Mat rgb_image) {
-  cv::Mat gray_image(cv::Size(rgb_image.cols, rgb_image.rows), CV_8UC1);
+  cv::Mat binarized_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);
@@ -23,14 +21,14 @@
       // 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;
+        binarized_image.at<uint8_t>(row, col) = 255;
       } else {
-        gray_image.at<uint8_t>(row, col) = 0;
+        binarized_image.at<uint8_t>(row, col) = 0;
       }
     }
   }
 
-  return gray_image;
+  return binarized_image;
 }
 
 std::vector<std::vector<cv::Point>> BlobDetector::FindBlobs(
@@ -44,26 +42,60 @@
   return contours;
 }
 
+std::vector<BlobDetector::BlobStats> BlobDetector::ComputeStats(
+    std::vector<std::vector<cv::Point>> blobs) {
+  std::vector<BlobDetector::BlobStats> blob_stats;
+  for (auto blob : blobs) {
+    // Make the blob convex before finding bounding box
+    std::vector<cv::Point> convex_blob;
+    cv::convexHull(blob, convex_blob);
+    auto blob_size = cv::boundingRect(convex_blob).size();
+    cv::Moments moments = cv::moments(convex_blob);
+
+    const auto centroid =
+        cv::Point(moments.m10 / moments.m00, moments.m01 / moments.m00);
+    const double aspect_ratio =
+        static_cast<double>(blob_size.width) / blob_size.height;
+    const double area = moments.m00;
+    const size_t points = blob.size();
+
+    blob_stats.emplace_back(BlobStats{centroid, aspect_ratio, area, points});
+  }
+  return blob_stats;
+}
+
 // 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>> blobs,
+    std::vector<BlobDetector::BlobStats> blob_stats) {
   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";
+  auto blob_it = blobs.begin();
+  auto stats_it = blob_stats.begin();
+  while (blob_it < blobs.end() && stats_it < blob_stats.end()) {
+    constexpr int kMaxY = 400;
+    constexpr double kTapeAspectRatio = 5.0 / 2.0;
+    constexpr double kAspectRatioThreshold = 1.5;
+    constexpr double kMinArea = 10;
+    constexpr size_t kMinPoints = 2;
+    // Remove all blobs that are at the bottom of the image, have a different
+    // aspect ratio than the tape, or have too little area or points
+    if ((stats_it->centroid.y <= kMaxY) &&
+        (std::abs(kTapeAspectRatio - stats_it->aspect_ratio) <
+         kAspectRatioThreshold) &&
+        (stats_it->area >= kMinArea) && (stats_it->points >= kMinPoints)) {
+      filtered_blobs.push_back(*blob_it);
     }
+    blob_it++;
+    stats_it++;
   }
   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) {
+    cv::Mat view_image,
+    const std::vector<std::vector<cv::Point>> &unfiltered_blobs,
+    const std::vector<std::vector<cv::Point>> &filtered_blobs,
+    const std::vector<BlobStats> &blob_stats) {
   CHECK_GT(view_image.cols, 0);
   if (unfiltered_blobs.size() > 0) {
     // Draw blobs unfilled, with red color border
@@ -72,26 +104,23 @@
 
   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;
+  for (auto stats : blob_stats) {
+    cv::circle(view_image, stats.centroid, 2, cv::Scalar(255, 0, 0),
+               cv::FILLED);
+  }
 }
 
 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) {
+    std::vector<BlobStats> &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);
+  blob_stats = ComputeStats(unfiltered_blobs);
+  filtered_blobs = FilterBlobs(unfiltered_blobs, blob_stats);
+  DrawBlobs(blob_image, unfiltered_blobs, filtered_blobs, blob_stats);
 }
 
 }  // namespace vision
diff --git a/y2022/vision/blob_detector.h b/y2022/vision/blob_detector.h
index 3da3d4b..84e504d 100644
--- a/y2022/vision/blob_detector.h
+++ b/y2022/vision/blob_detector.h
@@ -1,6 +1,7 @@
 #ifndef Y2022_BLOB_DETECTOR_H_
 #define Y2022_BLOB_DETECTOR_H_
 
+#include <opencv2/features2d.hpp>
 #include <opencv2/imgproc.hpp>
 
 namespace y2022 {
@@ -8,6 +9,13 @@
 
 class BlobDetector {
  public:
+  struct BlobStats {
+    cv::Point centroid;
+    double aspect_ratio;
+    double area;
+    size_t points;
+  };
+
   BlobDetector() {}
   // Given an image, threshold it to find "green" pixels
   // Input: Color image
@@ -17,25 +25,28 @@
   // Given binary image, extract blobs
   static std::vector<std::vector<cv::Point>> FindBlobs(cv::Mat threshold_image);
 
+  // Extract stats for each blob
+  static std::vector<BlobStats> ComputeStats(
+      std::vector<std::vector<cv::Point>> blobs);
+
   // 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);
+      std::vector<std::vector<cv::Point>> blobs,
+      std::vector<BlobStats> blob_stats);
 
   // 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 DrawBlobs(
+      cv::Mat view_image,
+      const std::vector<std::vector<cv::Point>> &filtered_blobs,
+      const std::vector<std::vector<cv::Point>> &unfiltered_blobs,
+      const std::vector<BlobStats> &blob_stats);
 
   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);
+      std::vector<BlobStats> &blob_stats);
 };
 }  // namespace vision
 }  // namespace y2022
diff --git a/y2022/vision/viewer.cc b/y2022/vision/viewer.cc
index 8959f53..6573308 100644
--- a/y2022/vision/viewer.cc
+++ b/y2022/vision/viewer.cc
@@ -15,9 +15,7 @@
               "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_string(png_dir, "LED_Ring_exp", "Path to a set of images to display.");
 DEFINE_bool(show_features, true, "Show the blobs.");
 
 namespace y2022 {
@@ -52,8 +50,8 @@
   }
 
   cv::Mat binarized_image, ret_image;
-  std::vector<std::vector<cv::Point>> unfiltered_blobs, filtered_blobs,
-      blob_stats;
+  std::vector<std::vector<cv::Point>> unfiltered_blobs, filtered_blobs;
+  std::vector<BlobDetector::BlobStats> blob_stats;
   BlobDetector::ExtractBlobs(rgb_image, binarized_image, ret_image,
                              filtered_blobs, unfiltered_blobs, blob_stats);
 
@@ -112,8 +110,8 @@
   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;
+    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 =
@@ -127,7 +125,7 @@
     cv::imshow("image", rgb_image);
     cv::imshow("blobs", ret_image);
 
-    int keystroke = cv::waitKey(1000);
+    int keystroke = cv::waitKey(0);
     if ((keystroke & 0xFF) == static_cast<int>('q')) {
       return;
     }