leon 1 月之前
父節點
當前提交
2192aa92d3
共有 1 個文件被更改,包括 1 次插入1 次删除
  1. 1 1
      src/infer/trt/yolo.hpp

+ 1 - 1
src/infer/trt/yolo.hpp

@@ -56,7 +56,7 @@ namespace yolo
         void preprocess(int ibatch, affine::LetterBoxMatrix &affine, void *stream = nullptr);
         
     
-        bool load(const std::string &engine_file, ModelType model_type, float confidence_threshold, float nms_threshold);
+        bool load(const std::string &engine_file, ModelType model_type, const std::vector<std::string>& names, float confidence_threshold, float nms_threshold);
     
         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);