diff --git a/fastdeploy/vision/evaluation/classify.py b/fastdeploy/vision/evaluation/classify.py index e9cc421fe..66ab47bbd 100644 --- a/fastdeploy/vision/evaluation/classify.py +++ b/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 diff --git a/fastdeploy/vision/ultralytics/__init__.py b/fastdeploy/vision/ultralytics/__init__.py index f7ea91f60..9c3ed915c 100644 --- a/fastdeploy/vision/ultralytics/__init__.py +++ b/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." diff --git a/fastdeploy/vision/ultralytics/ultralytics_pybind.cc b/fastdeploy/vision/ultralytics/ultralytics_pybind.cc index abca7e488..bfb81b0c2 100644 --- a/fastdeploy/vision/ultralytics/ultralytics_pybind.cc +++ b/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_( ultralytics_module, "YOLOv5") - .def(pybind11::init()) .def(pybind11::init()) .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 diff --git a/fastdeploy/vision/ultralytics/yolov5.cc b/fastdeploy/vision/ultralytics/yolov5.cc index c62ab05c7..79158dd60 100644 --- a/fastdeploy/vision/ultralytics/yolov5.cc +++ b/fastdeploy/vision/ultralytics/yolov5.cc @@ -41,25 +41,16 @@ void LetterBox(Mat* mat, std::vector size, std::vector 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; diff --git a/fastdeploy/vision/ultralytics/yolov5.h b/fastdeploy/vision/ultralytics/yolov5.h index 711ceee9f..7dd901db8 100644 --- a/fastdeploy/vision/ultralytics/yolov5.h +++ b/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>& im_info, - float conf_threshold, float nms_iou_threshold); + virtual bool Postprocess( + FDTensor& infer_result, DetectionResult* result, + const std::map>& 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 diff --git a/requirements.txt b/requirements.txt index 0dd006bbc..276109fc1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ opencv-python +tqdm