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,17 +22,12 @@ 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
+        # 初始化后的option保存在self._runtime_option
-        super(YOLOv5, self).__init__(backend_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.initialized判断整个模型的初始化是否成功

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,

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后端
+  if (model_format == Frontend::ONNX) {
-  valid_gpu_backends = {Backend::PDINFER};  // 指定可用的GPU后端
+    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格式模型的输入
+  // 当model_format为ONNX时，无需指定params_file
-  YOLOv5(const std::string& model_file,
+  // 当model_format为Paddle时，则需同时指定model_file & params_file
-         const RuntimeOption& option = RuntimeOption(),
+  YOLOv5(const std::string& model_file, const std::string& params_file = "",
-         const Frontend& model_format = Frontend::ONNX);
-
-  // 在X2Paddle转成Paddle后，支持Paddle格式模型的输入
-  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,8 +48,8 @@ class FASTDEPLOY_DECL YOLOv5 : public FastDeployModel {
   // im_info 为预处理记录的信息，后处理用于还原box
   // conf_threshold 后处理时过滤box的置信度阈值
   // nms_iou_threshold 后处理时NMS设定的iou阈值
-  virtual bool
+  virtual bool Postprocess(
-  Postprocess(FDTensor& infer_result, DetectionResult* result,
+      FDTensor& infer_result, DetectionResult* result,
       const std::map<std::string, std::array<float, 2>>& im_info,
       float conf_threshold, float nms_iou_threshold);
 

requirements.txt

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