@@ -172,6 +172,13 @@ public:
}
};
+std::shared_ptr<Infer> load(const std::string &engine_file, YoloType yolo_type, int gpu_id, float confidence_threshold, float nms_threshold)
+{
+ // checkRuntime(cudaSetDevice(gpu_id));
+ return std::shared_ptr<Yolov5InferImpl>((Yolov5InferImpl *)(new Yolov5InferImpl(engine_file, yolo_type, confidence_threshold, nms_threshold)));
+}
+
#endif // YOLO_HPP__
