[RKNN] 4. 基于零拷贝接口封装

系列文章目录

[RKNN] 1. 入门介绍
[RKNN] 2. 模型转换和推理–API介绍&以yolox为例
[RKNN] 3. 零拷贝接口推理
[RKNN] 4. 基于零拷贝接口封装


文章目录

  • 系列文章目录
  • 前言
  • 一、思路
    • 推理基类
    • 后处理类
    • 实际任务类
  • 二、推理
  • 三、改进


前言

前面几篇文章基本把rknpu的api和推理都介绍了,并且也都通过代码测试了。但是前面的一些实践都是一些顺序执行的结构,适合简单测试,但是在机器人上用起来就很不方便。本文通过封装整个推理过程,并且利用C++的特点,让程序方便后续增加其他功能,例如yolo系列其它检测器、分类任务、分割任务等,另外封装也有利于后续机器人的实际部署,提供一个简易接口方便使用。
完整代码在Github。

一、思路

C++的特性还是比较有利于扩展功能和实现简单接口的。
1. 检测基类:首先需要一个基类通过调用rknn的api完成推理模型的构建、推理设置和内存的分配等工作。
2. 后处理类:图像推理任务的预处理一般比较简单,后处理则比较麻烦,同时其复用的程度也比较大,于是封装后处理类。
3. 实际任务类:继承基类,重写推理函数,实现图像预处理,创建后处理类对象,完成完整推理流程。

推理基类

推理基类我设置成一个模板类,这是因为我希望子类需要重写的infer接口可以根据任务不同实现不同的输出。由于子类需要读取一些模型相关信息,成员变量设置为protected
一般来说子类最少需要重写两个函数分别是init_io_tensor_meminfer。其中infer是纯虚函数,子类必须重写推理过程。考虑到不同模型需要的输入输出数据类型可能不同,所以建议根据自己的模型重写init_io_tensor_mem函数,主要用于零拷贝API中设定输入输出数据类型和需要内存的大小。此外子类在构造函数里必须调用Init()进行初始化过程。

template<class OUTPUT>
class RknnInferBase {
public:
    virtual ~RknnInferBase(); // 销毁相关资源
    virtual void Init(const std::string &model_path);	// 初始化 子类构造函数必须调用
    ...
    virtual void init_io_tensor_mem();   	 // 申请输入输出内存 一般都需要重写
    virtual OUTPUT infer(cv::Mat &img) = 0;	 // 推理 必须重写

protected:
    ...
};

后处理类

后处理过程我也进行封装,但是考虑到实际中部分任务可能不需要复杂的后处理过程,所以没有和基类混合。类的结构也比较简单,主要包括初始化和处理过程,在构造函数里实现模型相关数据的获取,后处理过程部分数据的运算等,然后根据需要在process完成处理过程。
下面程序是yolox推理的后处理类的基本结构。

class YoloxPostProcess {

public:
    YoloxPostProcess(int input_size, 
                        float prob_threshold, 
                        float nms_threshold, 
                        std::vector<rknn_tensor_attr> &output_attrs,
                        std::vector<int32_t> &zps, 
                        std::vector<float> &scales);
    ~YoloxPostProcess() = default;
    void process(int8_t *src, std::vector<ObjBox> &results, float img_scale);

private:
    ...
};

实际任务类

本文还是以前几篇文章写的yolox的推理过程为例,首先继承基类,重写推理过程,完成图像预处理和后处理并返回结果。

// 预处理
static cv::Mat static_resize(cv::Mat& img, int input_w, int input_h);
// 模板类设置
using RknnInferBaseObjBox = RknnInferBase<std::vector<ObjBox>>;

class RknnYolox :public RknnInferBaseObjBox{
public:
    RknnYolox(const std::string &model_path, const float nms_threshold, const float conf_threshold);
    virtual ~RknnYolox();
    virtual std::vector<ObjBox> infer(cv::Mat &img) override; // yolox推理

private:
    std::shared_ptr<YoloxPostProcess> postprocess_;
    std::vector<ObjBox> results_;
};
// 创建模型
std::shared_ptr<RknnInferBaseObjBox> create_infer_yolox(const std::string &model_path, const float nms_threshold, const float conf_threshold);

二、推理

经过封装,可以使用很少的代码量完成推理过程,同时相关资源也会自动完成释放,有利于后续进一步开发。
完整代码在Github。

#include "yolox/rknn_yolox.h"

using namespace std;

int main() {
    string model_name = "../../1convert/yolox_relu_nodecode.rknn";
	string image_name = "../img/1.jpg";
	const float nms_threshold      = 0.65;
	const float box_conf_threshold = 0.45;

    cv::Mat img = cv::imread(image_name);
    auto model = create_infer_yolox(model_name, nms_threshold, box_conf_threshold);

    auto res = model->infer(img);
    for (auto a : res) {
		std::cout<<"ans: "<<a.x1<<" "<<a.y1<<" "<<a.x2<<" "<<a.y2 <<" "<<a.score<<" "<<a.category<< std::endl; 
		cv::rectangle(img, cv::Point(a.x1, a.y1), cv::Point(a.x2, a.y2), cv::Scalar(255, 0, 0, 255), 3);
		cv::putText(img, std::to_string(a.category), cv::Point(a.x1, a.y1 + 12), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));
	}
	return 0;
}

三、改进

之后还是会继续改进的,所以仓库的代码也会更新,所以过段时间可能上面的代码也会发生改变,尽量记录吧。未来主要需要实现:

  • 预处理在RGA上完成
  • 异步 并发 多线程等
  • 封装ros pankage
  • 还没想到

你可能感兴趣的:(RKNN,计算机视觉,边缘计算,人工智能)