yolo-inference多后端+多任务+多算法+多精度模型 框架开发记录(cpp版)

本文主要是介绍yolo-inference多后端+多任务+多算法+多精度模型 框架开发记录(cpp版),希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

先贴出github地址,欢迎大家批评指正:https://github.com/taifyang/yolo-inference
不知不觉LZ已经快工作两年了,由于之前的工作内容主要和模型部署相关,想着利用闲暇时间写一些推理方面的经验总结,于是有了这个工程。其实本来也是自己写了玩的,不过已经陆续迭代半年多了,期间也通过借签优秀代码吸收了经验,索性总结一下心得~

1.0 初始版本
1.1 支持多精度模型
1.2 支持tensorrt的cuda前后处理
1.3 支持onnxruntime的int8推理
1.4 onnxruntime推理代码采用cpp风格接口
1.5 采用抽象工厂和单例模式重构代码
1.6 增加cmake编译支持
1.7 增加Linux系统编译支持
2.0 增加yolov8检测器支持
2.1 增加cmake条件编译选项和自动化测试脚本
3.0 增加分类和分割算法支持
3.1 重构代码结构和缺陷修复

最初版的接口头文件部分主要如下:

#pragma once#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>...enum Device_Type
{CPU,GPU,
};class YOLOv5
{
public:void infer(const std::string image_path){m_image = cv::imread(image_path);m_result = m_image.clone();pre_process();process();post_process();cv::imwrite("result.jpg", m_result);cv::imshow("result", m_result);cv::waitKey(0);}cv::Mat m_image;cv::Mat m_result;private:virtual void pre_process() = 0;virtual void process() = 0;virtual void post_process() = 0;
};

该接口类的思路很简单,即一个名为YOLOv5的基类,定义了抽象业务接口如前处理pre_process()、模型推理process()和后处理post_process()需要在派生类中进行具体实现。由基类YOLO根据后端推理框架种类派生出五个子类YOLO_Libtorch 、YOLO_ONNXRuntime、YOLO_OpenCV、YOLO_OpenVINO和YOLO_TensorRT。

#pragma once#include "yolov5.h"
#include <torch/script.h>
#include <torch/torch.h>class YOLOv5_Libtorch : public YOLOv5
{	
public:YOLOv5_Libtorch(std::string model_path, Device_Type device_type);~YOLOv5_Libtorch();private:void pre_process();void process();void post_process();torch::DeviceType m_device;torch::jit::script::Module module;std::vector<torch::jit::IValue> m_inputs;torch::jit::IValue m_outputs;
};

调用时初始化传入模型路径和Device_Type,并指定图片路径即可推理,demo如下:

#include "yolov5_libtorch.h"int main(int argc, char* argv[])
{YOLOv5* yolov5 = new YOLOv5_Libtorch("yolov5n_cpu.torchscript", CPU);yolov5->infer("bus.jpg");return 0;
}

后续扩充了支持不同的模型Model_Type,并使用抽象工厂和单例模式时对外接口在被调用时表现形式更统一:

#pragma once#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>...enum Algo_Type
{Libtorch,ONNXRuntime,OpenCV,OpenVINO,TensorRT,
};enum Device_Type
{CPU,GPU,
};enum Model_Type
{FP32,FP16,INT8,
};class YOLOv5
{
public:virtual void init(const std::string model_path, const Device_Type device_type, Model_Type model_type) = 0;void infer(const std::string image_path);virtual void release() {};protected:virtual void pre_process() = 0;virtual void process() = 0;virtual void post_process();cv::Mat m_image;cv::Mat m_result;float* m_outputs_host;
};class  AlgoFactory
{
public:typedef std::unique_ptr<YOLOv5>(*CreateFunction)();static AlgoFactory& instance();void register_algo(const Algo_Type& algo_type, CreateFunction create_function);std::unique_ptr<YOLOv5> create(const Algo_Type& algo_type);private:AlgoFactory();std::map<Algo_Type, CreateFunction> m_algo_registry;
};

AlgoFactory类中m_algo_registry用来储存算法的唯一全局注册表,register_algo接口用来注册算法,create接口用来返回抽象工厂创建的算法,具体实现如下:

