infer 类加锁。保证模型复用

src/infer/infer.hpp

@@ -17,6 +17,9 @@ class Infer{
 public:
     virtual data::BoxArray forward(const tensor::Image &image, int slice_width, int slice_height, float overlap_width_ratio, float overlap_height_ratio, void *stream = nullptr) = 0;
     virtual data::BoxArray forward(const tensor::Image &image, void *stream = nullptr) = 0;
+
+protected:
+    std::mutex mutex_;
 };
 
 std::shared_ptr<Infer> load(const std::string& model_path, ModelType model_type, const std::vector<std::string>& names, int gpu_id=0, float confidence_threshold=0.5f, float nms_threshold=0.45f);

src/infer/trt/depth_any/depth.cu

@@ -0,0 +1,30 @@
+#include "infer/trt/depth_any/depth.hpp"
+#include <vector>
+#include <memory>
+#include "infer/trt/affine.hpp"
+#include "common/check.hpp"
+
+
+namespace depth
+{
+
+
+bool DepthModelImpl::load(const std::string &engine_file, ModelType model_type, const std::vector<std::string>& names, float confidence_threshold, float nms_threshold)
+{
+    trt_ = TensorRT::load(engine_file);
+    if (trt_ == nullptr) return false;
+
+    trt_->print();
+
+    auto input_dim = trt_->static_dims(0);
+    bbox_head_dims_ = trt_->static_dims(1);
+    network_input_width_ = input_dim[3];
+    network_input_height_ = input_dim[2];
+    isdynamic_model_ = trt_->has_dynamic_dim();
+
+    normalize_ = affine::Norm::alpha_beta(1 / 255.0f, 0.0f, affine::ChannelType::SwapRB);
+    return true;
+}
+
+
+}   // namespace depth

src/infer/trt/depth_any/depth.hpp

@@ -0,0 +1,55 @@
+#ifndef DEPTH_HPP__
+#define DEPTH_HPP__
+
+#include <vector>
+#include <iomanip>
+#include "common/memory.hpp"
+#include "infer/infer.hpp"
+#include "common/image.hpp"
+#include "common/data.hpp"
+#include "infer/trt/affine.hpp"
+
+namespace depth
+{
+
+    class DepthModelImpl : public Infer 
+    {
+    public:
+        ModelType model_type_;
+    
+        std::shared_ptr<TensorRT::Engine> trt_;
+        std::string engine_file_;
+    
+        tensor::Memory<int> box_count_;
+    
+        tensor::Memory<float> affine_matrix_;
+        tensor::Memory<float>  input_buffer_, bbox_predict_, output_boxarray_;
+    
+        int network_input_width_, network_input_height_;
+        affine::Norm normalize_;
+
+        bool isdynamic_model_ = false;
+    
+        DepthModelImpl() = default;
+    
+        virtual ~DepthModelImpl() = default;
+    
+        void adjust_memory(int batch_size);
+    
+        void preprocess(int ibatch, affine::LetterBoxMatrix &affine, void *stream = nullptr);
+        
+    
+        bool load(const std::string &engine_file, ModelType model_type, const std::vector<std::string>& names, float confidence_threshold, float nms_threshold);
+    
+        virtual cv::Mat forward(const tensor::Image &image, void *stream = nullptr);
+    
+};
+
+Infer *loadraw(const std::string &engine_file, ModelType model_type, const std::vector<std::string>& names, float confidence_threshold,
+    float nms_threshold);
+
+std::shared_ptr<Infer> load_depth(const std::string &engine_file, ModelType model_type, const std::vector<std::string>& names, int gpu_id, float confidence_threshold, float nms_threshold);
+
+}   // namespace depth
+
+#endif

src/infer/trt/yolo.cu → src/infer/trt/yolo/yolo.cu

@@ -1,6 +1,7 @@
-#include "infer/trt/yolo.hpp"
 #include <vector>
+#include <mutex>
 #include <memory>
+#include "infer/trt/yolo/yolo.hpp"
 #include "infer/slice/slice.hpp"
 #include "infer/trt/affine.hpp"
 #include "common/check.hpp"
@@ -399,12 +400,14 @@ bool YoloModelImpl::load(const std::string &engine_file, ModelType model_type, c
 
 data::BoxArray YoloModelImpl::forward(const tensor::Image &image, int slice_width, int slice_height, float overlap_width_ratio, float overlap_height_ratio, void *stream)
 {
+    std::lock_guard<std::mutex> lock(mutex_); // 自动加锁/解锁
     slice_->slice(image, slice_width, slice_height, overlap_width_ratio, overlap_height_ratio, stream);
     return forwards(stream);
 }
 
 data::BoxArray YoloModelImpl::forward(const tensor::Image &image, void *stream)
 {
+    std::lock_guard<std::mutex> lock(mutex_); // 自动加锁/解锁
     slice_->autoSlice(image, stream);
     return forwards(stream);
 }

src/main.cpp

@@ -56,8 +56,7 @@ int main()
 }
 
 
-// TODO: 模型多路复用
-// TODO: 通过配置文件创建 pipeline
-// TODO：日志系统
-// TODO: 结果推送接口 http mqtt ...
-// TODO: 电子围栏
+// TODO
+// 通过配置文件创建 pipeline
+// 日志
+// 设置电子围栏