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

相关文章

一文详解Python如何开发游戏

《一文详解Python如何开发游戏》Python是一种非常流行的编程语言,也可以用来开发游戏模组,:本文主要介绍Python如何开发游戏的相关资料,文中通过代码介绍的非常详细,需要的朋友可以参考下... 目录一、python简介二、Python 开发 2D 游戏的优劣势优势缺点三、Python 开发 3D

深入理解Mysql OnlineDDL的算法

《深入理解MysqlOnlineDDL的算法》本文主要介绍了讲解MysqlOnlineDDL的算法,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小... 目录一、Online DDL 是什么?二、Online DDL 的三种主要算法2.1COPY(复制法)

基于Python开发Windows自动更新控制工具

《基于Python开发Windows自动更新控制工具》在当今数字化时代,操作系统更新已成为计算机维护的重要组成部分,本文介绍一款基于Python和PyQt5的Windows自动更新控制工具,有需要的可... 目录设计原理与技术实现系统架构概述数学建模工具界面完整代码实现技术深度分析多层级控制理论服务层控制注

Java中的分布式系统开发基于 Zookeeper 与 Dubbo 的应用案例解析

《Java中的分布式系统开发基于Zookeeper与Dubbo的应用案例解析》本文将通过实际案例,带你走进基于Zookeeper与Dubbo的分布式系统开发,本文通过实例代码给大家介绍的非常详... 目录Java 中的分布式系统开发基于 Zookeeper 与 Dubbo 的应用案例一、分布式系统中的挑战二

Java 缓存框架 Caffeine 应用场景解析

《Java缓存框架Caffeine应用场景解析》文章介绍Caffeine作为高性能Java本地缓存框架,基于W-TinyLFU算法,支持异步加载、灵活过期策略、内存安全机制及统计监控,重点解析其... 目录一、Caffeine 简介1. 框架概述1.1 Caffeine的核心优势二、Caffeine 基础2

Linux五种IO模型的使用解读

《Linux五种IO模型的使用解读》文章系统解析了Linux的五种IO模型(阻塞、非阻塞、IO复用、信号驱动、异步),重点区分同步与异步IO的本质差异,强调同步由用户发起,异步由内核触发,通过对比各模... 目录1.IO模型简介2.五种IO模型2.1 IO模型分析方法2.2 阻塞IO2.3 非阻塞IO2.4

docker编写java的jar完整步骤记录

《docker编写java的jar完整步骤记录》在平常的开发工作中,我们经常需要部署项目,开发测试完成后,最关键的一步就是部署,:本文主要介绍docker编写java的jar的相关资料,文中通过代... 目录all-docker/生成Docker打包部署文件配置服务A的Dockerfile (a/Docke

基于Go语言开发一个 IP 归属地查询接口工具

《基于Go语言开发一个IP归属地查询接口工具》在日常开发中,IP地址归属地查询是一个常见需求,本文将带大家使用Go语言快速开发一个IP归属地查询接口服务,有需要的小伙伴可以了解下... 目录功能目标技术栈项目结构核心代码(main.go)使用方法扩展功能总结在日常开发中,IP 地址归属地查询是一个常见需求:

MySQL使用EXISTS检查记录是否存在的详细过程

《MySQL使用EXISTS检查记录是否存在的详细过程》EXISTS是SQL中用于检查子查询是否返回至少一条记录的运算符,它通常用于测试是否存在满足特定条件的记录,从而在主查询中进行相应操作,本文给大... 目录基本语法示例数据库和表结构1. 使用 EXISTS 在 SELECT 语句中2. 使用 EXIS

基于 Cursor 开发 Spring Boot 项目详细攻略

《基于Cursor开发SpringBoot项目详细攻略》Cursor是集成GPT4、Claude3.5等LLM的VSCode类AI编程工具,支持SpringBoot项目开发全流程,涵盖环境配... 目录cursor是什么?基于 Cursor 开发 Spring Boot 项目完整指南1. 环境准备2. 创建