AlgoFactory& AlgoFactory::instance()
{static AlgoFactory algo_factory;return algo_factory;
}void AlgoFactory::register_algo(const Algo_Type& algo_type, CreateFunction create_function)
{m_algo_registry[algo_type] = create_function;
}std::unique_ptr<YOLOv5> AlgoFactory::create(const Algo_Type& algo_type)
{assert(("algo type not exists!", m_algo_registry.find(algo_type) != m_algo_registry.end()));return m_algo_registry[algo_type]();
}AlgoFactory::AlgoFactory()
{register_algo(Algo_Type::Libtorch, []() -> std::unique_ptr<YOLOv5> { return std::make_unique<YOLOv5_Libtorch>(); });register_algo(Algo_Type::ONNXRuntime, []() -> std::unique_ptr<YOLOv5> { return std::make_unique<YOLOv5_ONNXRuntime>(); });register_algo(Algo_Type::OpenCV, []() -> std::unique_ptr<YOLOv5> { return std::make_unique<YOLOv5_OpenCV>(); });register_algo(Algo_Type::OpenVINO, []() -> std::unique_ptr<YOLOv5> { return std::make_unique<YOLOv5_OpenVINO>(); });register_algo(Algo_Type::TensorRT, []() -> std::unique_ptr<YOLOv5> { return std::make_unique<YOLOv5_TensorRT>(); });
}

此时调用时需要先创建算法实例,并依次调用init()、infer()和release()接口,demo的表现形式如下:

#include "yolov5.h"int main(int argc, char* argv[])
{std::unique_ptr<YOLOv5> yolov5 = AlgoFactory::instance().create(Algo_Type::Libtorch);yolov5->init("yolov5n_cpu_fp32.torchscript", CPU, FP32);yolov5->infer("test.mp4");yolov5->release();return 0;
}

2.x版本中主要增加了对yolov8检测器的支持,接口头文件除了增加Algo_Type枚举类型用来表示不同算法:

enum Algo_Type
{YOLOv5,YOLOv8,
};

3.x版本主要增加了对分类、分割算法的支持,头文件增加了Task_Type枚举类型,由于维度的扩充将算法注册表替换为二维向量来储存:

#pragma once#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>enum Backend_Type
{Libtorch,ONNXRuntime,OpenCV,OpenVINO,TensorRT,
};enum Task_Type
{Classification,Detection,Segmentation,
};enum Algo_Type
{YOLOv5,YOLOv8,
};enum Device_Type
{CPU,GPU,
};enum Model_Type
{FP32,FP16,INT8,
};class YOLO
{
public:virtual ~YOLO() {};	//不加此句会导致虚拟继承内存泄漏virtual void init(const Algo_Type algo_type, const Device_Type device_type, const Model_Type model_type, const std::string model_path) = 0;void infer(const std::string file_path, char* argv[], bool save_result = true, bool show_result = true);virtual void release() {};protected:virtual void pre_process() = 0;virtual void process() = 0;virtual void post_process() = 0;cv::Mat m_image;cv::Mat m_result;int m_input_width = 640;int m_input_height = 640;int m_input_numel = 1 * 3 * m_input_width * m_input_height;
};class  CreateFactory
{
public:typedef std::unique_ptr<YOLO>(*CreateFunction)();static CreateFactory& instance();void register_class(const Backend_Type& backend_type, const Task_Type& task_type, CreateFunction create_function);std::unique_ptr<YOLO> create(const Backend_Type& backend_type, const Task_Type& task_type);private:CreateFactory();std::vector<std::vector<CreateFunction>> m_create_registry;
};
CreateFactory& CreateFactory::instance()
{static CreateFactory create_factory;return create_factory;
}void CreateFactory::register_class(const Backend_Type& backend_type, const Task_Type& task_type, CreateFunction create_function)
{m_create_registry[backend_type][task_type] = create_function;
}std::unique_ptr<YOLO> CreateFactory::create(const Backend_Type& backend_type, const Task_Type& task_type)
{if (backend_type >= m_create_registry.size()){std::cout << "unsupported backend type!" << std::endl;std::exit(-1);}if (task_type >= m_create_registry[task_type].size()){std::cout << "unsupported task type!" << std::endl;std::exit(-1);}return m_create_registry[backend_type][task_type]();
}CreateFactory::CreateFactory()
{m_create_registry.resize(5, std::vector<CreateFunction>(3));#ifdef _YOLO_LIBTORCHregister_class(Backend_Type::Libtorch, Task_Type::Classification, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_Libtorch_Classification>(); });register_class(Backend_Type::Libtorch, Task_Type::Detection, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_Libtorch_Detection>(); });register_class(Backend_Type::Libtorch, Task_Type::Segmentation, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_Libtorch_Segmentation>(); });
#endif // _YOLO_Libtorch#ifdef _YOLO_ONNXRUNTIMEregister_class(Backend_Type::ONNXRuntime, Task_Type::Classification, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_ONNXRuntime_Classification>(); });register_class(Backend_Type::ONNXRuntime, Task_Type::Detection, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_ONNXRuntime_Detection>(); });register_class(Backend_Type::ONNXRuntime, Task_Type::Segmentation, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_ONNXRuntime_Segmentation>(); });
#endif // _YOLO_ONNXRuntime#ifdef _YOLO_OPENCVregister_class(Backend_Type::OpenCV, Task_Type::Classification, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_OpenCV_Classification>(); });register_class(Backend_Type::OpenCV, Task_Type::Detection, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_OpenCV_Detection>(); });register_class(Backend_Type::OpenCV, Task_Type::Segmentation, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_OpenCV_Segmentation>(); });
#endif // _YOLO_OpenCV#ifdef _YOLO_OPENVINOregister_class(Backend_Type::OpenVINO, Task_Type::Classification, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_OpenVINO_Classification>(); });register_class(Backend_Type::OpenVINO, Task_Type::Detection,[]() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_OpenVINO_Detection>(); });register_class(Backend_Type::OpenVINO, Task_Type::Segmentation, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_OpenVINO_Segmentation>(); });
#endif // _YOLO_OpenVINO#ifdef _YOLO_TENSORRTregister_class(Backend_Type::TensorRT, Task_Type::Classification, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_TensorRT_Classification>(); });register_class(Backend_Type::TensorRT, Task_Type::Detection, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_TensorRT_Detection>(); });register_class(Backend_Type::TensorRT, Task_Type::Segmentation, []() -> std::unique_ptr<YOLO> { return std::make_unique<YOLO_TensorRT_Segmentation>(); });
#endif // _YOLO_TensorRT
}

