瀏覽代碼

add config for draw image

leon 2 周之前
父節點
當前提交
14ff8f00f7
共有 3 個文件被更改,包括 130 次插入49 次删除
  1. 1 1
      src/main.cpp
  2. 81 42
      src/nodes/draw/drawNode.cpp
  3. 48 6
      src/pipeline/pipeline.hpp

+ 1 - 1
src/main.cpp

@@ -30,7 +30,7 @@ int main(int argc , char* argv[])
             std::cout << "  - Node Name: " << node->get_name() << std::endl;
         }
     }
-    pipeline_manager.start_pipelines();
+    pipeline_manager.start_all_pipelines();
     pipeline_manager.show_all_pipeline();
     getchar();
     return 0;

+ 81 - 42
src/nodes/draw/drawNode.cpp

@@ -120,15 +120,27 @@ static float box_iou(const data::Box& box1, const data::Box& box2)
 void DrawNode::handle_data(std::shared_ptr<meta::BaseData>& meta_data) 
 {
     auto frame_data = std::dynamic_pointer_cast<meta::FrameData>(meta_data);
-    bool show_track = true;
+    if (config_data_ == nullptr) 
+    {
+        PLOGE.printf("DrawNode : [%s] config data is null, show nothing", name_.c_str());
+        return;
+    }
+    if (config_data_->data_name != "DrawConfigData") 
+    {
+        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;
+
     cv::Mat image = frame_data->image.clone();
     int image_width = image.cols;
     int image_height = image.rows;
     PositionManager<float> pm(getFontSize);
-    for (auto& box : frame_data->boxes)
+    if (show_original_result) 
     {
-        if (show_track)
-        {
+      for (auto& box : frame_data->boxes)
+      {
           float max_iou = 0.0f;
           for (const auto& track_box : frame_data->track_boxes)
           {
@@ -139,52 +151,79 @@ void DrawNode::handle_data(std::shared_ptr<meta::BaseData>& meta_data)
               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);
+  
+          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);
+                  }
+              }
+          }
+      }
+    }
 
-        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())
+    if (show_final_result) 
+    {
+        for (auto& box : frame_data->result)
         {
-            for (const auto& point : box.keypoints)
+            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) 
             {
-                cv::circle(image, cv::Point(point.x, point.y), 5, cv::Scalar(b, g, r), -1);
+                text = str_format("%s %.2f id=%d", box.label.c_str(), box.score, box.class_id);
             }
-            for (const auto& pair : coco_pairs) 
+            else
             {
-                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);
-                }
+                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);
         }
     }
+    
     frame_data->draw_image = image;
 }
 

+ 48 - 6
src/pipeline/pipeline.hpp

@@ -45,6 +45,46 @@ public:
         return shared_models_;
     }
 
+    bool set_config(const std::string& pipeline_id, const std::string& node_id, const meta::DrawConfigData& config) 
+    {
+        auto instance = std::find_if(configured_pipelines_.begin(), configured_pipelines_.end(), 
+            [&pipeline_id](const PipelineInstance& instance) { return instance.pipeline_id == pipeline_id; });
+        if (instance == configured_pipelines_.end())
+        {
+            std::cerr << "Pipeline with ID " << pipeline_id << " not found." << std::endl;
+            return false;
+        }
+        auto node = std::find_if(instance->nodes.begin(), instance->nodes.end(), 
+            [&node_id](const std::shared_ptr<GNode::BaseNode>& node) { return node->get_name() == node_id; });
+        if (node == instance->nodes.end())
+        {
+            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));
+        return true;
+    }
+
+    bool set_config(const std::string& pipeline_id, const std::string& node_id, const meta::AnalyzeConfigData& config) 
+    {
+        auto instance = std::find_if(configured_pipelines_.begin(), configured_pipelines_.end(), 
+            [&pipeline_id](const PipelineInstance& instance) { return instance.pipeline_id == pipeline_id; });
+        if (instance == configured_pipelines_.end())
+        {
+            std::cerr << "Pipeline with ID " << pipeline_id << " not found." << std::endl;
+            return false;
+        }
+        auto node = std::find_if(instance->nodes.begin(), instance->nodes.end(), 
+            [&node_id](const std::shared_ptr<GNode::BaseNode>& node) { return node->get_name() == node_id; });
+        if (node == instance->nodes.end())
+        {
+            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));
+        return true;
+    }
+
     void start_all_pipelines() 
     {
       for (auto& instance : configured_pipelines_) 
@@ -111,14 +151,15 @@ public:
 
     void show_all_pipeline()
     {
-        for (const auto& pipeline : configured_pipelines_)
+        for (const auto& instance : configured_pipelines_)
         {
-            printf("Pipeline ID: %s\nDescription: %s\n", pipeline.pipeline_id.c_str(), pipeline.description.c_str());
-            for (const auto& node : instance->nodes)
+            printf("Pipeline ID: %s\nDescription: %s\n", instance.pipeline_id.c_str(), instance.description.c_str());
+            printf("start --> ");
+            for (const auto& node : instance.nodes)
             {
                 printf("%s --> ", node->get_name().c_str());
             }
-            printf("\n");
+            printf(" end\n");
         }
     }
 
@@ -131,12 +172,13 @@ public:
             std::cerr << "Pipeline with ID " << pipeline_id << " not found." << std::endl;
             return;
         }
-        printf("Pipeline ID: %s\nDescription: %s\n", pipeline.pipeline_id.c_str(), pipeline.description.c_str());
+        printf("Pipeline ID: %s\nDescription: %s\n", instance->pipeline_id.c_str(), instance->description.c_str());
+        printf("start --> ");
         for (const auto& node : instance->nodes)
         {
             printf("%s --> ", node->get_name().c_str());
         }
-        printf("\n");
+        printf(" end\n");
     }