优化展示

src/nodes/draw/drawNode.cpp

@@ -117,6 +117,120 @@ static float box_iou(const data::Box& box1, const data::Box& box2)
     return inter_area / (box1_area + box2_area - inter_area + 1e-6f); // Avoid division by zero
 }
 
+void DrawNode::draw_track_boxes(cv::Mat& image, const data::BoxArray& track_boxes)
+{
+  auto now = std::chrono::steady_clock::now();
+
+  for (const auto& box : track_boxes) 
+  {
+    uint8_t b, g, r;
+    std::tie(b, g, r) = random_color(box.class_id);
+
+    // Draw bounding box
+    // cv::rectangle(image, cv::Point(box.left, box.top), cv::Point(box.right, box.bottom), cv::Scalar(b, g, r), 2);
+
+    // Calculate center point of the box
+    cv::Point center((box.left + box.right) / 2, (box.top + box.bottom) / 2);
+    std::string track_text = str_format("ID: %d", box.class_id);
+    cv::putText(image, track_text, center, cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(b, g, r), 2);
+    track_history_[box.class_id].push_back(center);
+
+    if (track_history_[box.class_id].size() > 50) 
+    {
+      track_history_[box.class_id].erase(track_history_[box.class_id].begin());
+    }
+
+    // Draw the track lines
+    const auto& history = track_history_[box.class_id];
+    for (size_t i = 1; i < history.size(); ++i) 
+    {
+      cv::line(image, history[i - 1], history[i], cv::Scalar(b, g, r), 2);
+    }
+
+    // Update the last seen time for this class ID
+    last_track_update_time_[box.class_id] = now;
+  }
+
+  // Remove outdated tracks
+  for (auto it = track_history_.begin(); it != track_history_.end(); ) 
+  {
+    if (last_track_update_time_.find(it->first) != last_track_update_time_.end() &&
+      std::chrono::duration_cast<std::chrono::seconds>(now - last_track_update_time_[it->first]).count() > 5) 
+    {
+      last_track_update_time_.erase(it->first);
+      it = track_history_.erase(it);
+    } 
+    else 
+    {
+      ++it;
+    }
+  }
+}
+
+void DrawNode::draw_boxes(cv::Mat& image, const data::BoxArray& boxes)
+{
+  int image_width = image.cols;
+  int image_height = image.rows;
+  PositionManager<float> pm(getFontSize);
+  for (const auto& box : boxes) 
+  {
+    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);
+    std::tuple<int, int, int, int> pbox = std::make_tuple(box.left, box.top, box.right, box.bottom);
+    int x, y;
+    std::string text = str_format("%s %.2f", box.label.c_str(), box.score);
+    std::tie(x, y) = pm.selectOptimalPosition(pbox, image_width, image_height, text);
+    
+    cv::putText(image, text, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(b, g, r), 2);
+    if (!box.seg_mask.empty())
+    {
+        overlay_mask(image, box.seg_mask, box.left, box.top, cv::Scalar(b, g, r), 0.6);
+    }
+    if (!box.keypoints.empty())
+    {
+        for (const auto& point : box.keypoints)
+        {
+            cv::circle(image, cv::Point(point.x, point.y), 5, cv::Scalar(b, g, r), -1);
+        }
+        for (const auto& pair : coco_pairs) 
+        {
+            int startIdx = pair.first;
+            int endIdx = pair.second;
+
+            if (startIdx < box.keypoints.size() && endIdx < box.keypoints.size()) 
+            {
+                int x1 = (int)box.keypoints[startIdx].x;
+                int y1 = (int)box.keypoints[startIdx].y;
+                int x2 = (int)box.keypoints[endIdx].x;
+                int y2 = (int)box.keypoints[endIdx].y;
+                cv::line(image, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(b, g, r), 2);
+            }
+        }
+    }
+  }  
+}
+
+void DrawNode::draw_result(cv::Mat& image, const data::BoxArray& result)
+{
+  int image_width = image.cols;
+  int image_height = image.rows;
+  for (auto& box : result)
+  {
+      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);
+
+      std::tuple<int, int, int, int> pbox = std::make_tuple(box.left, box.top, box.right, box.bottom);
+      int x, y;
+      std::string text = str_format("%s %.2f", box.label.c_str(), box.score);
+      std::tie(x, y) = pm.selectOptimalPosition(pbox, image_width, image_height, text);
+      
+      cv::putText(image, text, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 1.0f, cv::Scalar(b, g, r), 2);
+  }
+}
+
+
 void DrawNode::handle_data(std::shared_ptr<meta::BaseData>& meta_data) 
 {
     auto frame_data = std::dynamic_pointer_cast<meta::FrameData>(meta_data);
@@ -130,98 +244,23 @@ void DrawNode::handle_data(std::shared_ptr<meta::BaseData>& meta_data)
         PLOGE.printf("DrawNode : [%s] Invalid config data %s", name_.c_str(), config_data_->data_name.c_str());
         return;
     }
