add draw_bboxes function

StrayedCats · Mar 30, 2024 · e30b85f · e30b85f
1 parent aee6e36
commit e30b85f
Showing 1 changed file with 25 additions and 0 deletions.
diff --git a/detector2d_base/include/detector2d_base/detector2d_base.hpp b/detector2d_base/include/detector2d_base/detector2d_base.hpp
@@ -28,6 +28,31 @@ class Detector
   virtual vision_msgs::msg::Detection2DArray detect(const cv::Mat & image) = 0;
   virtual ~Detector() {}
 
+  cv::Mat3b draw_bboxes(
+      const cv::Mat & frame,
+      const vision_msgs::msg::Detection2DArray & boxes)
+  {
+      cv::Mat3b frame_out;
+      if (frame.channels() == 1)
+      {
+          cv::cvtColor(frame, frame_out, cv::COLOR_GRAY2BGR);
+      }
+      else {
+          frame_out = frame;
+      }
+
+      for (auto detection : boxes.detections)
+      {
+          cv::Rect bbox;
+          bbox.x = detection.bbox.center.position.x - detection.bbox.size_x / 2;
+          bbox.y = detection.bbox.center.position.y - detection.bbox.size_y / 2;
+          bbox.width = detection.bbox.size_x;
+          bbox.height = detection.bbox.size_y;
+          cv::rectangle(frame_out, bbox, cv::Scalar(0, 255, 0), 2);
+      }
+      return frame_out;
+  }
+
 protected:
   Detector() {}
 };