通过基类YOLO根据任务类型派生出三个子类YOLO_Classification、YOLO_Detection ,并由YOLO_Detection 派生出YOLO_Segmentation:

#pragma once#include "yolo.h"
#include "utils.h"class YOLO_Classification : virtual public YOLO
{
protected:void draw_result(std::string label);int class_num = 1000;
};
#pragma once#include "yolo.h"
#include "utils.h"class YOLO_Detection : virtual public YOLO
{
protected:void LetterBox(cv::Mat& input_image, cv::Mat& output_image, cv::Vec4d& params, cv::Size shape = cv::Size(640, 640), cv::Scalar color = cv::Scalar(114, 114, 114));void nms(std::vector<cv::Rect> & boxes, std::vector<float> & scores, float score_threshold, float nms_threshold, std::vector<int> & indices);void scale_box(cv::Rect& box, cv::Size size);void draw_result(std::string label, cv::Rect box);int class_num = 80;float score_threshold = 0.2;float nms_threshold = 0.5;float confidence_threshold = 0.2;cv::Vec4d m_params;int m_output_numprob;int m_output_numbox;int m_output_numdet;
};
#pragma once#include "yolo_detection.h"//网络输出相关参数
struct OutputSeg
{int id;             //结果类别idfloat confidence;   //结果置信度cv::Rect box;       //矩形框cv::Mat boxMask;    //矩形框内mask,节省内存空间和加快速度
};//掩膜相关参数
struct MaskParams
{int segChannels = 32;int segWidth = 160;int segHeight = 160;int netWidth = 640;int netHeight = 640;float maskThreshold = 0.5;cv::Size srcImgShape;cv::Vec4d params;
};class YOLO_Segmentation : public YOLO_Detection
{
protected:void GetMask(const cv::Mat& maskProposals, const cv::Mat& mask_protos, OutputSeg& output, const MaskParams& maskParams);void draw_result(std::vector<OutputSeg> result);MaskParams m_mask_params;int m_output_numseg;
};

另一方面和之前版本类似,由基类YOLO根据后端推理框架种类派生出五个子类YOLO_Libtorch 、YOLO_ONNXRuntime、YOLO_OpenCV、YOLO_OpenVINO和YOLO_TensorRT。最终的具体实现子类需要派生自任务种类的父类和推理框架种类的父类,如下所示:

class YOLO_Libtorch : virtual public YOLO
{	
public:void init(const Algo_Type algo_type, const Device_Type device_type, const Model_Type model_type, const std::string model_path);protected:Algo_Type m_algo;torch::DeviceType m_device;Model_Type m_model;torch::jit::script::Module module;std::vector<torch::jit::IValue> m_input;torch::jit::IValue m_output;
};class YOLO_Libtorch_Classification : public YOLO_Libtorch, public YOLO_Classification
{
public:void init(const Algo_Type algo_type, const Device_Type device_type, const Model_Type model_type, const std::string model_path);private:void pre_process();void process();void post_process();float* m_output_host;
};

这篇关于yolo-inference多后端+多任务+多算法+多精度模型 框架开发记录(cpp版)的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/1056836

相关文章

国内环境搭建私有知识问答库踩坑记录(ollama+deepseek+ragflow)

《国内环境搭建私有知识问答库踩坑记录(ollama+deepseek+ragflow)》本文给大家利用deepseek模型搭建私有知识问答库的详细步骤和遇到的问题及解决办法,感兴趣的朋友一起看看吧... 目录1. 第1步大家在安装完ollama后,需要到系统环境变量中添加两个变量2. 第3步 “在cmd中

如何通过Golang的container/list实现LRU缓存算法

《如何通过Golang的container/list实现LRU缓存算法》文章介绍了Go语言中container/list包实现的双向链表,并探讨了如何使用链表实现LRU缓存,LRU缓存通过维护一个双向... 目录力扣:146. LRU 缓存主要结构 List 和 Element常用方法1. 初始化链表2.

Python结合Flask框架构建一个简易的远程控制系统

《Python结合Flask框架构建一个简易的远程控制系统》这篇文章主要为大家详细介绍了如何使用Python与Flask框架构建一个简易的远程控制系统,能够远程执行操作命令(如关机、重启、锁屏等),还... 目录1.概述2.功能使用系统命令执行实时屏幕监控3. BUG修复过程1. Authorization

基于Python开发PDF转Doc格式小程序

《基于Python开发PDF转Doc格式小程序》这篇文章主要为大家详细介绍了如何基于Python开发PDF转Doc格式小程序,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下... 用python实现PDF转Doc格式小程序以下是一个使用Python实现PDF转DOC格式的GUI程序,采用T

使用Python开发一个图像标注与OCR识别工具

《使用Python开发一个图像标注与OCR识别工具》:本文主要介绍一个使用Python开发的工具,允许用户在图像上进行矩形标注,使用OCR对标注区域进行文本识别,并将结果保存为Excel文件,感兴... 目录项目简介1. 图像加载与显示2. 矩形标注3. OCR识别4. 标注的保存与加载5. 裁剪与重置图像

Python中如何控制小数点精度与对齐方式

《Python中如何控制小数点精度与对齐方式》在Python编程中,数据输出格式化是一个常见的需求,尤其是在涉及到小数点精度和对齐方式时,下面小编就来为大家介绍一下如何在Python中实现这些功能吧... 目录一、控制小数点精度1. 使用 round() 函数2. 使用字符串格式化二、控制对齐方式1. 使用

SpringBoot集成图片验证码框架easy-captcha的详细过程

《SpringBoot集成图片验证码框架easy-captcha的详细过程》本文介绍了如何将Easy-Captcha框架集成到SpringBoot项目中,实现图片验证码功能,Easy-Captcha是... 目录SpringBoot集成图片验证码框架easy-captcha一、引言二、依赖三、代码1. Ea

C#集成DeepSeek模型实现AI私有化的流程步骤(本地部署与API调用教程)

《C#集成DeepSeek模型实现AI私有化的流程步骤(本地部署与API调用教程)》本文主要介绍了C#集成DeepSeek模型实现AI私有化的方法,包括搭建基础环境,如安装Ollama和下载DeepS... 目录前言搭建基础环境1、安装 Ollama2、下载 DeepSeek R1 模型客户端 ChatBo

Gin框架中的GET和POST表单处理的实现

《Gin框架中的GET和POST表单处理的实现》Gin框架提供了简单而强大的机制来处理GET和POST表单提交的数据,通过c.Query、c.PostForm、c.Bind和c.Request.For... 目录一、GET表单处理二、POST表单处理1. 使用c.PostForm获取表单字段:2. 绑定到结

Spring Retry 实现乐观锁重试实践记录

《SpringRetry实现乐观锁重试实践记录》本文介绍了在秒杀商品SKU表中使用乐观锁和MybatisPlus配置乐观锁的方法,并分析了测试环境和生产环境的隔离级别对乐观锁的影响,通过简单验证,... 目录一、场景分析 二、简单验证 2.1、可重复读 2.2、读已提交 三、最佳实践 3.1、配置重试模板