-    bool show_final_result = config_data_->show_final_result;
-    bool show_original_result = config_data_->show_original_result;
+    auto config_data = std::dynamic_pointer_cast<meta::DrawConfigData>(config_data_);
+    bool show_final_result = config_data->show_final_result;
+    bool show_original_result = config_data->show_original_result;
 
     cv::Mat image = frame_data->image.clone();
     int image_width = image.cols;
     int image_height = image.rows;
-    PositionManager<float> pm(getFontSize);
+    
     if (show_original_result) 
     {
-      for (auto& box : frame_data->boxes)
-      {
-          float max_iou = 0.0f;
-          for (const auto& track_box : frame_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);
-  
-          std::tuple<int, int, int, int> pbox = std::make_tuple(box.left, box.top, box.right, box.bottom);
-          int x, y;
-          std::string text;
-          if (box.class_id >= 0) 
-          {
-              text = str_format("%s %.2f id=%d", box.label.c_str(), box.score, box.class_id);
-          }
-          else
-          {
-              text = str_format("%s %.2f", box.label.c_str(), box.score);
-          }
-          std::tie(x, y) = pm.selectOptimalPosition(pbox, image_width, image_height, text);
-          
-          cv::putText(image, text, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(b, g, r), 2);
-  
-          if (!box.seg_mask.empty())
-          {
-              overlay_mask(image, box.seg_mask, box.left, box.top, cv::Scalar(b, g, r), 0.6);
-          }
-          if (!box.keypoints.empty())
-          {
-              for (const auto& point : box.keypoints)
-              {
-                  cv::circle(image, cv::Point(point.x, point.y), 5, cv::Scalar(b, g, r), -1);
-              }
-              for (const auto& pair : coco_pairs) 
-              {
-                  int startIdx = pair.first;
-                  int endIdx = pair.second;
-  
-                  if (startIdx < box.keypoints.size() && endIdx < box.keypoints.size()) 
-                  {
-                      int x1 = (int)box.keypoints[startIdx].x;
-                      int y1 = (int)box.keypoints[startIdx].y;
-                      int x2 = (int)box.keypoints[endIdx].x;
-                      int y2 = (int)box.keypoints[endIdx].y;
-                      cv::line(image, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(b, g, r), 2);
-                  }
-              }
-          }
-      }
+        draw_boxes(image, frame_data->boxes);
+        draw_track_boxes(image, frame_data->track_boxes);
     }
 
     if (show_final_result) 
     {
-        for (auto& box : frame_data->result)
-        {
-            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);
-  
-            std::tuple<int, int, int, int> pbox = std::make_tuple(box.left, box.top, box.right, box.bottom);
-            int x, y;
-            std::string text;
-            if (box.class_id >= 0) 
-            {
-                text = str_format("%s %.2f id=%d", box.label.c_str(), box.score, box.class_id);
-            }
-            else
-            {
-                text = str_format("%s %.2f", box.label.c_str(), box.score);
-            }
-            std::tie(x, y) = pm.selectOptimalPosition(pbox, image_width, image_height, text);
-            
-            cv::putText(image, text, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 1.0f, cv::Scalar(b, g, r), 2);
-        }
+        draw_result(image, frame_data->result);
     }
     
     frame_data->draw_image = image;

src/nodes/draw/drawNode.hpp

@@ -2,6 +2,8 @@
 #define DRAWNODE_HPP__
 #include "nodes/base/base.hpp"
 #include <opencv2/opencv.hpp>
+#include <unordered_map>
+
 // 日志库
 #include "plog/Log.h"
 #include "plog/Initializers/RollingFileInitializer.h"
@@ -18,6 +20,15 @@ public:
     virtual ~DrawNode() { stop(); };
 
     void handle_data(std::shared_ptr<meta::BaseData>& meta_data) override;
+
+private:
+    void draw_boxes(cv::Mat& image, const data::BoxArray& boxes);
+    void draw_track_boxes(cv::Mat& image, const data::BoxArray& track_boxes);
+    void draw_result(cv::Mat& image, const data::BoxArray& result);
+
+private:
+    std::unordered_map<int, std::vector<cv::Point>> track_history_;
+    std::unordered_map<int, std::chrono::steady_clock::time_point> last_track_update_time_;
 };
 
 } // namespace Node

src/pipeline/pipeline.hpp

@@ -61,7 +61,7 @@ public:
             std::cerr << "Node with ID " << node_id << " not found in pipeline " << pipeline_id << "." << std::endl;
             return false;
         }
-        node->set_config_data(std::make_shared<meta::DrawConfigData>(config));
+        (*node)->set_config_data(std::make_shared<meta::DrawConfigData>(config));
         return true;
     }
 
@@ -81,7 +81,7 @@ public:
             std::cerr << "Node with ID " << node_id << " not found in pipeline " << pipeline_id << "." << std::endl;
             return false;
         }
-        node->set_config_data(std::make_shared<meta::AnalyzeConfigData>(config));
+        (*node)->set_config_data(std::make_shared<meta::AnalyzeConfigData>(config));
         return true;
     }