leon 1 miesiąc temu
rodzic
commit
c7174cc1e7
2 zmienionych plików z 21 dodań i 3 usunięć
  1. 2 2
      src/infer/trt/yolo.cu
  2. 19 1
      src/main.cpp

+ 2 - 2
src/infer/trt/yolo.cu

@@ -1,8 +1,8 @@
 #include "model/yolo.hpp"
 #include <vector>
 #include <memory>
-#include "slice/slice.hpp"
-#include "model/affine.hpp"
+#include "infer/slice/slice.hpp"
+#include "infer/trt/affine.hpp"
 #include "common/check.hpp"
 
 #define GPU_BLOCK_THREADS 512

+ 19 - 1
src/main.cpp

@@ -5,9 +5,25 @@
 #include "nodes/httpPush/httpPush.hpp"
 
 
-
 int main()
 {
+    // create trt model
+    std::vector<std::string> names = {
+        "person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck",
+        "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
+        "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
+        "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
+        "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+        "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
+        "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
+        "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa",
+        "pottedplant", "bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse",
+        "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
+        "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
+        "toothbrush"
+    }; 
+    std::shared_ptr<Infer> model = load("model/yolo11s.engine", ModelType::YOLO11, names, 0, 0.5, 0.45);
+
     std::shared_ptr<Node::StreamNode> src_node0   = std::make_shared<Node::StreamNode>("src0", "rtsp://admin:lww123456@172.16.22.16:554/Streaming/Channels/101");
     src_node0->set_skip_frame(10);
 
@@ -19,6 +35,8 @@ int main()
 
 
     std::shared_ptr<Node::InferNode> infer_node   = std::make_shared<Node::InferNode>("infer");
+    infer_node->set_model_instance(model);
+
     std::shared_ptr<Node::DrawNode> draw_node     = std::make_shared<Node::DrawNode>("draw");
     std::shared_ptr<Node::HttpPushNode> push_node = std::make_shared<Node::HttpPushNode>("push", "172.16.20.168", 8080, "/push");
     Node::LinkNode(src_node0, infer_node);