leon 3 долоо хоног өмнө
parent
commit
7d798cc44f

+ 14 - 0
src/nodes/draw/drawNode.cpp

@@ -104,12 +104,26 @@ static std::tuple<uint8_t, uint8_t, uint8_t> random_color(int id)
 
 void DrawNode::handle_data(std::shared_ptr<meta::MetaData>& meta_data) 
 {
+    bool show_track = true;
     cv::Mat image = meta_data->image.clone();
     int image_width = image.cols;
     int image_height = image.rows;
     PositionManager<float> pm(getFontSize);
     for (auto& box : meta_data->boxes)
     {
+        if (show_track)
+        {
+          float max_iou = 0.0f;
+          for (const auto& track_box : meta_data->track_boxes)
+          {
+            float iou = box_iou(box, track_box);
+            if (iou > max_iou && iou > 0.7f)
+            {
+              max_iou = iou;
+              box.class_id = track_box.class_id;
+            }   
+          }
+        }
         uint8_t b, g, r;
         std::tie(b, g, r) = random_color(box.class_id);
         cv::rectangle(image, cv::Point(box.left, box.top), cv::Point(box.right, box.bottom), cv::Scalar(b, g, r), 2);