[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_mem和infer。其中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
• 还没想到