Make api of yolov5 be same bettwen Python and C++

fastdeploy/vision/evaluation/classify.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
-import fastdeploy as fd
 import cv2
 import os
 import re

fastdeploy/vision/ultralytics/__init__.py

@@ -22,19 +22,14 @@ class YOLOv5(fd.FastDeployModel):
     def __init__(self,
                  model_file,
                  params_file=None,
-                 backend_option=None,
+                 runtime_option=None,
                  model_format=fd.Frontend.ONNX):
         # 调用基函数进行backend_option的初始化
-        # 初始化后的option保存在self._option
-        super(YOLOv5, self).__init__(backend_option)
+        # 初始化后的option保存在self._runtime_option
+        super(YOLOv5, self).__init__(runtime_option)
 
-        if model_format == fd.Frontend.ONNX:
-            # 加载的模型需要保存在成员变量self._model中
-            self._model = C.vision.ultralytics.YOLOv5(
-                model_file, self._runtime_option, model_format)
-        elif model_format == fd.Frontend.PADDLE:
-            self._model = C.vision.ultralytics.YOLOv5(
-                model_file, params_file, self._runtime_option, model_format)
+        self._model = C.vision.ultralytics.YOLOv5(
+            model_file, params_file, self._runtime_option, model_format)
         # 通过self.initialized判断整个模型的初始化是否成功
         assert self.initialized, "YOLOv5 initialize failed."
 

fastdeploy/vision/ultralytics/ultralytics_pybind.cc

@@ -20,7 +20,6 @@ void BindUltralytics(pybind11::module& m) {
       m.def_submodule("ultralytics", "https://github.com/ultralytics/yolov5");
   pybind11::class_<vision::ultralytics::YOLOv5, FastDeployModel>(
       ultralytics_module, "YOLOv5")
-      .def(pybind11::init<std::string, RuntimeOption, Frontend>())
       .def(pybind11::init<std::string, std::string, RuntimeOption, Frontend>())
       .def("predict",
            [](vision::ultralytics::YOLOv5& self, pybind11::array& data,
@@ -37,4 +36,4 @@ void BindUltralytics(pybind11::module& m) {
       .def_readwrite("is_no_pad", &vision::ultralytics::YOLOv5::is_no_pad)
       .def_readwrite("is_scale_up", &vision::ultralytics::YOLOv5::stride);
 }
-} // namespace fastdeploy
+}  // namespace fastdeploy

fastdeploy/vision/ultralytics/yolov5.cc

@@ -41,25 +41,16 @@ void LetterBox(Mat* mat, std::vector<int> size, std::vector<float> color,
   }
 }
 
-YOLOv5::YOLOv5(const std::string& model_file,
-               const RuntimeOption& custom_option,
-               const Frontend& model_format) {
-  valid_cpu_backends = {Backend::ORT};  // 指定可用的CPU后端
-  valid_gpu_backends = {Backend::ORT};  // 指定可用的GPU后端
-  runtime_option = custom_option;
-  runtime_option.model_format = model_format;  // 指定模型格式
-  runtime_option.model_file = model_file;
-  // initialized用于标记模型是否初始化成功
-  // C++或Python中可调用YOLOv5.Intialized() /
-  // YOLOv5.initialized()判断模型是否初始化成功
-  initialized = Initialize();
-}
-
 YOLOv5::YOLOv5(const std::string& model_file, const std::string& params_file,
                const RuntimeOption& custom_option,
                const Frontend& model_format) {
-  valid_cpu_backends = {Backend::PDINFER};  // 指定可用的CPU后端
-  valid_gpu_backends = {Backend::PDINFER};  // 指定可用的GPU后端
+  if (model_format == Frontend::ONNX) {
+    valid_cpu_backends = {Backend::ORT};  // 指定可用的CPU后端
+    valid_gpu_backends = {Backend::ORT, Backend::TRT};  // 指定可用的GPU后端
+  } else {
+    valid_cpu_backends = {Backend::PDINFER, Backend::ORT};
+    valid_gpu_backends = {Backend::PDINFER, Backend::ORT, Backend::TRT};
+  }
   runtime_option = custom_option;
   runtime_option.model_format = model_format;
   runtime_option.model_file = model_file;

fastdeploy/vision/ultralytics/yolov5.h

@@ -14,8 +14,8 @@
 
 #pragma once
 #include "fastdeploy/fastdeploy_model.h"
-#include "fastdeploy/vision/common/result.h"
 #include "fastdeploy/vision/common/processors/transform.h"
+#include "fastdeploy/vision/common/result.h"
 
 namespace fastdeploy {
 namespace vision {
@@ -23,15 +23,11 @@ namespace ultralytics {
 
 class FASTDEPLOY_DECL YOLOv5 : public FastDeployModel {
  public:
-  // 支持ONNX格式模型的输入
-  YOLOv5(const std::string& model_file,
-         const RuntimeOption& option = RuntimeOption(),
-         const Frontend& model_format = Frontend::ONNX);
-
-  // 在X2Paddle转成Paddle后，支持Paddle格式模型的输入
-  YOLOv5(const std::string& model_file, const std::string& params_file,
+  // 当model_format为ONNX时，无需指定params_file
+  // 当model_format为Paddle时，则需同时指定model_file & params_file
+  YOLOv5(const std::string& model_file, const std::string& params_file = "",
          const RuntimeOption& custom_option = RuntimeOption(),
-         const Frontend& model_format = Frontend::PADDLE);
+         const Frontend& model_format = Frontend::ONNX);
 
   // 定义模型的名称
   virtual std::string ModelName() const { return "ultralytics/yolov5"; }
@@ -52,10 +48,10 @@ class FASTDEPLOY_DECL YOLOv5 : public FastDeployModel {
   // im_info 为预处理记录的信息，后处理用于还原box
   // conf_threshold 后处理时过滤box的置信度阈值
   // nms_iou_threshold 后处理时NMS设定的iou阈值
-  virtual bool
-  Postprocess(FDTensor& infer_result, DetectionResult* result,
-              const std::map<std::string, std::array<float, 2>>& im_info,
-              float conf_threshold, float nms_iou_threshold);
+  virtual bool Postprocess(
+      FDTensor& infer_result, DetectionResult* result,
+      const std::map<std::string, std::array<float, 2>>& im_info,
+      float conf_threshold, float nms_iou_threshold);
 
   // 模型预测接口，即用户调用的接口
   // im 为用户的输入数据，目前对于CV均定义为cv::Mat
@@ -84,6 +80,6 @@ class FASTDEPLOY_DECL YOLOv5 : public FastDeployModel {
   // padding stride, for is_mini_pad
   int stride;
 };
-} // namespace ultralytics
-} // namespace vision
-} // namespace fastdeploy
+}  // namespace ultralytics
+}  // namespace vision
+}  // namespace fastdeploy

requirements.txt

@@ -1 +1,2 @@
 opencv-python
+tqdm