├── .gitignore
├── 3rdparty
    ├── dlog
    │   ├── include
    │   │   ├── dlog.h
    │   │   └── log_manage.h
    │   ├── log.properties
    │   └── src
    │   │   ├── log.cpp
    │   │   └── log_manage.cpp
    └── rknn
    │   └── include
    │       ├── rknn_api.h
    │       └── rknn_matmul_api.h
├── CMakeLists.txt
├── LICENSE
├── README.md
├── assets
    └── architecture.png
├── rknn_infer
    ├── main.cpp
    ├── plugin_ctrl.cpp
    ├── plugin_ctrl.h
    ├── rknn_infer.cpp
    ├── rknn_infer.h
    ├── rknn_model.cpp
    ├── rknn_model.h
    ├── utils.h
    └── utils_log.h
├── rknn_infer_api
    └── rknn_infer_api.h
├── rknn_plugins
    ├── image_op_utils.h
    ├── mpp_video_decoder.cpp
    ├── mpp_video_decoder.h
    ├── mpp_video_encoder.cpp
    ├── mpp_video_encoder.h
    ├── mpp_video_utils.h
    ├── plugin_common.h
    ├── rknn_mobilenet
    │   └── rknn_mobilenet.cpp
    ├── rknn_plugin_template
    │   └── rknn_plugin_template.cpp
    └── rknn_yolo_v5
    │   ├── drm_func.h
    │   ├── postprocess.cc
    │   ├── postprocess.h
    │   ├── rga_func.h
    │   ├── rknn_yolo_v5.cpp
    │   ├── rknn_yolo_v5_video.cpp
    │   └── src
    │       └── main_video.cc
└── unit_test
    ├── test_mpp_video_decoder.cpp
    └── test_rknn_model.cpp


/.gitignore:
--------------------------------------------------------------------------------
 1 | # Prerequisites
 2 | *.d
 3 | 
 4 | # Compiled Object files
 5 | *.slo
 6 | *.lo
 7 | *.o
 8 | *.obj
 9 | 
10 | # Precompiled Headers
11 | *.gch
12 | *.pch
13 | 
14 | # Compiled Dynamic libraries
15 | *.so
16 | *.dylib
17 | *.dll
18 | 
19 | # Fortran module files
20 | *.mod
21 | *.smod
22 | 
23 | # Compiled Static libraries
24 | *.lai
25 | *.la
26 | *.a
27 | *.lib
28 | 
29 | # Executables
30 | *.exe
31 | *.out
32 | *.app
33 | 
34 | # self define
35 | cmake-build-debug/
36 | .idea/
37 | 3rdparty/opencv/**
38 | cmake-build-release/
39 | 3rdparty/rga/**
40 | 3rdparty/mpp/**
41 | 


--------------------------------------------------------------------------------
/3rdparty/dlog/include/dlog.h:
--------------------------------------------------------------------------------
 1 | #ifndef LOG_LOG_H
 2 | #define LOG_LOG_H
 3 | 
 4 | #include <string>
 5 | 
 6 | #define LOGGER_CONFIG "log.properties"
 7 | 
 8 | // 注册日志模块，返回实例
 9 | #define LOG_MODULE_INIT(module) log_module_init(#module)
10 | 
11 | // 对外提供的一些必要的定义
12 | #define LOG_MAX_BUFFER 409600
13 | 
14 | #define CHECK_VAL(x,handle) if((x)){        \
15 |                            handle;          \
16 |                          }
17 | 
18 | #define CHECK(x,m,handle) if((x) == (m)){   \
19 |                            handle;          \
20 |                          }
21 | 
22 | /* 日志等级定义 */
23 | enum LOG_LEVEL{
24 |     LOG_DEBUG = 0,
25 |     LOG_INFO,
26 |     LOG_WARN,
27 |     LOG_ERROR,
28 |     LOG_FATAL
29 | };
30 | 
31 | /* 初始化logger */
32 | extern void *log_module_init(const char *module_name);
33 | 
34 | /* log内容接口 */
35 | extern void log(void *logger, LOG_LEVEL logLevel, const char *format = "", ... );
36 | 
37 | #endif //LOG_LOG_H
38 | 


--------------------------------------------------------------------------------
/3rdparty/dlog/include/log_manage.h:
--------------------------------------------------------------------------------
 1 | #ifndef LOG_MANAGE_H
 2 | #define LOG_MANAGE_H
 3 | 
 4 | #include <string>
 5 | #include <fstream>
 6 | #include <cstring>
 7 | #include <iostream>
 8 | #include <ctime>
 9 | #include <list>
10 | #include <map>
11 | 
12 | #include "dlog.h"
13 | 
14 | #define DLOG_TEMP_BUFFER_SIZE 4096
15 | 
16 | /* log方式定义 */
17 | enum LOG_TYPE {
18 |     OUTPUT_FILE = 0,
19 |     OUTPUT_SCREEN,
20 |     OUTPUT_PLAIN,
21 |     OUTPUT_NONE
22 | };
23 | 
24 | class Logger {
25 | public:
26 |     Logger(const std::string &logger_name, LOG_LEVEL log_level, LOG_TYPE log_type, std::string &log_filename);
27 |     void log_message(short level, std::string &message);
28 | 
29 | private:
30 |     static std::string get_level_str(short level);
31 |     bool is_greater_than_level(short level);
32 | 
33 |     // log to file
34 |     void log_msg_file(short level, std::string &message);
35 | 
36 |     // log to file
37 |     void log_msg_plain(short level, std::string &message);
38 | 
39 |     // log to screen
40 |     static void log_msg_screen(short level, std::string &message);
41 | 
42 | private:
43 |     LOG_LEVEL log_level;
44 |     LOG_TYPE log_type;
45 |     std::string log_filename;
46 | };
47 | 
48 | #define LoggerMap std::map<std::string, Logger *>
49 | class LoggerCtl {
50 | public:
51 |     static LoggerCtl* instance();
52 |     LoggerCtl();
53 |     ~LoggerCtl();
54 | 
55 | public:
56 |     // 获取配置并初始化一个log实例
57 |     // 注册模块并初始化实例
58 |     void *register_logger(const std::string &module_name);
59 | 
60 | private:
61 |     static LoggerCtl* _instance;
62 | 
63 | private:
64 |     std::string m_dlog_module_path;
65 |     std::string m_dlog_config_path;
66 |     LoggerMap logger_map;
67 | 
68 |     void get_logger_config(const std::string &name, LOG_TYPE &log_type, LOG_LEVEL &log_level, std::string &filename);
69 | };
70 | 
71 | 
72 | #endif //LOG_MANAGE_H
73 | 


--------------------------------------------------------------------------------
/3rdparty/dlog/log.properties:
--------------------------------------------------------------------------------
 1 | # 每个日志模块分为不同的输出方式
 2 | # 日志等级：DEBUG < INFO < WARN < ERROR < FATAL
 3 | # group 模块
 4 | logger.d_group.log_level = DEBUG
 5 | logger.d_group.log_type = SCREEN
 6 | 
 7 | # mysql 模块
 8 | logger.d_sql.log_level = INFO
 9 | logger.d_sql.log_type = SCREEN
10 | 
11 | 
12 | 


--------------------------------------------------------------------------------
/3rdparty/dlog/src/log.cpp:
--------------------------------------------------------------------------------
 1 | #include <cstdarg>
 2 | 
 3 | #include "dlog.h"
 4 | #include "log_manage.h"
 5 | 
 6 | void *log_module_init(const char *module_name) {
 7 |     if(module_name == nullptr){
 8 |         printf("module name is NULL");
 9 |         return nullptr;
10 |     }
11 |     return LoggerCtl::instance()->register_logger(module_name);
12 | }
13 | 
14 | std::string get_format_str(const char *format, va_list args) {
15 |     char buffer[LOG_MAX_BUFFER]{};
16 |     vsnprintf(buffer, LOG_MAX_BUFFER, format, args);
17 |     return buffer;
18 | }
19 | 
20 | /* log内容接口 */
21 | void log(void *logger, LOG_LEVEL logLevel, const char *format, ... ){
22 |     if (logger == nullptr){
23 |         printf("logger is NULL");
24 |         return;
25 |     }
26 |     if(format == nullptr){
27 |         printf("format is NULL");
28 |         return;
29 |     }
30 |     va_list vaList;
31 |     va_start(vaList, format);
32 |     std::string message = get_format_str(format, vaList);
33 |     va_end(vaList);
34 |     auto *logManage = (Logger *)logger;
35 |     logManage->log_message(logLevel, message);
36 | }


--------------------------------------------------------------------------------
/3rdparty/dlog/src/log_manage.cpp:
--------------------------------------------------------------------------------
  1 | #include <fstream>
  2 | #include <cstring>
  3 | #ifdef __linux__
  4 | #include <unistd.h>
  5 | #else
  6 | #include <Windows.h>
  7 | #include <tchar.h>
  8 | #endif
  9 | #include "log_manage.h"
 10 | 
 11 | using namespace std;
 12 | 
 13 | void Logger::log_message(short level, std::string &message) {
 14 |     if(!is_greater_than_level(level)){
 15 |         return;
 16 |     }
 17 | 
 18 |     switch (log_type){
 19 |         case OUTPUT_FILE:
 20 |             log_msg_file(level, message);
 21 |             break;
 22 |         case OUTPUT_SCREEN:
 23 |             log_msg_screen(level, message);
 24 |             break;
 25 |         case OUTPUT_PLAIN:
 26 |             log_msg_plain(level, message);
 27 |             break;
 28 |         case OUTPUT_NONE:
 29 |             break;
 30 |     }
 31 | }
 32 | 
 33 | void Logger::log_msg_file(short level, std::string &message) {
 34 |     std::fstream output_file_fd;
 35 |     output_file_fd.open(log_filename, std::ios_base::app);
 36 | 
 37 |     time_t now = time(nullptr);
 38 |     struct tm *local_time = localtime(&now);
 39 |     char time_buffer[32];
 40 |     strftime(time_buffer, 80, "%Y-%m-%d %H:%M:%S", local_time);
 41 |     output_file_fd << time_buffer << " " << get_level_str(level) << " " << message << std::endl;
 42 |     output_file_fd.close();
 43 | }
 44 | 
 45 | void Logger::log_msg_plain(short level, std::string &message) {
 46 |     std::fstream output_file_fd;
 47 |     output_file_fd.open(log_filename, std::ios_base::app);
 48 |     output_file_fd << message << std::endl;
 49 |     output_file_fd.close();
 50 | }
 51 | 
 52 | void Logger::log_msg_screen(short level, std::string &message) {
 53 |     time_t now = time(nullptr);
 54 |     struct tm *local_time = localtime(&now);
 55 |     char time_buffer[32];
 56 |     strftime(time_buffer, 80, "%Y-%m-%d %H:%M:%S", local_time);
 57 |     std::cout << time_buffer << " " << get_level_str(level) << " " << message << std::endl;
 58 | }
 59 | 
 60 | bool Logger::is_greater_than_level(short level) {
 61 |     return this->log_level <= level;
 62 | }
 63 | 
 64 | std::string Logger::get_level_str(short level) {
 65 |     if(level == LOG_ERROR){
 66 |         return "[ERROR]";
 67 |     }else if(level == LOG_WARN){
 68 |         return "[WARN] ";
 69 |     }else if(level == LOG_INFO){
 70 |         return "[INFO] ";
 71 |     }else if(level == LOG_DEBUG){
 72 |         return "[DEBUG]";
 73 |     }else if(level == LOG_FATAL){
 74 |         return "[FATAL]";
 75 |     }
 76 |     return "[NONE] ";
 77 | }
 78 | 
 79 | Logger::Logger(const string &logger_name, LOG_LEVEL log_level, LOG_TYPE log_type, string &log_filename) {
 80 |     this->log_level = log_level;
 81 |     this->log_type = log_type;
 82 |     if ((log_type == OUTPUT_FILE || log_type == OUTPUT_PLAIN) && log_filename.empty()){
 83 |         this->log_filename = logger_name + "_default.log";
 84 |     }else{
 85 |         this->log_filename = log_filename;
 86 |     }
 87 | }
 88 | 
 89 | 
 90 | LoggerCtl* LoggerCtl::_instance = nullptr;
 91 | LoggerCtl *LoggerCtl::instance() {
 92 |     if(_instance == nullptr){
 93 |         _instance = new LoggerCtl();
 94 |     }
 95 |     return _instance;
 96 | }
 97 | 
 98 | void *LoggerCtl::register_logger(const std::string &module_name) {
 99 |     auto it = logger_map.find(module_name);
100 |     // 如果已经注册过直接返回注册过的
101 |     if(it != logger_map.end()){
102 |         return it->second;
103 |     }
104 | 
105 |     // 根据module_name获取配置并注册
106 |     LOG_TYPE log_type = OUTPUT_SCREEN;
107 |     LOG_LEVEL log_level = LOG_INFO;
108 |     string filename = "default.log";
109 |     get_logger_config(module_name, log_type, log_level, filename);
110 | 
111 |     //printf("name :%s, type:%d level:%d file:%s\n", module_name.c_str(), log_type, log_level, filename.c_str());
112 |     // 生成对象并注册到列表里
113 |     auto *logger = new Logger(module_name, log_level, log_type, filename);
114 |     logger_map.insert(make_pair(module_name, logger));
115 |     return logger;
116 | }
117 | 
118 | LoggerCtl::LoggerCtl() {
119 |     // 临时路径缓存
120 |     char path_buffer[DLOG_TEMP_BUFFER_SIZE];
121 | #ifdef __linux__
122 |     // Linux 获取绝对路径
123 |     char *ret = getcwd(path_buffer, DLOG_TEMP_BUFFER_SIZE);
124 |     if(ret == nullptr){
125 |         printf("getcwd error!!!");
126 |         snprintf(path_buffer, 1024, "%s", "./");
127 |     }
128 | #else
129 |     // Window 获取绝对路径
130 |     ::GetModuleFileName(nullptr, path_buffer, DLOG_TEMP_BUFFER_SIZE);
131 |     (_tcsrchr(path_buffer, '\\'))[1] = 0;
132 | #endif
133 |     m_dlog_module_path = path_buffer;
134 |     m_dlog_module_path += "/";
135 |     m_dlog_config_path = m_dlog_module_path + LOGGER_CONFIG;
136 | }
137 | 
138 | LoggerCtl::~LoggerCtl() {
139 |     // 清理内存占用
140 |     for(auto it : logger_map){
141 |         delete it.second;
142 |         it.second = nullptr;
143 |     }
144 |     logger_map.clear();
145 | }
146 | 
147 | void LoggerCtl::get_logger_config(const std::string &name, LOG_TYPE &log_type, LOG_LEVEL &log_level, std::string &filename) {
148 |     FILE *fp = fopen(m_dlog_config_path.c_str(), "r");
149 |     if (fp == nullptr) {
150 |         // 文件不存在
151 |         printf("config file %s not exist\n", m_dlog_config_path.c_str());
152 |         return;
153 |     }
154 | 
155 |     char read_buffer[DLOG_TEMP_BUFFER_SIZE];
156 |     while(fgets(read_buffer, DLOG_TEMP_BUFFER_SIZE, fp) != nullptr){
157 |         // 清理空格
158 |         int blank_count = 0;
159 |         unsigned int read_len = strlen(read_buffer);
160 |         for(int i = 0; i < read_len; i++){
161 |             if(read_buffer[i] == ' '){
162 |                 blank_count++;
163 |             }
164 |             read_buffer[i - blank_count] = read_buffer[i];
165 |         }
166 |         read_buffer[read_len - blank_count] = '\0';
167 |         // 跳过注释
168 |         if(read_buffer[0] == '#'){
169 |             continue;
170 |         }
171 | 
172 |         // 解析 key 和 value
173 |         char c_key[256], c_value[256];
174 |         sscanf(read_buffer, "%s=%s", c_key, c_value);
175 |         string key = c_key, value = c_value;
176 | 
177 |         // 跳过非当前模块的配置
178 |         if(key.find("." + name + ".") == string::npos){
179 |             continue;
180 |         }
181 | 
182 |         if(key.find(".log_level") != string::npos){
183 |             if(value == "DEBUG") {
184 |                 log_level = LOG_DEBUG;
185 |             } else if(value == "WARN") {
186 |                 log_level = LOG_WARN;
187 |             } else if(value == "INFO") {
188 |                 log_level = LOG_INFO;
189 |             } else if(value == "ERROR") {
190 |                 log_level = LOG_ERROR;
191 |             } else if(value == "FATAL") {
192 |                 log_level = LOG_FATAL;
193 |             } else {
194 |                 // 默认的日志等级
195 |                 log_level = LOG_INFO;
196 |             }
197 |         } else if(key.find(".log_type") != string::npos){
198 |             if(value == "SCREEN") {
199 |                 log_type = OUTPUT_SCREEN;
200 |             } else if(value == "FILE") {
201 |                 log_type = OUTPUT_FILE;
202 |             } else if(value == "PLAIN") {
203 |                 log_type = OUTPUT_PLAIN;
204 |             } else if(value == "NONE") {
205 |                 log_type = OUTPUT_NONE;
206 |             } else {
207 |                 // 默认的输出类型
208 |                 log_type = OUTPUT_SCREEN;
209 |             }
210 |         } else if(key.find(".log_file") != string::npos){
211 |             // 获取输出为文件类型时的位置
212 |             filename += m_dlog_module_path;
213 |             filename += value;
214 |         }
215 |     }
216 |     fclose(fp);
217 | }
218 | 


--------------------------------------------------------------------------------
/3rdparty/rknn/include/rknn_matmul_api.h:
--------------------------------------------------------------------------------
  1 | /****************************************************************************
  2 |  *
  3 |  *    Copyright (c) 2017 - 2018 by Rockchip Corp.  All rights reserved.
  4 |  *
  5 |  *    The material in this file is confidential and contains trade secrets
  6 |  *    of Rockchip Corporation. This is proprietary information owned by
  7 |  *    Rockchip Corporation. No part of this work may be disclosed,
  8 |  *    reproduced, copied, transmitted, or used in any way for any purpose,
  9 |  *    without the express written permission of Rockchip Corporation.
 10 |  *
 11 |  *****************************************************************************/
 12 | 
 13 | #ifndef _RKNN_MATMUL_API_H
 14 | #define _RKNN_MATMUL_API_H
 15 | 
 16 | #ifdef __cplusplus
 17 | extern "C" {
 18 | #endif
 19 | 
 20 | #include "rknn_api.h"
 21 | 
 22 | typedef rknn_context rknn_matmul_ctx;
 23 | 
 24 | typedef struct _rknn_matmul_tensor_attr
 25 | {
 26 |   char name[RKNN_MAX_NAME_LEN];
 27 | 
 28 |   // indicate A(M, K) or B(K, N) or C(M, N)
 29 |   uint32_t n_dims;
 30 |   uint32_t dims[RKNN_MAX_DIMS];
 31 | 
 32 |   // matmul tensor size
 33 |   uint32_t size;
 34 | 
 35 |   // matmul tensor data type
 36 |   // int8 : A, B
 37 |   // int32: C
 38 |   rknn_tensor_type type;
 39 | } rknn_matmul_tensor_attr;
 40 | 
 41 | typedef struct _rknn_matmul_io_attr
 42 | {
 43 |   // indicate A(M, K) or B(K, N) or C(M, N)
 44 |   rknn_matmul_tensor_attr A;
 45 |   rknn_matmul_tensor_attr B;
 46 |   rknn_matmul_tensor_attr C;
 47 | } rknn_matmul_io_attr;
 48 | 
 49 | /*
 50 |   matmul information struct
 51 |  */
 52 | typedef struct rknn_matmul_info_t
 53 | {
 54 |   int32_t M;
 55 |   int32_t K; // limit: must be aligned with 32byte
 56 |   int32_t N; // limit: must be aligned with 16byte
 57 | 
 58 |   // matmul data type, only support int8 currently
 59 |   // int8: int8(A) x int8(B) -> int32(C)
 60 |   rknn_tensor_type type;
 61 | 
 62 |   // matmul native layout for B
 63 |   // 0: normal layout
 64 |   // 1: native layout
 65 |   int32_t native_layout;
 66 | 
 67 |   // matmul perf layout for A and C
 68 |   // 0: normal layout
 69 |   // 1: perf layout
 70 |   int32_t perf_layout;
 71 | } rknn_matmul_info;
 72 | 
 73 | /*  rknn_matmul_create
 74 | 
 75 |     params:
 76 |         rknn_matmul_ctx *rk_model_ctx           the handle of context.
 77 |         rknn_matmul_info *info         the matmal information.
 78 |         rknn_matmul_io_attr *io_attr   inputs/output attribute
 79 |     return:
 80 |         int                         error code
 81 | */
 82 | int rknn_matmul_create(rknn_matmul_ctx* ctx, rknn_matmul_info* info, rknn_matmul_io_attr* io_attr);
 83 | 
 84 | /* rknn_matmul_set_io_mem
 85 | 
 86 |     params:
 87 |         rknn_matmul_ctx rk_model_ctx            the handle of context.
 88 |         rknn_tensor_mem *mem           the pointer of tensor memory information.
 89 |         rknn_matmul_tensor_attr *attr  the attribute of input or output tensor buffer.
 90 |     return:
 91 |         int                         error code.
 92 | 
 93 |     formula:
 94 |       C = A * B,
 95 | 
 96 |     limit:
 97 |       K <= 4096
 98 |       K align with 32bytes
 99 |       N align with 16bytes
100 | 
101 |     A shape: M x K
102 |       normal layout: (M, K)
103 |               [M1K1, M1K2, ..., M1Kk,
104 |                M2K1, M2K2, ..., M2Kk,
105 |                ...
106 |                MmK1, MmK2, ..., MmKk]
107 |       perf layout: (K / 8, M, 8)
108 |               [K1M1, K2M1,  ..., K8M1,
109 |                K9M2, K10M2, ..., K16M2,
110 |                ...
111 |                K(k-7)Mm, K(k-6)Mm, ..., KkMm]
112 |     B shape: K x N
113 |       normal layout: (K, N)
114 |               [K1N1, K1N2, ..., K1Nn,
115 |                K2N1, K2N2, ..., K2Nn,
116 |                ...
117 |                KkN1, KkN2, ..., KkNn]
118 |       native layout: (N / 16, K / 32, 16, 32)
119 |               [K1N1,  K2N1,  ..., K32N1,
120 |                K1N2,  K2N2,  ..., K32N2,
121 |                ...
122 |                K1N16, K2N16, ..., K32N16,
123 |                K33N1, K34N1, ..., K64N1,
124 |                K33N2, K34N2, ..., K64N2,
125 |                ...
126 |                K(k-31)N16, K(k-30)N16, ..., KkN16,
127 |                K1N17, K2N17, ..., K32N17,
128 |                K1N18, K2N18, ..., K32N18,
129 |                ...
130 |                K(k-31)Nn, K(k-30)Nn, ..., KkNn]
131 |     C shape: M x N
132 |       normal layout: (M, N)
133 |               [M1N1, M1N2, ..., M1Nn,
134 |                M2N1, M2N2, ..., M2Nn,
135 |                ...
136 |                MmN1, MmN2, ..., MmNn]
137 |       perf layout: (N / 4, M, 4)
138 |               [N1M1, N2M1, ..., N4M1,
139 |                N5M2, N6M2, ..., N8M2,
140 |                ...
141 |                N(n-3)Mm, N(n-2)Mm, ..., NnMm]
142 |  */
143 | int rknn_matmul_set_io_mem(rknn_matmul_ctx ctx, rknn_tensor_mem* mem, rknn_matmul_tensor_attr* attr);
144 | 
145 | /*  rknn_matmul_run
146 | 
147 |     run the matmul in blocking mode
148 | 
149 |     params:
150 |         rknn_matmul_ctx rk_model_ctx         the handle of context.
151 |     return:
152 |         int                         error code.
153 |  */
154 | int rknn_matmul_run(rknn_matmul_ctx ctx);
155 | 
156 | /*  rknn_matmul_destroy
157 | 
158 |     destroy the matmul context
159 | 
160 |     params:
161 |         rknn_matmul_ctx rk_model_ctx         the handle of context.
162 |     return:
163 |         int                         error code.
164 |  */
165 | int rknn_matmul_destroy(rknn_matmul_ctx ctx);
166 | 
167 | #ifdef __cplusplus
168 | } // extern "C"
169 | #endif
170 | 
171 | #endif // _RKNN_MATMUL_API_H


--------------------------------------------------------------------------------
/CMakeLists.txt:
--------------------------------------------------------------------------------
  1 | # CMake版本
  2 | cmake_minimum_required(VERSION 3.16.3)
  3 | # C++版本
  4 | set(CMAKE_CXX_STANDARD 17)
  5 | # 指定运行时加载路径位置
  6 | SET(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
  7 | SET(CMAKE_CXX_FLAGS "-Wl,-rpath=./:./lib")
  8 | SET(CMAKE_CXX_FLAGS "-Wl,-E")
  9 | set(CMAKE_INSTALL_RPATH "./:./lib")
 10 | # 选项-本地测试（不依赖 rknn，空转）
 11 | SET(ENABLE_LOCAL_TEST FALSE)
 12 | # 选项-性能统计
 13 | SET(ENABLE_PERFORMANCE_STATISTIC TRUE)
 14 | 
 15 | # RKNN 推理系统
 16 | project(rknn_infer)
 17 | 
 18 | # 第三方工程
 19 | ## RKNN 推理库
 20 | SET(RKNN_DIR ${CMAKE_SOURCE_DIR}/3rdparty/rknn/)
 21 | link_directories(${RKNN_DIR}/lib)
 22 | include_directories(${RKNN_DIR}/include)
 23 | 
 24 | ## opencv
 25 | if (CMAKE_SYSTEM_NAME STREQUAL "Android")
 26 |     set(OpenCV_DIR ${CMAKE_SOURCE_DIR}/../3rdparty/opencv/OpenCV-android-sdk/sdk/native/jni/abi-${CMAKE_ANDROID_ARCH_ABI})
 27 | else()
 28 |   if(LIB_ARCH STREQUAL "armhf")
 29 |     set(OpenCV_DIR ${CMAKE_SOURCE_DIR}/../3rdparty/opencv/opencv-linux-armhf/share/OpenCV)
 30 |   else()
 31 |     set(OpenCV_DIR ${CMAKE_SOURCE_DIR}/../3rdparty/opencv/opencv-linux-aarch64/share/OpenCV)
 32 |   endif()
 33 | endif()
 34 | find_package(OpenCV REQUIRED)
 35 | 
 36 | ## rga
 37 | SET(RGA_DIR ${CMAKE_SOURCE_DIR}/3rdparty/rga/RK356X/)
 38 | include_directories(${RGA_DIR}/include)
 39 | link_directories(${RGA_DIR}/lib/Linux/aarch64)
 40 | 
 41 | ## mpp
 42 | SET(MPP_DIR ${CMAKE_SOURCE_DIR}/3rdparty/mpp/)
 43 | include_directories(${MPP_DIR}/include/rockchip)
 44 | link_directories(${MPP_DIR}/Linux/aarch64)
 45 | 
 46 | ## 日志记录功能
 47 | SET(DLOG_DIR ${CMAKE_SOURCE_DIR}/3rdparty/dlog/)
 48 | include_directories(${DLOG_DIR}/include)
 49 | SET(DLOG_SRC
 50 |     ${DLOG_DIR}/src/log.cpp
 51 |     ${DLOG_DIR}/src/log_manage.cpp
 52 | )
 53 | 
 54 | # 系统内部统计
 55 | if (${ENABLE_PERFORMANCE_STATISTIC})
 56 |     ADD_DEFINITIONS(-DPERFORMANCE_STATISTIC)
 57 | endif ()
 58 | 
 59 | # 工程包含目录设置
 60 | include_directories(
 61 |     ${CMAKE_SOURCE_DIR}/rknn_infer
 62 |     ${CMAKE_SOURCE_DIR}/rknn_infer_api
 63 |     ${CMAKE_SOURCE_DIR}/rknn_plugins
 64 | )
 65 | 
 66 | #link_libraries(${RKNN_DIR}/lib/librknn_api.so)
 67 | #link_libraries(${RKNN_DIR}/lib/librknnrt.so)
 68 | #link_libraries(pthread dl)
 69 | 
 70 | # 推理系统
 71 | add_executable(rknn_infer
 72 |     ${CMAKE_SOURCE_DIR}/rknn_infer/main.cpp
 73 |     ${CMAKE_SOURCE_DIR}/rknn_infer/rknn_infer.cpp
 74 |     ${CMAKE_SOURCE_DIR}/rknn_infer/rknn_model.cpp
 75 |     ${CMAKE_SOURCE_DIR}/rknn_infer/plugin_ctrl.cpp
 76 |     ${DLOG_SRC}
 77 | )
 78 | target_link_libraries(rknn_infer
 79 |     ${RKNN_DIR}/lib/librknn_api.so
 80 |     ${RKNN_DIR}/lib/librknnrt.so
 81 |     pthread
 82 |     dl
 83 | )
 84 | 
 85 | # 测试
 86 | project(test_rknn_model)
 87 | add_executable(test_rknn_model
 88 |         ${CMAKE_SOURCE_DIR}/unit_test/test_rknn_model.cpp
 89 |         ${CMAKE_SOURCE_DIR}/rknn_infer/rknn_model.cpp
 90 |         ${DLOG_SRC}
 91 |         )
 92 | target_link_libraries(test_rknn_model
 93 |         ${OpenCV_LIBS}
 94 |         ${RKNN_DIR}/lib/librknn_api.so
 95 |         ${RKNN_DIR}/lib/librknnrt.so
 96 |         )
 97 | 
 98 | project(test_mpp_video_decoder)
 99 | add_executable(test_mpp_video_decoder
100 |         ${CMAKE_SOURCE_DIR}/unit_test/test_mpp_video_decoder.cpp
101 |         ${CMAKE_SOURCE_DIR}/rknn_plugins/mpp_video_decoder.cpp
102 |         ${CMAKE_SOURCE_DIR}/rknn_plugins/mpp_video_encoder.cpp
103 |         ${DLOG_SRC}
104 |         )
105 | target_link_libraries(test_mpp_video_decoder
106 |         ${OpenCV_LIBS}
107 |         ${RKNN_DIR}/lib/librknn_api.so
108 |         ${RKNN_DIR}/lib/librknnrt.so
109 |         ${MPP_DIR}/Linux/aarch64/librockchip_mpp.so.0
110 |         tbb
111 |         )
112 | 
113 | # 图像图例插件示例
114 | ## rknn_plugin_template
115 | include_directories(${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_plugin_template/)
116 | add_library(rknn_plugin_template SHARED
117 |         ${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_plugin_template/rknn_plugin_template.cpp
118 |         ${DLOG_SRC}
119 | )
120 | target_link_libraries(rknn_plugin_template
121 |         ${RKNN_DIR}/lib/librknn_api.so
122 |         ${RKNN_DIR}/lib/librknnrt.so
123 | )
124 | ## rknn_mobilenet_demo
125 | # https://github.com/rockchip-linux/rknpu2/tree/master/examples/rknn_mobilenet_demo
126 | # include_directories(${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_mobilenet/)
127 | add_library(rknn_mobilenet SHARED
128 |     ${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_mobilenet/rknn_mobilenet.cpp
129 |     ${DLOG_SRC}
130 | )
131 | target_link_libraries(rknn_mobilenet
132 |     ${OpenCV_LIBS}
133 |     ${RKNN_DIR}/lib/librknn_api.so
134 |     ${RKNN_DIR}/lib/librknnrt.so
135 | )
136 | ## rknn_yolov5_demo
137 | # https://github.com/rockchip-linux/rknpu2/tree/master/examples/rknn_yolov5_demo
138 | include_directories(${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_yolo_v5/)
139 | add_library(rknn_yolo_v5 SHARED
140 |         ${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_yolo_v5/rknn_yolo_v5.cpp
141 |         ${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_yolo_v5/postprocess.cc
142 |         ${DLOG_SRC}
143 |         )
144 | target_link_libraries(rknn_yolo_v5
145 |         ${OpenCV_LIBS}
146 |         ${RKNN_DIR}/lib/librknn_api.so
147 |         ${RKNN_DIR}/lib/librknnrt.so
148 |         ${RGA_DIR}/lib/Linux/aarch64/librga.so
149 |         )
150 | add_library(rknn_yolo_v5_video SHARED
151 |         ${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_yolo_v5/rknn_yolo_v5_video.cpp
152 |         ${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_yolo_v5/postprocess.cc
153 |         ${CMAKE_SOURCE_DIR}/rknn_plugins/mpp_video_decoder.cpp
154 |         ${DLOG_SRC}
155 |         )
156 | target_link_libraries(rknn_yolo_v5_video
157 |         ${OpenCV_LIBS}
158 |         ${RKNN_DIR}/lib/librknn_api.so
159 |         ${RKNN_DIR}/lib/librknnrt.so
160 |         ${RGA_DIR}/lib/Linux/aarch64/librga.so
161 |         ${MPP_DIR}/Linux/aarch64/librockchip_mpp.so
162 |         )


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # RknnInferTemplate
  2 | RKNN模型推理部署模板：将数据的输入输出（不同模型适配不同的输入源和做输出结果的处理）做成易于替换的插件的形式，将模型调度部分（通用部分）单独封装起来，减小模型部署工作量。
  3 | 
  4 | # 一、简介
  5 | 
  6 | > 首先需要说明的是采购的时 RK3568 的板子，所以部署模板的设计会以 RK3568 为基准，后续会做其它型号的兼容。
  7 | 
  8 | 为什么使用 NPU 推理部署：板卡上包含了 CPU 算力和 NPU 算力，其中 NPU 算力**仅**可用于加速推理（作为边缘设备，推理应用更广泛一些）。于是为了充分使用板卡上的算力，可以考虑将 CPU 算力用于模型的训练（对于有需要在板卡上训练的场景来说）；将 NPU 算力用于模型的推理应用（一般放在边缘上的模型都会做推理，所以推理场景更广泛），而且将模型训练和模型推理使用的算力分开，也能让推理更顺畅（不会因为突如其来的模型训练占用算力导致不能推理，另外 NPU 的速度相对 CPU 来说更快一些，大约快了 CPU 的 50%（来源于[网友的测试](https://zhuanlan.zhihu.com/p/529861266)））。
  9 | 
 10 | ## 1.1 需求简介
 11 | 
 12 | 有通用模板的需求是因为在 6G 项目中购置了一批板卡，后续将在板卡上做训练和运行模型；并且加了一个 ROCK PI 的开发群，群内大多都对模型的 NPU 推理部署有需求但是对 C++ 和多线程不是很了解，于是有做一个通用模板来部署模型的想法，**将 NPU 推理调度流程与输入输出处理流程分离，降低在推理调度上的工作量**。
 13 | 
 14 | ROCK PI 上使用 NPU 推理的模型需要做模型转换，转换后的模型也都基本是一致的，于是调用模型推理的方式也基本是一样的，所以可以将==调度的部分==封装起来做成一个模板。在经过在开发群的一段实际窥屏后，整理了对模板的以下**要求**：
 15 | 
 16 | * 模型推理调度功能：能够同时支持多个输入源，和多线程推理（这也是比较有挑战性的一点，需要考虑线程之间的协同工作）
 17 | * 模板对不同需求的适应性：不同的模型有不同的数据源，包括文本数据，图像数据，视频数据（实时的/非实时的），甚至语音数据，所以需要考虑不同类型输入的接入，已经异类输入的可扩展性。
 18 | 
 19 | ## 1.2 需求分析
 20 | 
 21 | 从目标上来看，推理模板的**要点在于推理调度与输入输出流程分离**，将输入输出流程剥离出去之后就剩下了模型调度问题。
 22 | 
 23 | 以之前的经验来看，推理调度与输入输出分离的设计可以参考 [FIO 软件中的插件式设计](http://www.delta1037.cn/2021/C_C++/C-C++%E6%8F%92%E4%BB%B6%E5%BC%8F%E8%AE%BE%E8%AE%A1/)，将输入输出的部分做成插件，推理调度部分做成可以加载插件的主程序，这样就能实现同一套推理调度，但是可以有不同的数据源和输出处理。
 24 | 
 25 | > FIO 软件中的插件式设计我已经调研了好久了，终于有了用武之地。我很中意这个设计，本来这个项目可以给学弟来练手，但是我等不及就开始做了。
 26 | 
 27 | 推理调度部分是多线程之间的协同，需要封装好协同处理的队列（主要是数据结构的定义，使其具有可扩展性）。队列处理线程部分也可以使用线程池来做，但是我想这样队列产生线程是零散的，处理线程单独封装有点不雅观，索性可以散装放在推理调度部分。
 28 | 
 29 | # 二、设计
 30 | 
 31 | 以下主要对架构设计和细节设计做详细说明，而设计的主要原则依赖于具体的需求。具体的需求我参考了 RK 官方给出的一些 RKNPU 使用样例，包括一个比较简单的 [rknn_mobilenet_demo](https://github.com/rockchip-linux/rknpu2/tree/master/examples/rknn_mobilenet_demo) 和 [rknn_yolov5_demo](https://github.com/rockchip-linux/rknpu2/tree/master/examples/rknn_yolov5_demo)；以及官方的 SDK 使用说明，使用说明里包含了 RKNN 模型的推理流程。
 32 | 
 33 | ## 2.1 架构
 34 | 
 35 | 架构设计图在：[RKNN 部署模板设计-架构图](siyuan://plugins/siyuan-plugin-excalidrawexcalidraw?icon=iconExcalidraw&title=RKNN%E9%83%A8%E7%BD%B2%E6%A8%A1%E6%9D%BF%E8%AE%BE%E8%AE%A1-%E6%9E%B6%E6%9E%84%E5%9B%BE&data=%7B%22name%22:%22RKNN%E9%83%A8%E7%BD%B2%E6%A8%A1%E6%9D%BF%E8%AE%BE%E8%AE%A1-%E6%9E%B6%E6%9E%84%E5%9B%BE%22%7D)
 36 | 
 37 | 整体的架构图如下，主要划分为调度程序和插件程序两个部分。其中插件程序主要完成数据的输入处理和推理结果的输出处理；调度程序包括了核心的推理调度模块，对推理流程进一步封装的模型管理模块，最后是一个插件管理模块做插件的加载和校验。调度程序在插件加载之后通过插件接口与插件程序通信。
 38 | 
 39 | ![image](assets/architecture.png)
 40 | 
 41 | ## 2.2 细节
 42 | 
 43 | ### 插件接口
 44 | 
 45 | 插件接口部分是整个程序的核心部分，这部分关系到插件程序如何被调度程序调用。首先对调度程序如何使用插件做简要分析，插件中最重要的部分就是获取输入数据和设置推理的输出结果；另外还有插件程序向调度程序设置调度配置（这部分主要是调度程序的线程个数，以及输出结果类型的定制），调度程序向插件发送插件运行配置（这部分主要是加载的模型参数）；最后需要包含插件初始化部分和反初始化做模型内部针对每个线程的配置管理。根据 FIO 中的插件式设计，插件接口定义一个函数结构体来做插件中函数的定义。结构体定义如下：
 46 | 
 47 | ```cpp
 48 | struct PluginStruct {
 49 |     // 插件名称
 50 |     const char *plugin_name;
 51 |     // 插件版本
 52 |     int plugin_version;
 53 | 
 54 |     // 从插件中获取调度配置
 55 |     int (*get_config)(PluginConfigGet *plugin_config);
 56 | 
 57 |     // 给插件设置运行配置
 58 |     int (*set_config)(PluginConfigSet *plugin_config);
 59 | 
 60 |     // 插件初始化
 61 |     int (*init)(struct ThreadData *);
 62 |     // 插件反初始化
 63 |     int (*uninit)(struct ThreadData *);
 64 | 
 65 |     // 收集输入数据和释放输入资源
 66 |     int (*rknn_input)(struct ThreadData *, struct InputUnit *);
 67 |     int (*rknn_input_release)(struct ThreadData *, struct InputUnit *);
 68 | 
 69 |     int (*rknn_output)(struct ThreadData *, struct OutputUnit *);
 70 | };
 71 | ```
 72 | 
 73 | `struct ThreadData` 是每个调度线程独有的数据，上面所有包含该结构的函数都会被数据输入线程和结果输出线程调用；未包含该结构体的函数是全局类型的，只会调用一次。
 74 | 
 75 | ### 插件程序
 76 | 
 77 | 在插件程序部分主要介绍如何使用插件接口来完成一个完整插件程序的撰写（只涉及到接口的功能方面）。
 78 | 
 79 | #### 名称和版本
 80 | 
 81 | 名称和版本是为了做插件的规范化管理。
 82 | 
 83 | #### 配置设置与获取
 84 | 
 85 | 配置是双向的，插件程序需要给调度程序提供必要的参数来启动调度（调度线程的个数），调度程序也必须给插件提供必要的模型信息来让插件给出合适的输入格式。`get_config` 是给调度程序设置配置，`set_config` 是给插件程序设置配置，其中 `get_config` 先于 `set_config` 调用。
 86 | 
 87 | #### 初始化与反初始化
 88 | 
 89 | 初始化 `init` 和反初始化 `uninit` 是为了让插件给每一个线程做私有域的初始化。例如在多个摄像头捕获输入的情况下，需要给每一个线程设置捕获的来源；在多个模型同时被调度的情况下，需要给每一个输出线程设置输出的处理（做统一处理也可以，这里举个不恰当例子）。
 90 | 
 91 | #### 数据输入和结果输出
 92 | 
 93 | 数据输入 `rknn_input` 是调度程序获取推理数据源，可以根据线程数据中的私有化数据（在初始化中设置）来获取不同的输入源；数据输入释放 `rknn_input_release` 是在推理结束后释放掉 `rknn_input` 中申请的动态内存。结果输出 `rknn_output` 是调度程序返回推理的结果，有进一步的处理可以放在这里处理。
 94 | 
 95 | ### 插件管理
 96 | 
 97 | 插件管理部分通过 `dlopen` 和 `dlsym` 来注册插件，并使用 STL 中的 `map` 来做插件的管理。插件注册成功之后，插件管理部分会检测插件中各个函数的可用性（一些非必要函数可以不给出定义）。
 98 | 
 99 | ### 模型管理
100 | 
101 | 模型管理部分是对 RKNN 模型推理的流程做了简要的封装（参考了 RKNN SDK 文档和示例模型）。
102 | 
103 | ### 推理调度
104 | 
105 | 推理调度部分主要的部分是数据获取线程、模型推理线程和这两类线程间的数据队列缓存。工作模式是数据获取线程调用插件的数据获取接口获取模型的数据，将获取的数据放入到任务队列中；推理线程从队列中获取需要处理的数据，送入到模型管理部分得到推理的结果，并调用插件的结果输出接口返回推理结果。
106 | 
107 | # 三、使用
108 | 
109 | 使用此模板做新模型推理时，仅需编写针对新模型的插件，也就是实现插件中的各个接口；另外需要修改 `CMakeList` 使插件能够编译出来。
110 | 
111 | ## 3.1 插件模板
112 | 
113 | ### 插件功能汇聚
114 | 
115 | 以下对插件模板 `rknn_plugin_template` 中的各个位置做出解释，首先是插件模板文件的后半部分：
116 | 
117 | ```cpp
118 | // 注册所有本插件的相关函数到插件结构体中
119 | // 注意：插件名称和结构体定义名称必须和插件动态库名称一致
120 | static struct PluginStruct rknn_plugin_name = {
121 |         .plugin_name 		= "rknn_plugin_name",
122 |         .plugin_version 	= 1,
123 |         .get_config         = get_config,
124 |         .set_config         = set_config,
125 |         .init				= rknn_plugin_init,
126 |         .uninit 			= rknn_plugin_uninit,
127 |         .rknn_input 		= rknn_plugin_input,
128 |         .rknn_input_release = rknn_plugin_input_release,
129 |         .rknn_output		= rknn_plugin_output,
130 | };
131 | 
132 | // 插件动态库在加载时会自动调用该函数
133 | static void plugin_init plugin_auto_register(){
134 |     d_rknn_plugin_info("auto register plugin %p, name: %s", &rknn_plugin_name, rknn_plugin_name.plugin_name)
135 |     plugin_register(&rknn_plugin_name);
136 | }
137 | 
138 | // 插件动态库在关闭时会自动调用该函数
139 | static void plugin_exit plugin_auto_unregister(){
140 |     d_rknn_plugin_info("auto unregister plugin %p, name: %s", &rknn_plugin_name, rknn_plugin_name.plugin_name)
141 |     plugin_unregister(&rknn_plugin_name);
142 | }
143 | 
144 | ```
145 | 
146 | 后半部分是对插件自身的功能做汇聚，并向外部注册。需要注意的有一点：**结构体的名称需要和插件的名称一致**。
147 | 
148 | ### 插件配置处理
149 | 
150 | 插件的配置主要有两个部分，一个是插件向调度程序设置配置；另一个是调度程序向插件设置配置。
151 | 
152 | ```cpp
153 | static int get_config(PluginConfigGet *plugin_config){
154 |     // 输入线程个数
155 |     plugin_config->input_thread_nums = 2;
156 |     // 输出线程个数
157 |     plugin_config->output_thread_nums = 2;
158 |     // 是否需要输出float类型的输出结果
159 |     plugin_config->output_want_float = true;
160 |     return 0;
161 | }
162 | 
163 | PluginConfigSet g_plugin_config_set;
164 | static int set_config(PluginConfigSet *plugin_config){
165 |     // 注意拷贝构造函数
166 |     memcpy(&g_plugin_config_set, plugin_config, sizeof(PluginConfigSet));
167 |     d_rknn_plugin_info("plugin config set success")
168 |     return 0;
169 | }
170 | ```
171 | 
172 | `get_config` 是插件向调度程序设置配置，函数名是调度程序向插件程序获取配置的意思（调度程序是主，插件程序是从）。在这部分插件程序需要向调度程序声明输入线程的个数，推理（输出）线程的个数，以及输出的格式。
173 | 
174 | `set_config` 是调度程序向插件设置配置。在这部分，调度程序将模型相关的信息（在 `rknn_infer_api` 中定义）传递给插件，例如模型版本，输入输出 `tensor ​` 的个数，输入输出 `tensor ​` 特征等等。
175 | 
176 | ### 插件初始化和反初始化
177 | 
178 | 插件初始化时需要对每一个线程做单独的处理，例如对输入线程单独配置独立的输入源（媒体数据等）。
179 | 
180 | ```cpp
181 | static int rknn_plugin_init(struct ThreadData *td) {
182 |     // 设置输入线程的输入源
183 |     return 0;
184 | }
185 | 
186 | static int rknn_plugin_uninit(struct ThreadData *td) {
187 |     // 释放输入线程的输入源
188 |     return 0;
189 | }
190 | ```
191 | 
192 | 每个线程独立的数据源可以存放在自定义的结构体中，并将结构体的地址传给 `ThreadData` 中的 `plugin_private_data` 变量，以便调用其它接口时能够获取到每个线程的私有数据（线程访问变量“全局化”：各个线程私有数据互不干扰，每次调用接口都能拿到该线程私有数据）。注意线程私有数据范围限制在了插件内部，因为访问 `plugin_private_data` 变量需要知道该变量的结构信息，而结构信息仅有插件内部知道。
193 | 
194 | ### 插件输入和输出
195 | 
196 | 插件的输入是插件程序从数据源封装需要推理的数据；插件的输出是插件程序处理推理结果。
197 | 
198 | ```cpp
199 | static int rknn_plugin_input(struct ThreadData *td, struct InputUnit *input_unit) {
200 |     // 根据数据源采集数据，使用 动态 的内存做封装
201 |     return 0;
202 | }
203 | 
204 | static int rknn_plugin_input_release(struct ThreadData *td, struct InputUnit *input_unit) {
205 |     // 释放 rknn_plugin_input 中申请的内存
206 |      return 0;
207 | }
208 | 
209 | static int rknn_plugin_output(struct ThreadData *td, struct OutputUnit *output_unit) {
210 |     // 处理输出数据
211 |     d_rknn_plugin_info("plugin print output data, thread_id: %d", td->thread_id)
212 |     return 0;
213 | }
214 | ```
215 | 
216 | 在多源推理的情况下，输入接口可以从 `ThreadData` 中的 `plugin_sync_data` 获取到输入线程的私有数据源，然后从该私有数据源装配需要推理的数据。
217 | 
218 | ## 3.2 插件编译
219 | 
220 | ### 配置
221 | 
222 | 插件的编译是获取插件动态库文件，需要在 CMakeList 中添加以下内容（以编译 `rknn_plugin_template` 为例）：
223 | 
224 | ```cmake
225 | ## rknn_plugin_template
226 | include_directories(${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_plugin_template/)
227 | add_library(rknn_plugin_template SHARED
228 |         ${CMAKE_SOURCE_DIR}/rknn_plugins/rknn_plugin_template/rknn_plugin_template.cpp
229 |         ${DLOG_SRC}
230 | )
231 | target_link_libraries(rknn_plugin_template
232 |         ${RKNN_DIR}/lib/librknn_api.so
233 |         ${RKNN_DIR}/lib/librknnrt.so
234 | )
235 | ```
236 | 
237 | **其中：**
238 | 
239 | * include_directories：添加头文件包含路径（如果插件不依赖其它文件可以不加）
240 | * add_library：将插件源代码装配成动态库文件，`rknn_plugin_template` 是动态库名称，`SHARED` 表示装配成动态库而不是静态库，后续是动态库中的文件（`${DLOG_SRC}` 是日志相关的代码）
241 | * target_link_libraries：添加插件动态库依赖的其它库文件
242 | 
243 | ### 编译
244 | 
245 | 在根目录下执行如下命令：
246 | 
247 | ```cmake
248 | # 生成编译目录
249 | mkdir build
250 | cd build
251 | 
252 | # 生成Makefile
253 | cmake ..
254 | 
255 | # 编译全部文件,-j8是多路同时编译
256 | make -j 8
257 | # 编译部分文件
258 | make -j 8 project_name
259 | ```
260 | 
261 | 将依赖的库文件和模型文件拷贝到编译目录。
262 | 
263 | ### 运行
264 | 
265 | 在编译目录执行如下命令运行：
266 | 
267 | ```cmake
268 | ./rknn_infer -m <model_path> -p <plugin_name>
269 | ```


--------------------------------------------------------------------------------
/assets/architecture.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/delta1037/RknnInferTemplate/fc86761424f096fc67de7289045ece89bac31836/assets/architecture.png


--------------------------------------------------------------------------------
/rknn_infer/main.cpp:
--------------------------------------------------------------------------------
 1 | /**
 2 |  * @author: bo.liu
 3 |  * @mail: geniusrabbit@qq.com
 4 |  * @date: 2023.08.03
 5 |  * @brief: 推理模板主文件
 6 |  */
 7 | #include <string>
 8 | #include "rknn_infer.h"
 9 | #include "utils_log.h"
10 | 
11 | bool g_system_running;
12 | 
13 | #ifdef __linux__
14 | // 处理linux信号
15 | #include <csignal>
16 | 
17 | void quit_handler(int sig_num){
18 |     if(sig_num == SIGQUIT){
19 |         g_system_running = false;
20 |         d_rknn_infer_warn("system received signal SIGQUIT")
21 |     }else if(sig_num == SIGINT){
22 |         g_system_running = false;
23 |         d_rknn_infer_warn("system received signal SIGINT")
24 |     }else if (sig_num == SIGSTOP){
25 |         g_system_running = false;
26 |         d_rknn_infer_warn("system received signal SIGSTOP")
27 |     }
28 | }
29 | #endif
30 | 
31 | int main(int argc, char *argv[]) {
32 | #ifdef __linux__
33 |     // 注册信号处理函数
34 |     signal(SIGQUIT, quit_handler);
35 | #endif
36 |     // 读取配置
37 |     const std::string usage = "Usage: ./rknn_infer -m <model_path> -p <plugin_name>";
38 |     std::string model_path = "./model/RK3566_RK3568/mobilenet_v1.rknn";
39 |     std::string plugin_name = "rknn_mobilenet";
40 |     for(int idx = 0; idx < argc; idx++){
41 |         std::string args = argv[idx];
42 |         if (args == "-m" || args == "--model"){
43 |             model_path = argv[++idx];
44 |         }
45 |         if (args == "-p" || args == "--plugin"){
46 |             plugin_name = argv[++idx];
47 |         }
48 |     }
49 |     if (model_path.empty()){
50 |         d_rknn_infer_error("model path is empty!")
51 |         d_rknn_plugin_info("%s", usage.c_str())
52 |         return -1;
53 |     }
54 |     d_rknn_infer_info("model path: %s", model_path.c_str())
55 | 
56 |     if (plugin_name.empty()){
57 |         d_rknn_infer_error("plugin name is empty!")
58 |         d_rknn_plugin_info("%s", usage.c_str())
59 |         return -1;
60 |     }
61 |     d_rknn_infer_info("plugin path: %s", plugin_name.c_str())
62 | 
63 |     // 启动推理
64 |     g_system_running = true;
65 |     RknnInfer rknn_infer(model_path, plugin_name);
66 |     if (!rknn_infer.check_init()){
67 |         d_rknn_infer_error("rknn infer init fail!")
68 |         return -1;
69 |     }
70 |     rknn_infer.stop();
71 |     d_rknn_infer_info("rknn infer stop!")
72 | 
73 |     // 输出时间统计
74 | #ifdef PERFORMANCE_STATISTIC
75 |     d_rknn_infer_info("performance statistic:")
76 |     rknn_infer.print_statistic();
77 | #endif
78 |     return 0;
79 | }


--------------------------------------------------------------------------------
/rknn_infer/plugin_ctrl.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.03
  5 |  * @brief: 插件控制
  6 |  */
  7 | #include <map>
  8 | #include <thread>
  9 | #include <mutex>
 10 | #include <cstring>
 11 | #include "plugin_ctrl.h"
 12 | #include "utils_log.h"
 13 | 
 14 | static std::mutex m_plugin_map_lock;
 15 | static std::map<std::string, struct PluginStruct *> g_plugin_map;
 16 | 
 17 | PluginStruct* map_plugin_find(const std::string &plugin_name){
 18 |     // 从 map 中查找
 19 |     std::lock_guard<std::mutex> map_lock(m_plugin_map_lock);
 20 |     auto iter = g_plugin_map.find(plugin_name);
 21 |     if (iter != g_plugin_map.end()){
 22 |         return iter->second;
 23 |     }
 24 |     return nullptr;
 25 | }
 26 | 
 27 | void map_plugin_insert(const std::string &plugin_name, PluginStruct* plugin){
 28 |     // 从 map 中查找
 29 |     std::lock_guard<std::mutex> map_lock(m_plugin_map_lock);
 30 |     auto iter = g_plugin_map.find(plugin_name);
 31 |     if (iter == g_plugin_map.end()){
 32 |         g_plugin_map.insert(std::make_pair(plugin_name, plugin));
 33 |     }
 34 | }
 35 | 
 36 | void map_plugin_remove(const std::string &plugin_name){
 37 |     // 从 map 中查找
 38 |     std::lock_guard<std::mutex> map_lock(m_plugin_map_lock);
 39 |     auto iter = g_plugin_map.find(plugin_name);
 40 |     if (iter != g_plugin_map.end()){
 41 |         iter->second = nullptr;
 42 |         g_plugin_map.erase(plugin_name);
 43 |     }
 44 | }
 45 | 
 46 | PluginStruct* load_plugin(const std::string &plugin_name){
 47 |     struct PluginStruct *plugin;
 48 |     // 从 map 中查找
 49 |     plugin = map_plugin_find(plugin_name);
 50 |     if (plugin){
 51 |         return plugin;
 52 |     }
 53 | 
 54 |     char path_buffer[1024];
 55 | #ifdef __linux__
 56 |     // Linux 获取绝对路径
 57 |     char *ret = getcwd(path_buffer, 1024);
 58 |     if(ret == nullptr){
 59 |         d_rknn_plugin_error("getcwd error!!!");
 60 |         snprintf(path_buffer, 1024, "%s", "./");
 61 |     }
 62 | #endif
 63 |     std::string lib_path = std::string(path_buffer) + "/lib" + plugin_name + ".so";
 64 |     d_rknn_plugin_info("dlopen plugin %s, path:%s", plugin_name.c_str(), lib_path.c_str())
 65 |     void *dll_handle = dlopen(lib_path.c_str(), RTLD_LAZY);
 66 |     if (!dll_handle) {
 67 |         d_rknn_plugin_info("dlopen plugin %s, error: %s", plugin_name.c_str(), dlerror())
 68 |         return nullptr;
 69 |     }
 70 |     // 动态库打开后就可能已经注册过了（自动注册接口），再判断一次
 71 |     plugin = map_plugin_find(plugin_name);
 72 |     if (plugin){
 73 |         return plugin;
 74 |     }
 75 | 
 76 |     // 如果没有注册过，就查找结构体定义位置
 77 |     std::string plugin_symbol = plugin_name;
 78 | #ifdef __cplusplus
 79 |     // C++ 中符号修饰略显复杂，所以需要使用系统命令查询修饰后的符号
 80 |     std::string cmd = "nm " + lib_path + " | grep " + plugin_name + " | grep -i \" d \" | awk '{print $3}'";
 81 |     d_rknn_plugin_info("dlsym plugin %s, cmd:%s", plugin_name.c_str(), cmd.c_str())
 82 |     FILE *fp = popen(cmd.c_str(), "r");
 83 |     if (fp == nullptr){
 84 |         d_rknn_plugin_error("popen error!!!")
 85 |         return nullptr;
 86 |     }
 87 |     char symbol_buffer[1024];
 88 |     memset(symbol_buffer, 0, 1024);
 89 |     if (nullptr == fgets(symbol_buffer, 1024, fp)){
 90 |         d_rknn_plugin_error("fgets error!!!")
 91 |         return nullptr;
 92 |     }
 93 |     pclose(fp);
 94 |     if (strlen(symbol_buffer) == 0){
 95 |         d_rknn_plugin_error("popen error!!!")
 96 |         return nullptr;
 97 |     }
 98 |     // 去掉最后一个换行符
 99 |     symbol_buffer[strlen(symbol_buffer) - 1] = '\0';
100 |     plugin_symbol = std::string(symbol_buffer);
101 | #endif
102 |     plugin = static_cast<PluginStruct *>(dlsym(dll_handle, plugin_symbol.c_str()));
103 |     d_rknn_plugin_info("dlsym plugin %s, symbol:%s, symbol_ptr:%p",
104 |                        plugin_name.c_str(),
105 |                        plugin_symbol.c_str(),
106 |                        plugin)
107 |     // 将查找到的结构体插入到 map 中
108 |     if (plugin) {
109 |         map_plugin_insert(plugin_name, plugin);
110 |     }else{
111 |         plugin = map_plugin_find(plugin_name);
112 |     }
113 |     return plugin;
114 | }
115 | 
116 | void plugin_register(struct PluginStruct *plugin){
117 |     if(!plugin){
118 |         d_rknn_plugin_error("plugin register error, plugin is null")
119 |         return;
120 |     }
121 |     d_rknn_plugin_info("plugin register, plugin_name:%s, plugin_version", plugin->plugin_name, plugin->plugin_version)
122 |     map_plugin_insert(plugin->plugin_name, plugin);
123 | }
124 | 
125 | void plugin_unregister(struct PluginStruct *plugin){
126 |     if(!plugin){
127 |         d_rknn_plugin_error("plugin unregister error, plugin is null")
128 |         return;
129 |     }
130 |     d_rknn_plugin_info("plugin unregister, plugin_name:%s, plugin_version %d", plugin->plugin_name, plugin->plugin_version)
131 |     map_plugin_remove(plugin->plugin_name);
132 | }
133 | 
134 | struct PluginStruct *get_plugin(const std::string &plugin_name){
135 |     PluginStruct *it_find = map_plugin_find(plugin_name);
136 |     if (it_find == nullptr){
137 |         it_find = load_plugin(plugin_name);
138 |     }
139 |     // 判断接口是否是可用的
140 |     if(it_find == nullptr){
141 |         d_rknn_plugin_error("plugin is nullptr! plugin_name=%s", plugin_name.c_str())
142 |         return nullptr;
143 |     }
144 |     if (it_find->plugin_name == nullptr || strlen(it_find->plugin_name) == 0){
145 |         d_rknn_plugin_error("plugin_name is nullptr! plugin_name=%s", plugin_name.c_str())
146 |         return nullptr;
147 |     }
148 | 
149 | //    if (it_find->get_config == nullptr){
150 | //        d_rknn_plugin_error("plugin get_config is nullptr! plugin_name=%s", plugin_name.c_str())
151 | //        return nullptr;
152 | //    }
153 |     if (it_find->set_config == nullptr){
154 |         d_rknn_plugin_error("plugin set_config is nullptr! plugin_name=%s", plugin_name.c_str())
155 |         return nullptr;
156 |     }
157 | 
158 |     if(it_find->init == nullptr){
159 |         d_rknn_plugin_error("plugin init is nullptr! plugin_name=%s", plugin_name.c_str())
160 |         return nullptr;
161 |     }
162 |     if (it_find->uninit == nullptr){
163 |         d_rknn_plugin_error("plugin uninit is nullptr! plugin_name=%s", plugin_name.c_str())
164 |         return nullptr;
165 |     }
166 | 
167 |     if (it_find->rknn_input == nullptr){
168 |         d_rknn_plugin_error("plugin rknn_input is nullptr! plugin_name=%s", plugin_name.c_str())
169 |         return nullptr;
170 |     }
171 |     if (it_find->rknn_input_release == nullptr){
172 |         d_rknn_plugin_error("plugin rknn_input_release is nullptr! plugin_name=%s", plugin_name.c_str())
173 |         return nullptr;
174 |     }
175 | 
176 |     if (it_find->rknn_output == nullptr){
177 |         d_rknn_plugin_error("plugin rknn_infer is nullptr! plugin_name=%s", plugin_name.c_str())
178 |         return nullptr;
179 |     }
180 |     d_rknn_plugin_info("plugin is find! plugin_name=%s, version:%d", plugin_name.c_str(), it_find->plugin_version)
181 |     return it_find;
182 | }


--------------------------------------------------------------------------------
/rknn_infer/plugin_ctrl.h:
--------------------------------------------------------------------------------
 1 | /**
 2 |  * @author: bo.liu
 3 |  * @mail: geniusrabbit@qq.com
 4 |  * @date: 2023.08.03
 5 |  * @brief: 插件控制
 6 |  */
 7 | #ifndef RKNN_INFER_PLUGIN_CTRL_H
 8 | #define RKNN_INFER_PLUGIN_CTRL_H
 9 | #include <string>
10 | #include <dlfcn.h>
11 | #include "rknn_infer_api.h"
12 | 
13 | // 获取动态库接口
14 | struct PluginStruct *get_plugin(const std::string &plugin_name);
15 | 
16 | #endif // RKNN_INFER_PLUGIN_CTRL_H


--------------------------------------------------------------------------------
/rknn_infer/rknn_infer.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.03
  5 |  * @brief: 推理调度实现
  6 |  */
  7 | #include <cstring>
  8 | #include "rknn_infer.h"
  9 | #include "utils_log.h"
 10 | 
 11 | extern bool g_system_running;
 12 | 
 13 | RknnInfer::RknnInfer(const std::string &model_name, const std::string &plugin_name) {
 14 |     // 初始化变量
 15 |     m_init = false;
 16 |     // 加载插件
 17 |     PluginStruct *plugin = get_plugin(plugin_name);
 18 |     if (plugin == nullptr) {
 19 |         d_rknn_infer_error("load plugin failed")
 20 |         return;
 21 |     }
 22 | 
 23 |     // 获取插件配置
 24 |     if (plugin->get_config != nullptr && 0 != plugin->get_config(&m_plugin_get_config)) {
 25 |         d_rknn_infer_error("get_config failed")
 26 |         return;
 27 |     }
 28 |     d_rknn_infer_info("rknn config, input_thread_nums:%d, output_thread_nums:%d",
 29 |                       m_plugin_get_config.input_thread_nums,
 30 |                       m_plugin_get_config.output_thread_nums)
 31 |     d_rknn_infer_info("rknn config, output_want_float:%d",
 32 |                       m_plugin_get_config.output_want_float)
 33 | 
 34 |     // 初始化模型
 35 | #ifdef PERFORMANCE_STATISTIC
 36 |     time_unit t_model_init = get_time_of_ms();
 37 | #endif
 38 |     for(int idx = 0; idx < m_plugin_get_config.output_thread_nums; ++idx){
 39 |         if(idx == 0){
 40 |             m_rknn_models.emplace_back(new RknnModel(model_name, m_plugin_set_config));
 41 |         }
 42 |         else{
 43 |             m_rknn_models.emplace_back(m_rknn_models[0]->model_infer_dup());
 44 |         }
 45 |     }
 46 | #ifdef PERFORMANCE_STATISTIC
 47 |     m_statistic.s_model_init_count = m_plugin_get_config.output_thread_nums;
 48 |     m_statistic.s_model_init_ms = get_time_of_ms() - t_model_init;
 49 | #endif
 50 |     for(int idx = 0; idx < m_plugin_get_config.output_thread_nums; ++idx){
 51 |         if(!m_rknn_models[idx]->check_init()){
 52 |             d_rknn_infer_error("rknn model %d init failed", idx)
 53 |             return;
 54 |         }
 55 |     }
 56 | 
 57 |     // 调度之前给插件传递配置信息
 58 |     if (0 != plugin->set_config(&m_plugin_set_config)) {
 59 |         d_rknn_infer_error("set_config failed")
 60 |         return;
 61 |     }
 62 | 
 63 |     // 启动输出处理线程
 64 |     for (int idx = 0; idx < m_plugin_get_config.output_thread_nums; ++idx) {
 65 |         ThreadData td_data{};
 66 |         td_data.plugin = plugin;
 67 |         td_data.thread_id = idx;
 68 |         td_data.thread_type = THREAD_TYPE_OUTPUT;
 69 |         m_infer_proc_meta.push_back(td_data);
 70 |         m_infer_proc_ctrl.emplace_back([this, idx] { infer_proc_thread(idx); });
 71 |     }
 72 | 
 73 |     // 启动输入接收线程
 74 |     for (int idx = 0; idx < m_plugin_get_config.input_thread_nums; ++idx) {
 75 |         ThreadData td_data{};
 76 |         td_data.plugin = plugin;
 77 |         td_data.thread_id = idx;
 78 |         td_data.thread_type = THREAD_TYPE_INPUT;
 79 |         m_input_data_meta.push_back(td_data);
 80 |         m_input_data_ctrl.emplace_back([this, idx] { input_data_thread(idx); });
 81 |     }
 82 | 
 83 |     m_init = true;
 84 | }
 85 | 
 86 | RetStatus RknnInfer::stop() {
 87 |     // 先数据线程
 88 |     for (auto &item : m_input_data_ctrl) {
 89 |         item.join();
 90 |     }
 91 |     // 停推理线程
 92 |     for (auto &item : m_infer_proc_ctrl) {
 93 |         item.join();
 94 |     }
 95 |     return RET_STATUS_SUCCESS;
 96 | }
 97 | 
 98 | 
 99 | bool RknnInfer::check_init() const {
100 |     return m_init;
101 | }
102 | 
103 | RetStatus RknnInfer::get_input_unit(QueuePack &pack) {
104 |     std::unique_lock<std::mutex> proc_queue_lock(m_infer_queue_mutex);
105 |     while (m_infer_queue.empty()) {
106 |         m_infer_queue_not_empty.wait(proc_queue_lock);  //如果队列为空，线程就在此阻塞挂起，等待唤醒
107 |     }
108 |     if(!m_infer_queue.empty()){
109 |         // 从头部开始拿
110 |         pack = m_infer_queue.front();
111 |         m_infer_queue.pop_front();
112 |         return RetStatus::RET_STATUS_SUCCESS;
113 |     }else{
114 |         return RetStatus::RET_STATUS_FAILED;
115 |     }
116 | }
117 | 
118 | RetStatus RknnInfer::put_input_unit(QueuePack &pack) {
119 |     std::unique_lock<std::mutex> proc_queue_lock(m_infer_queue_mutex);
120 |     m_infer_queue.push_back(pack);
121 |     m_infer_queue_not_empty.notify_one();
122 |     return RetStatus::RET_STATUS_SUCCESS;
123 | }
124 | 
125 | uint32_t RknnInfer::get_queue_size() {
126 |     std::unique_lock<std::mutex> proc_queue_lock(m_infer_queue_mutex);
127 |     return m_infer_queue.size();
128 | }
129 | 
130 | void RknnInfer::input_data_thread(uint32_t idx) {
131 |     auto &td_data = m_input_data_meta[idx];
132 |     // 插件初始化
133 | #ifdef PERFORMANCE_STATISTIC
134 |     time_unit t_plugin_init = get_time_of_ms();
135 | #endif
136 |     if (0 != td_data.plugin->init(&td_data)) {
137 |         d_rknn_infer_error("rknn_infer_init failed")
138 |         return;
139 |     }
140 | #ifdef PERFORMANCE_STATISTIC
141 |     {
142 |         std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_plugin_init_mutex);
143 |         m_statistic.s_plugin_init_count++;
144 |         m_statistic.s_plugin_init_ms += get_time_of_ms() - t_plugin_init;
145 |     }
146 | #endif
147 | 
148 |     while(g_system_running){
149 |         // 限制队列长度，降低任务处理延时
150 |         if(m_plugin_get_config.task_queue_limit !=0){
151 |             uint32_t queue_size = get_queue_size();
152 |             if(queue_size >= uint32_t(m_plugin_get_config.task_queue_limit / 1.5)){
153 |                 sleepUS(100000);
154 |                 continue;
155 |             }else if (queue_size >= m_plugin_get_config.task_queue_limit){
156 |                 sleepUS(500000);
157 |                 continue;
158 |             }else if (queue_size >= uint32_t(m_plugin_get_config.task_queue_limit * 1.5)){
159 |                 sleepUS(100000);
160 |                 continue;
161 |             }
162 |         }
163 |         // 收集数据
164 | #ifdef PERFORMANCE_STATISTIC
165 |         time_unit t_plugin_input_ms = get_time_of_ms();
166 | #endif
167 |         auto *input_unit = new InputUnit();
168 |         if (0 != td_data.plugin->rknn_input(&td_data, input_unit)) {
169 |             d_rknn_infer_error("rknn_infer_get_input_data failed")
170 |             continue;
171 |         }
172 | #ifdef PERFORMANCE_STATISTIC
173 |         {
174 |             std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_plugin_input_mutex);
175 |             m_statistic.s_plugin_input_count++;
176 |             m_statistic.s_plugin_input_ms += get_time_of_ms() - t_plugin_input_ms;
177 |         }
178 | #endif
179 | 
180 |         // 放入队列
181 |         QueuePack pack{};
182 | #ifdef PERFORMANCE_STATISTIC
183 |         pack.s_pack_record_ms = get_time_of_ms();
184 | #endif
185 |         pack.input_unit = input_unit;
186 |         pack.plugin_sync_data = td_data.plugin_sync_data;
187 |         put_input_unit(pack);
188 |     }
189 | 
190 |     // 插件反初始化
191 | #ifdef PERFORMANCE_STATISTIC
192 |     time_unit t_plugin_uninit = get_time_of_ms();
193 | #endif
194 |     if (0 != td_data.plugin->uninit(&td_data)) {
195 |         d_rknn_infer_error("rknn_infer uninit failed")
196 |         return;
197 |     }
198 | #ifdef PERFORMANCE_STATISTIC
199 |     {
200 |         std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_plugin_uninit_mutex);
201 |         m_statistic.s_plugin_uninit_count++;
202 |         m_statistic.s_plugin_uninit_ms += get_time_of_ms() - t_plugin_uninit;
203 |     }
204 | #endif
205 | }
206 | 
207 | void RknnInfer::infer_proc_thread(uint32_t idx) {
208 |     auto &td_data = m_infer_proc_meta[idx];
209 |     // 插件初始化
210 | #ifdef PERFORMANCE_STATISTIC
211 |     time_unit t_plugin_init = get_time_of_ms();
212 | #endif
213 |     if (0 != td_data.plugin->init(&td_data)) {
214 |         d_rknn_infer_error("rknn_infer_init failed")
215 |         return;
216 |     }
217 | #ifdef PERFORMANCE_STATISTIC
218 |     {
219 |         std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_plugin_init_mutex);
220 |         m_statistic.s_plugin_init_count++;
221 |         m_statistic.s_plugin_init_ms += get_time_of_ms() - t_plugin_init;
222 |     }
223 | #endif
224 |     while(g_system_running){
225 |         // 获取数据
226 |         QueuePack pack{};
227 |         RetStatus ret = get_input_unit(pack);
228 |         if (ret != RetStatus::RET_STATUS_SUCCESS){
229 |             d_rknn_infer_error("get_input_unit failed")
230 |             continue;
231 |         }
232 | #ifdef PERFORMANCE_STATISTIC
233 |         {
234 |             std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_queue_mutex);
235 |             m_statistic.s_queue_count++;
236 |             m_statistic.s_queue_ms += get_time_of_ms() - pack.s_pack_record_ms;
237 |         }
238 | #endif
239 | //        auto output_unit = new OutputUnit();
240 |         auto *output_unit = new OutputUnit();
241 |         output_unit->n_outputs = m_plugin_set_config.io_num.n_output;
242 |         output_unit->outputs = (rknn_output*)malloc(output_unit->n_outputs * sizeof(rknn_output));
243 |         memset(output_unit->outputs, 0, output_unit->n_outputs * sizeof(rknn_output));
244 |         for(int i = 0; i < output_unit->n_outputs; i++){
245 |             output_unit->outputs[i].want_float = m_plugin_get_config.output_want_float ? 1 : 0;
246 |         }
247 | 
248 |         // 推理
249 | #ifdef PERFORMANCE_STATISTIC
250 |         time_unit t_model_infer = get_time_of_ms();
251 | #endif
252 |         ret = m_rknn_models[idx]->model_infer_sync(
253 |                 pack.input_unit->n_inputs,
254 |                 pack.input_unit->inputs,
255 |                 output_unit->n_outputs,
256 |                 output_unit->outputs);
257 |         if (ret != RetStatus::RET_STATUS_SUCCESS){
258 |             d_rknn_infer_error("model_infer_sync failed")
259 |             continue;
260 |         }
261 | #ifdef PERFORMANCE_STATISTIC
262 |         {
263 |             std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_model_infer_mutex);
264 |             m_statistic.s_model_infer_count++;
265 |             m_statistic.s_model_infer_ms += get_time_of_ms() - t_model_infer;
266 |         }
267 | #endif
268 | 
269 |         // 输出结果
270 |         // 转移同步数据
271 |         td_data.plugin_sync_data = pack.plugin_sync_data;
272 | #ifdef PERFORMANCE_STATISTIC
273 |         time_unit t_plugin_output = get_time_of_ms();
274 | #endif
275 |         if (0 != td_data.plugin->rknn_output(&td_data, output_unit)) {
276 |             d_rknn_infer_error("rknn_output failed")
277 |             continue;
278 |         }
279 | #ifdef PERFORMANCE_STATISTIC
280 |         {
281 |             std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_plugin_output_mutex);
282 |             m_statistic.s_plugin_output_count++;
283 |             m_statistic.s_plugin_output_ms += get_time_of_ms() - t_plugin_output;
284 |         }
285 | #endif
286 | 
287 |         // 释放输入资源
288 | #ifdef PERFORMANCE_STATISTIC
289 |         time_unit t_plugin_input_release = get_time_of_ms();
290 | #endif
291 |         if(0 != td_data.plugin->rknn_input_release(&td_data, pack.input_unit)){
292 |             d_rknn_infer_error("rknn_input_release failed")
293 |             continue;
294 |         }
295 | #ifdef PERFORMANCE_STATISTIC
296 |         {
297 |             std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_plugin_input_release_mutex);
298 |             m_statistic.s_plugin_input_release_count++;
299 |             m_statistic.s_plugin_input_release_ms += get_time_of_ms() - t_plugin_input_release;
300 |         }
301 | #endif
302 | 
303 |         // 释放资源
304 | #ifdef PERFORMANCE_STATISTIC
305 |         time_unit t_model_infer_release = get_time_of_ms();
306 | #endif
307 |         ret = m_rknn_models[idx]->model_infer_release(output_unit->n_outputs, output_unit->outputs);
308 |         if (ret != RetStatus::RET_STATUS_SUCCESS){
309 |             d_rknn_infer_error("model_infer_release failed")
310 |             continue;
311 |         }
312 | #ifdef PERFORMANCE_STATISTIC
313 |         {
314 |             std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_model_release_mutex);
315 |             m_statistic.s_model_release_count++;
316 |             m_statistic.s_model_release_ms += get_time_of_ms() - t_model_infer_release;
317 |         }
318 | #endif
319 | 
320 |         // 释放输出
321 |         free(output_unit->outputs);
322 |         free(output_unit);
323 |         // 释放输入
324 | //        delete pack.input_unit;
325 |     }
326 | 
327 |     // 插件反初始化
328 | #ifdef PERFORMANCE_STATISTIC
329 |     time_unit t_plugin_uninit = get_time_of_ms();
330 | #endif
331 |     if (0 != td_data.plugin->uninit(&td_data)) {
332 |         d_rknn_infer_error("rknn_infer uninit failed")
333 |         return;
334 |     }
335 | #ifdef PERFORMANCE_STATISTIC
336 |     {
337 |         std::lock_guard<std::mutex> proc_queue_lock(m_statistic.s_plugin_uninit_mutex);
338 |         m_statistic.s_plugin_uninit_count++;
339 |         m_statistic.s_plugin_uninit_ms += get_time_of_ms() - t_plugin_uninit;
340 |     }
341 | #endif
342 | }
343 | #ifdef PERFORMANCE_STATISTIC
344 | void RknnInfer::print_statistic() const {
345 |     d_time_info("queue_count: %d, queue_ms: %d, queue_avg_ms: %d",
346 |                 m_statistic.s_queue_count,
347 |                 m_statistic.s_queue_ms,
348 |                 m_statistic.s_queue_ms / m_statistic.s_queue_count)
349 | 
350 |     d_time_info("model_init_count: %d, model_init_ms: %d, model_init_avg_ms: %d",
351 |                 m_statistic.s_model_init_count,
352 |                 m_statistic.s_model_init_ms,
353 |                 m_statistic.s_model_init_ms / m_statistic.s_model_init_count)
354 |     d_time_info("model_infer_count: %d, model_infer_ms: %d, model_infer_avg_ms: %d",
355 |                 m_statistic.s_model_infer_count,
356 |                 m_statistic.s_model_infer_ms,
357 |                 m_statistic.s_model_infer_ms / m_statistic.s_model_infer_count)
358 |     d_time_info("model_release_count: %d, model_release_ms: %d, model_release_avg_ms: %d",
359 |                 m_statistic.s_model_release_count,
360 |                 m_statistic.s_model_release_ms,
361 |                 m_statistic.s_model_release_ms / m_statistic.s_model_release_count)
362 | 
363 |     d_time_info("plugin_init_count: %d, plugin_init_ms: %d, plugin_init_avg_ms: %d",
364 |                 m_statistic.s_plugin_init_count,
365 |                 m_statistic.s_plugin_init_ms,
366 |                 m_statistic.s_plugin_init_ms / m_statistic.s_plugin_init_count)
367 |     d_time_info("plugin_uninit_count: %d, plugin_uninit_ms: %d, plugin_uninit_avg_ms: %d",
368 |                 m_statistic.s_plugin_uninit_count,
369 |                 m_statistic.s_plugin_uninit_ms,
370 |                 m_statistic.s_plugin_uninit_ms / m_statistic.s_plugin_uninit_count)
371 |     d_time_info("plugin_output_count: %d, plugin_output_ms: %d, plugin_output_avg_ms: %d",
372 |                 m_statistic.s_plugin_output_count,
373 |                 m_statistic.s_plugin_output_ms,
374 |                 m_statistic.s_plugin_output_ms / m_statistic.s_plugin_output_count)
375 |     d_time_info("plugin_input_count: %d, plugin_input_ms: %d, plugin_input_avg_ms: %d",
376 |                 m_statistic.s_plugin_input_count,
377 |                 m_statistic.s_plugin_input_ms,
378 |                 m_statistic.s_plugin_input_ms / m_statistic.s_plugin_input_count)
379 |     d_time_info("plugin_input_release_count: %d, plugin_input_release_ms: %d, plugin_input_release_avg_ms: %d",
380 |                 m_statistic.s_plugin_input_release_count,
381 |                 m_statistic.s_plugin_input_release_ms,
382 |                 m_statistic.s_plugin_input_release_ms / m_statistic.s_plugin_input_release_count)
383 | }
384 | #endif


--------------------------------------------------------------------------------
/rknn_infer/rknn_infer.h:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.03
  5 |  * @brief: 推理调度头文件
  6 |  */
  7 | #ifndef RKNN_INFER_RKNN_INFER_H
  8 | #define RKNN_INFER_RKNN_INFER_H
  9 | #include <thread>
 10 | #include <mutex>
 11 | #include <list>
 12 | #include <vector>
 13 | #include <condition_variable>
 14 | #include "rknn_model.h"
 15 | #include "rknn_infer_api.h"
 16 | #include "plugin_ctrl.h"
 17 | 
 18 | struct QueuePack{
 19 | #ifdef PERFORMANCE_STATISTIC
 20 |     time_unit s_pack_record_ms;
 21 | #endif
 22 |     InputUnit* input_unit;
 23 |     void *plugin_sync_data;
 24 | };
 25 | #ifdef PERFORMANCE_STATISTIC
 26 | struct StaticStruct{
 27 |     // 模型初始化统计
 28 |     time_unit s_model_init_count;
 29 |     time_unit s_model_init_ms;
 30 |     // 模型推理统计
 31 |     std::mutex s_model_infer_mutex;
 32 |     time_unit s_model_infer_count;
 33 |     time_unit s_model_infer_ms;
 34 |     // 模型释放资源统计
 35 |     std::mutex s_model_release_mutex;
 36 |     time_unit s_model_release_count;
 37 |     time_unit s_model_release_ms;
 38 |     // 插件初始化统计
 39 |     std::mutex s_plugin_init_mutex;
 40 |     time_unit s_plugin_init_count;
 41 |     time_unit s_plugin_init_ms;
 42 |     // 插件反初始化统计
 43 |     std::mutex s_plugin_uninit_mutex;
 44 |     time_unit s_plugin_uninit_count;
 45 |     time_unit s_plugin_uninit_ms;
 46 |     // 插件输入调用统计
 47 |     std::mutex s_plugin_input_mutex;
 48 |     time_unit s_plugin_input_count;
 49 |     time_unit s_plugin_input_ms;
 50 |     // 插件输入释放统计
 51 |     std::mutex s_plugin_input_release_mutex;
 52 |     time_unit s_plugin_input_release_count;
 53 |     time_unit s_plugin_input_release_ms;
 54 |     // 插件输出统计
 55 |     std::mutex s_plugin_output_mutex;
 56 |     time_unit s_plugin_output_count;
 57 |     time_unit s_plugin_output_ms;
 58 |     // 队列调度统计
 59 |     std::mutex s_queue_mutex;
 60 |     time_unit s_queue_count;
 61 |     time_unit s_queue_ms;
 62 | 
 63 |     StaticStruct(){
 64 |         s_model_init_count = 0;
 65 |         s_model_init_ms = 0;
 66 | 
 67 |         s_model_infer_count = 0;
 68 |         s_model_infer_ms = 0;
 69 | 
 70 |         s_model_release_count = 0;
 71 |         s_model_release_ms = 0;
 72 | 
 73 |         s_plugin_init_count = 0;
 74 |         s_plugin_init_ms = 0;
 75 | 
 76 |         s_plugin_uninit_count = 0;
 77 |         s_plugin_uninit_ms = 0;
 78 | 
 79 |         s_plugin_input_count = 0;
 80 |         s_plugin_input_ms = 0;
 81 | 
 82 |         s_plugin_input_release_count = 0;
 83 |         s_plugin_input_release_ms = 0;
 84 | 
 85 |         s_plugin_output_count = 0;
 86 |         s_plugin_output_ms = 0;
 87 | 
 88 |         s_queue_count = 0;
 89 |         s_queue_ms = 0;
 90 |     }
 91 | };
 92 | #endif
 93 | class RknnInfer {
 94 | public:
 95 |     explicit RknnInfer(const std::string &model_name, const std::string &plugin_name);
 96 |     RetStatus stop();
 97 | 
 98 |     // 检查初始化
 99 |     [[nodiscard]] bool check_init() const;
100 | #ifdef PERFORMANCE_STATISTIC
101 |     void print_statistic() const;
102 | #endif
103 | private:
104 |     // 获取输入
105 |     RetStatus get_input_unit(QueuePack &pack);
106 |     // 填入输入
107 |     RetStatus put_input_unit(QueuePack &pack);
108 |     // 获取队列大小
109 |     uint32_t get_queue_size();
110 | 
111 |     // 输入处理线程
112 |     void input_data_thread(uint32_t idx);
113 |     // 输出处理线程
114 |     void infer_proc_thread(uint32_t idx);
115 | private:
116 |     bool m_init;
117 | #ifdef PERFORMANCE_STATISTIC
118 |     StaticStruct m_statistic;
119 | #endif
120 |     // 插件程序给调度程序的配置
121 |     PluginConfigGet m_plugin_get_config;
122 |     // 调度程序给插件程序的配置
123 |     PluginConfigSet m_plugin_set_config{};
124 |     // 调度队列
125 |     std::vector<std::thread> m_infer_proc_ctrl;
126 |     std::vector<ThreadData> m_infer_proc_meta;
127 |     // 输入调度
128 |     std::vector<RknnModel*> m_rknn_models;
129 | 
130 |     std::vector<std::thread> m_input_data_ctrl;
131 |     std::vector<ThreadData> m_input_data_meta;
132 | 
133 |     // 推理调度（和输出）
134 |     mutable std::mutex m_infer_queue_mutex;
135 |     std::condition_variable m_infer_queue_not_empty;
136 |     struct std::list<QueuePack> m_infer_queue;
137 | };
138 | 
139 | #endif //RKNN_INFER_RKNN_INFER_H
140 | 


--------------------------------------------------------------------------------
/rknn_infer/rknn_model.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.03
  5 |  * @brief: RKNN 模型接口封装
  6 |  */
  7 | #include <string>
  8 | #include <cstring>
  9 | #include <cstdio>
 10 | #include "rknn_model.h"
 11 | #include "utils_log.h"
 12 | 
 13 | static unsigned char* load_model(const char* filename, int* model_size) {
 14 |     FILE* fp = fopen(filename, "rb");
 15 |     if (fp == nullptr) {
 16 |         printf("fopen %s fail!\n", filename);
 17 |         return nullptr;
 18 |     }
 19 |     fseek(fp, 0, SEEK_END);
 20 |     int model_len = ftell(fp);
 21 |     auto* model = (unsigned char*)malloc(model_len);
 22 |     fseek(fp, 0, SEEK_SET);
 23 |     if (model_len != fread(model, 1, model_len, fp)) {
 24 |         printf("fread %s fail!\n", filename);
 25 |         free(model);
 26 |         return nullptr;
 27 |     }
 28 |     *model_size = model_len;
 29 |     fclose(fp);
 30 |     return model;
 31 | }
 32 | 
 33 | static void dump_tensor_attr(rknn_tensor_attr* attr) {
 34 |     d_rknn_model_info("index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
 35 |            "zp=%d, scale=%f",
 36 |            attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
 37 |            attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
 38 |            get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
 39 | }
 40 | 
 41 | RknnModel::RknnModel(const std::string &model_path, PluginConfigSet &plugin_config_set, bool show_model):m_plugin_config_set(plugin_config_set) {
 42 |     // 初始化变量
 43 |     init = false;
 44 |     rk_model_ctx = 0;
 45 | 
 46 |     // rknn 模型初始化
 47 |     int model_len = 0;
 48 |     m_model = load_model(model_path.c_str() , &model_len);
 49 |     if (m_model == nullptr){
 50 |         d_rknn_model_error("load m_model fail!")
 51 |         return;
 52 |     }
 53 |     int ret = rknn_init(&rk_model_ctx, m_model, model_len, 0, nullptr);
 54 |     if(ret != 0) {
 55 |         d_rknn_model_error("rknn_init fail! ret=%d", ret)
 56 |         return;
 57 |     }
 58 | 
 59 |     // 模型信息查询
 60 |     rknn_sdk_version version;
 61 |     ret = rknn_query(rk_model_ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version));
 62 |     if (ret < 0) {
 63 |         d_rknn_model_error("rknn_init error ret=%d\n", ret);
 64 |         return;
 65 |     }
 66 |     memcpy(&plugin_config_set.sdk_version, &version, sizeof(rknn_sdk_version));
 67 |     CHECK(show_model, true, d_rknn_model_info("sdk version: %s driver version: %s\n", version.api_version, version.drv_version));
 68 | 
 69 |     rknn_input_output_num io_num;
 70 |     ret = rknn_query(rk_model_ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
 71 |     if (ret != RKNN_SUCC) {
 72 |         d_rknn_model_error("rknn_query fail! ret=%d", ret)
 73 |         return;
 74 |     }
 75 |     plugin_config_set.io_num = io_num;
 76 |     CHECK(show_model, true, d_rknn_model_info("m_model input num: %d, output num: %d", io_num.n_input, io_num.n_output))
 77 | 
 78 |     CHECK(show_model, true, d_rknn_model_info("input tensors:"))
 79 |     plugin_config_set.input_attr = new rknn_tensor_attr[io_num.n_input];
 80 |     memset(plugin_config_set.input_attr, 0, sizeof(rknn_tensor_attr) * io_num.n_input);
 81 |     for (int i = 0; i < io_num.n_input; i++) {
 82 |         plugin_config_set.input_attr[i].index = i;
 83 |         ret = rknn_query(rk_model_ctx, RKNN_QUERY_INPUT_ATTR, &(plugin_config_set.input_attr[i]), sizeof(rknn_tensor_attr));
 84 |         if (ret != RKNN_SUCC) {
 85 |             d_rknn_model_error("rknn_query fail! ret=%d", ret);
 86 |             return;
 87 |         }
 88 |         CHECK(show_model, true, dump_tensor_attr(&(plugin_config_set.input_attr[i]));)
 89 |     }
 90 | 
 91 |     CHECK(show_model, true, d_rknn_model_info("output tensors:"))
 92 |     plugin_config_set.output_attr = new rknn_tensor_attr[io_num.n_output];
 93 |     memset(plugin_config_set.output_attr, 0, sizeof(rknn_tensor_attr) * io_num.n_output);
 94 |     for (int i = 0; i < io_num.n_output; i++) {
 95 |         plugin_config_set.output_attr[i].index = i;
 96 |         ret = rknn_query(rk_model_ctx, RKNN_QUERY_OUTPUT_ATTR, &(plugin_config_set.output_attr[i]), sizeof(rknn_tensor_attr));
 97 |         if (ret != RKNN_SUCC) {
 98 |             d_rknn_model_error("rknn_query fail! ret=%d", ret);
 99 |             return;
100 |         }
101 |         CHECK(show_model, true, dump_tensor_attr(&(plugin_config_set.output_attr[i]));)
102 |     }
103 |     d_rknn_model_info("rknn m_model init success! rk_model_ctx：%lu", rk_model_ctx)
104 |     is_dup = true;
105 |     init = true;
106 | }
107 | 
108 | RknnModel *RknnModel::model_infer_dup() const {
109 |     return new RknnModel(this->rk_model_ctx, this->m_plugin_config_set);
110 | }
111 | 
112 | RknnModel::RknnModel(rknn_context ctx, PluginConfigSet &plugin_config_set): m_plugin_config_set(plugin_config_set){
113 |     // 初始化变量
114 |     init = false;
115 |     this->rk_model_ctx = 0;
116 |     this->m_model = nullptr;
117 |     this->is_dup = true;
118 | 
119 |     // 复制 rknn 模型， 做权重复用
120 |     int ret = rknn_dup_context(&ctx, &this->rk_model_ctx);
121 |     if (ret != RKNN_SUCC){
122 |         d_rknn_model_error("rknn dup model fail! ret=%d", ret)
123 |         return;
124 |     }
125 |     d_rknn_model_error("rknn dup model success!")
126 |     init = true;
127 | }
128 | 
129 | 
130 | RknnModel::~RknnModel() {
131 |     // 销毁 rknn 模型
132 |     if(rk_model_ctx != 0){
133 |         rknn_destroy(rk_model_ctx);
134 |         rk_model_ctx = 0;
135 |     }
136 |     if(m_model != nullptr){
137 |         free(m_model);
138 |     }
139 | 
140 |     // 销毁配置信息
141 |     if(! is_dup){
142 |         if(m_plugin_config_set.input_attr != nullptr){
143 |             delete[] m_plugin_config_set.input_attr;
144 |             m_plugin_config_set.input_attr = nullptr;
145 |         }
146 |         if(m_plugin_config_set.output_attr != nullptr){
147 |             delete[] m_plugin_config_set.output_attr;
148 |             m_plugin_config_set.output_attr = nullptr;
149 |         }
150 |     }
151 | }
152 | 
153 | bool RknnModel::check_init() const {
154 |     return init;
155 | }
156 | 
157 | RetStatus RknnModel::model_infer_sync(
158 |         uint32_t n_inputs, rknn_input *inputs,
159 |         uint32_t n_outputs, rknn_output *outputs
160 |         ) const {
161 |     d_rknn_model_debug("rknn_inputs_set rk_model_ctx=%lu, n_inputs:%u, inputs:%p", rk_model_ctx, n_inputs, inputs)
162 |     d_rknn_model_debug("rknn_inputs index: %d", inputs[0].index)
163 |     d_rknn_model_debug("rknn_inputs type: %d", inputs[0].type)
164 |     d_rknn_model_debug("rknn_inputs size: %d", inputs[0].size)
165 |     d_rknn_model_debug("rknn_inputs fmt: %d", inputs[0].fmt)
166 |     int ret = rknn_inputs_set(rk_model_ctx, n_inputs, inputs);
167 |     if (ret < 0) {
168 |         d_rknn_model_error("rknn_input_set fail! ret=%d", ret);
169 |         return RET_STATUS_FAILED;
170 |     }
171 | 
172 |     d_rknn_model_debug("rknn_run");
173 |     ret = rknn_run(rk_model_ctx, nullptr);
174 |     if (ret < 0) {
175 |         d_rknn_model_error("rknn_run fail! ret=%d", ret);
176 |         return RET_STATUS_FAILED;
177 |     }
178 | 
179 |     ret = rknn_outputs_get(rk_model_ctx, n_outputs, outputs, nullptr);
180 |     if (ret < 0) {
181 |         d_rknn_model_error("rknn_outputs_get fail! ret=%d", ret);
182 |         return RET_STATUS_FAILED;
183 |     }
184 |     return RET_STATUS_SUCCESS;
185 | }
186 | 
187 | RetStatus RknnModel::model_infer_release(uint32_t n_outputs, rknn_output *outputs) const {
188 |     int ret = rknn_outputs_release(rk_model_ctx, n_outputs, outputs);
189 |     if(ret != RKNN_SUCC){
190 |         d_rknn_model_error("rknn_outputs_release fail! ret=%d", ret)
191 |     }
192 |     return RetStatus::RET_STATUS_SUCCESS;
193 | }
194 | 
195 | 


--------------------------------------------------------------------------------
/rknn_infer/rknn_model.h:
--------------------------------------------------------------------------------
 1 | /**
 2 |  * @author: bo.liu
 3 |  * @mail: geniusrabbit@qq.com
 4 |  * @date: 2023.08.03
 5 |  * @brief: RKNN 模型接口封装
 6 |  */
 7 | #ifndef PLUGIN_RKNN_IMAGE_RKNN_MODEL_H
 8 | #define PLUGIN_RKNN_IMAGE_RKNN_MODEL_H
 9 | 
10 | #include "rknn_api.h"
11 | #include "rknn_matmul_api.h"
12 | #include "utils.h"
13 | #include "rknn_infer_api.h"
14 | 
15 | class RknnModel {
16 | public:
17 |     explicit RknnModel(const std::string &model_path, PluginConfigSet &plugin_config_set, bool show_model=false);
18 |     ~RknnModel();
19 | 
20 |     // 检查初始化
21 |     [[nodiscard]] bool check_init() const;
22 | 
23 |     // 模型推理
24 |     RetStatus model_infer_sync(uint32_t n_inputs, rknn_input *inputs, uint32_t n_outputs, rknn_output *outputs) const;
25 |     // 释放推理资源
26 |     RetStatus model_infer_release(uint32_t n_outputs, rknn_output *outputs) const;
27 | 
28 |     // 模型复用
29 |     [[nodiscard]] RknnModel *model_infer_dup() const;
30 | 
31 | private:
32 |     // 内部模型上下文拷贝接口
33 |     explicit RknnModel(rknn_context ctx, PluginConfigSet &plugin_config_set);
34 | private:
35 |     // 初始化记录
36 |     bool init;
37 |     bool is_dup;
38 |     rknn_context rk_model_ctx;
39 |     unsigned char* m_model{};
40 |     PluginConfigSet &m_plugin_config_set;
41 | };
42 | #endif //PLUGIN_RKNN_IMAGE_RKNN_MODEL_H
43 | 


--------------------------------------------------------------------------------
/rknn_infer/utils.h:
--------------------------------------------------------------------------------
 1 | /**
 2 |  * @author: bo.liu
 3 |  * @mail: geniusrabbit@qq.com
 4 |  * @date: 2023.08.03
 5 |  * @brief: 对外开放插件接口
 6 |  */
 7 | #ifndef PLUGIN_RKNN_IMAGE_UTILS_H
 8 | #define PLUGIN_RKNN_IMAGE_UTILS_H
 9 | #include <ctime>
10 | #ifdef __linux__
11 | #include <sys/socket.h>
12 | #include <arpa/inet.h>
13 | #include<sys/time.h>
14 | #else
15 | #include <WS2tcpip.h>
16 | #endif
17 | 
18 | typedef uint64_t time_unit;
19 | 
20 | enum RetStatus {
21 |     // 成功
22 |     RET_STATUS_SUCCESS,
23 |     // 失败
24 |     RET_STATUS_FAILED,
25 |     // 未知
26 |     RET_STATUS_UNKNOWN,
27 | };
28 | 
29 | /* 获取NS时间 -9 */
30 | static time_unit getTimeOfNs() {
31 |     struct timespec tv{};
32 |     clock_gettime(CLOCK_REALTIME, &tv);
33 |     return tv.tv_sec*1000000000 + tv.tv_nsec;
34 | }
35 | 
36 | #if defined(_WIN32)
37 | #include <chrono>
38 | #include <utility>
39 | static int gettimeofday(struct timeval* tp, struct timezone* tzp) {
40 |     namespace sc = std::chrono;
41 |     sc::system_clock::duration d = sc::system_clock::now().time_since_epoch();
42 |     sc::seconds s = sc::duration_cast<sc::seconds>(d);
43 |     tp->tv_sec = long(s.count());
44 |     tp->tv_usec = long(sc::duration_cast<sc::microseconds>(d - s).count());
45 |     return 0;
46 | }
47 | #endif // _WIN32
48 | 
49 | /* 获取MS时间 -3 */
50 | static time_unit get_time_of_ms(){
51 |     struct timeval tv{};
52 |     gettimeofday(&tv, nullptr);
53 |     return tv.tv_sec * (time_unit)1000 + tv.tv_usec / 1000;
54 | }
55 | 
56 | /* 获取S时间 */
57 | static time_unit get_time_of_s(){
58 |     struct timeval tv{};
59 |     gettimeofday(&tv, nullptr);
60 |     return tv.tv_sec;
61 | }
62 | 
63 | /* 精确睡眠US时间 */
64 | static void sleepUS(time_unit usec){
65 | #if defined(_WIN32)
66 |     Sleep(usec/1000);
67 | #else
68 |     struct timeval tv{};
69 |     tv.tv_sec = long(usec / 1000000UL);
70 |     tv.tv_usec = long(usec % 1000000UL);
71 |     errno = 0;
72 |     select(0, nullptr, nullptr, nullptr, &tv);
73 | #endif
74 | }
75 | 
76 | #endif //PLUGIN_RKNN_IMAGE_UTILS_H
77 | 


--------------------------------------------------------------------------------
/rknn_infer/utils_log.h:
--------------------------------------------------------------------------------
 1 | /**
 2 |  * @author: bo.liu
 3 |  * @mail: geniusrabbit@qq.com
 4 |  * @date: 2023.08.03
 5 |  * @brief: 日志配置
 6 |  */
 7 | #ifndef RKNN_INFER_UTILS_LOG_H
 8 | #define RKNN_INFER_UTILS_LOG_H
 9 | 
10 | #include "dlog.h"
11 | 
12 | #ifndef __linux__
13 | #include "windows.h"
14 | #else
15 | #include "unistd.h"
16 | #endif
17 | 
18 | #ifndef __linux__
19 | static uint64_t get_thread_id(){
20 |     return GetCurrentThreadId();
21 | }
22 | static uint64_t get_process_id(){
23 |     return GetCurrentProcessId();
24 | }
25 | #else
26 | #if __GLIBC__ == 2 && __GLIBC_MINOR__ < 30
27 | #include <sys/syscall.h>
28 | #define gettid() syscall(SYS_gettid)
29 | #endif
30 | static uint64_t get_thread_id(){
31 |     return gettid();
32 | }
33 | static uint64_t get_process_id(){
34 |     return getpid();
35 | }
36 | #endif
37 | 
38 | #define DLOG_FORMAT_PREFIX "<%d, %s, %s, %d> "
39 | #ifndef __linux__
40 | #define __FILENAME__ ( __builtin_strrchr(__FILE__, '\\') ? __builtin_strrchr(__FILE__, '\\') + 1 : __FILE__ )
41 | #else
42 | #define __FILENAME__ ( __builtin_strrchr(__FILE__, '/') ? __builtin_strrchr(__FILE__, '/') + 1 : __FILE__ )
43 | #endif
44 | 
45 | #define d_rknn_infer_error(format, ...) log(LOG_MODULE_INIT(d_rknn_infer), LOG_ERROR, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
46 | #define d_rknn_infer_info(format, ...)  log(LOG_MODULE_INIT(d_rknn_infer), LOG_INFO, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
47 | #define d_rknn_infer_warn(format, ...) log(LOG_MODULE_INIT(d_rknn_infer), LOG_WARN, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
48 | #define d_rknn_infer_debug(format, ...) log(LOG_MODULE_INIT(d_rknn_infer), LOG_DEBUG, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
49 | 
50 | #define d_rknn_model_error(format, ...) log(LOG_MODULE_INIT(d_rknn_model), LOG_ERROR, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
51 | #define d_rknn_model_info(format, ...)  log(LOG_MODULE_INIT(d_rknn_model), LOG_INFO, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
52 | #define d_rknn_model_warn(format, ...) log(LOG_MODULE_INIT(d_rknn_model), LOG_WARN, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
53 | #define d_rknn_model_debug(format, ...) log(LOG_MODULE_INIT(d_rknn_model), LOG_DEBUG, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
54 | 
55 | #define d_rknn_plugin_error(format, ...) log(LOG_MODULE_INIT(d_rknn_plugin), LOG_ERROR, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
56 | #define d_rknn_plugin_info(format, ...)  log(LOG_MODULE_INIT(d_rknn_plugin), LOG_INFO, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
57 | #define d_rknn_plugin_warn(format, ...) log(LOG_MODULE_INIT(d_rknn_plugin), LOG_WARN, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
58 | #define d_rknn_plugin_debug(format, ...) log(LOG_MODULE_INIT(d_rknn_plugin), LOG_DEBUG, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
59 | 
60 | #define d_time_warn(format, ...)  log(LOG_MODULE_INIT(d_time), LOG_WARN, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
61 | #define d_time_info(format, ...)  log(LOG_MODULE_INIT(d_time), LOG_INFO, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
62 | #define d_time_debug(format, ...) log(LOG_MODULE_INIT(d_time), LOG_DEBUG, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
63 | 
64 | #define d_mpp_module_error(format, ...) log(LOG_MODULE_INIT(d_mpp_module), LOG_ERROR, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
65 | #define d_mpp_module_info(format, ...)  log(LOG_MODULE_INIT(d_mpp_module), LOG_INFO, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
66 | #define d_mpp_module_warn(format, ...) log(LOG_MODULE_INIT(d_mpp_module), LOG_WARN, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
67 | #define d_mpp_module_debug(format, ...) log(LOG_MODULE_INIT(d_mpp_module), LOG_DEBUG, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
68 | 
69 | #define d_unit_test_error(format, ...) log(LOG_MODULE_INIT(d_unit_test), LOG_ERROR, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
70 | #define d_unit_test_info(format, ...)  log(LOG_MODULE_INIT(d_unit_test), LOG_INFO, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
71 | #define d_unit_test_warn(format, ...) log(LOG_MODULE_INIT(d_unit_test), LOG_WARN, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
72 | #define d_unit_test_debug(format, ...) log(LOG_MODULE_INIT(d_unit_test), LOG_DEBUG, DLOG_FORMAT_PREFIX#format, get_thread_id(), __FILENAME__, __FUNCTION__ , __LINE__, ##__VA_ARGS__);
73 | 
74 | #endif //RKNN_INFER_UTILS_LOG_H
75 | 


--------------------------------------------------------------------------------
/rknn_infer_api/rknn_infer_api.h:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.03
  5 |  * @brief: 对外开放插件接口
  6 |  */
  7 | #ifndef RKNN_INFER_RKNN_INFER_API_H
  8 | #define RKNN_INFER_RKNN_INFER_API_H
  9 | 
 10 | #include "rknn_api.h"
 11 | #include "rknn_matmul_api.h"
 12 | 
 13 | // 插件在加载和关闭时，自动调用的构造函数和析构函数
 14 | #define plugin_init	__attribute__((constructor))
 15 | #define plugin_exit	__attribute__((destructor))
 16 | 
 17 | // 这里对线程类型做区分，因为本推理模板的线程调度是多对多的
 18 | enum ThreadType{
 19 |     // 输入类型线程，做数据采集和预处理
 20 |     THREAD_TYPE_INPUT,
 21 |     // 输出类型线程，做推理的输出
 22 |     THREAD_TYPE_OUTPUT,
 23 | };
 24 | 
 25 | // 输入单元，包含输入数据和输入数据数量
 26 | struct InputUnit{
 27 |     // 输入数据
 28 |     rknn_input *inputs;
 29 |     // 输入数据数量
 30 |     uint32_t n_inputs;
 31 | };
 32 | 
 33 | // 输出单元，包含输出数据和输出数据数量
 34 | struct OutputUnit{
 35 |     // 输出数据
 36 |     rknn_output *outputs;
 37 |     // 输出数据数量
 38 |     uint32_t n_outputs;
 39 | };
 40 | 
 41 | // 线程数据，保存一个线程用到的所有数据
 42 | struct PluginStruct;
 43 | struct ThreadData {
 44 |     // 线程数据(线程ID和线程类型)
 45 |     uint32_t thread_id;
 46 |     ThreadType thread_type;
 47 | 
 48 |     // 每个线程的私有数据（输入线程和输出线程）
 49 |     void *plugin_private_data;
 50 |     // 输入向输出线程的同步数据（输入线程申请，输出线程释放）
 51 |     void *plugin_sync_data;
 52 | 
 53 |     // 共享接口
 54 |     PluginStruct *plugin;
 55 | };
 56 | 
 57 | // 插件程序给调度程序的配置
 58 | struct PluginConfigGet{
 59 |     // 输入线程个数
 60 |     uint32_t input_thread_nums;
 61 |     // 输出线程个数
 62 |     uint32_t output_thread_nums;
 63 | 
 64 |     // 任务队列个数限制(0代表无限制)，降低任务处理延时
 65 |     uint32_t task_queue_limit;
 66 | 
 67 |     // 是否需要输出 float 类型的输出结果
 68 |     bool output_want_float;
 69 | 
 70 |     // 默认配置
 71 |     PluginConfigGet(){
 72 |         input_thread_nums = 1;
 73 | 
 74 |         output_thread_nums = 1;
 75 | 
 76 |         task_queue_limit = 100;
 77 | 
 78 |         output_want_float = true;
 79 |     }
 80 | };
 81 | 
 82 | // 调度程序给插件程序的配置
 83 | struct PluginConfigSet{
 84 |     // 模型版本
 85 |     rknn_sdk_version sdk_version;
 86 |     // 输入输出 tensor 个数
 87 |     rknn_input_output_num io_num;
 88 |     // 输入 tensor 特征
 89 |     rknn_tensor_attr *input_attr;
 90 |     // 输出 tensor 特征
 91 |     rknn_tensor_attr *output_attr;
 92 | };
 93 | 
 94 | // 插件接口格式定义
 95 | struct PluginStruct {
 96 |     // 插件名称
 97 |     const char *plugin_name;
 98 |     // 插件版本
 99 |     int plugin_version;
100 | 
101 |     // 从插件中获取调度配置
102 |     int (*get_config)(PluginConfigGet *plugin_config);
103 | 
104 |     // 给插件设置运行配置
105 |     int (*set_config)(PluginConfigSet *plugin_config);
106 | 
107 |     // 插件初始化
108 |     int (*init)(struct ThreadData *);
109 |     // 插件反初始化
110 |     int (*uninit)(struct ThreadData *);
111 | 
112 |     // 收集输入数据和释放输入资源
113 |     int (*rknn_input)(struct ThreadData *, struct InputUnit *);
114 |     int (*rknn_input_release)(struct ThreadData *, struct InputUnit *);
115 | 
116 |     int (*rknn_output)(struct ThreadData *, struct OutputUnit *);
117 | };
118 | 
119 | // 插件向主程序注册和反注册接口
120 | extern void plugin_register(struct PluginStruct *);
121 | extern void plugin_unregister(struct PluginStruct *);
122 | #endif //RKNN_INFER_RKNN_INFER_API_H
123 | 


--------------------------------------------------------------------------------
/rknn_plugins/image_op_utils.h:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.10
  5 |  * @brief: 图像处理函数，来源于官方的 drawing.cpp
  6 |  */
  7 | #ifndef RKNN_INFER_PLUGIN_IMAGE_OP_UTILS_H
  8 | #define RKNN_INFER_PLUGIN_IMAGE_OP_UTILS_H
  9 | 
 10 | #include "string.h"
 11 | 
 12 | #define max(a, b) (((a) > (b)) ? (a) : (b))
 13 | #define min(a, b) (((a) < (b)) ? (a) : (b))
 14 | 
 15 | static void draw_rectangle_c1(unsigned char* pixels, int w, int h, int stride, int rx, int ry, int rw, int rh, unsigned int color, int thickness)
 16 | {
 17 |     const unsigned char* pen_color = (const unsigned char*)&color;
 18 | 
 19 |     if (thickness == -1)
 20 |     {
 21 |         // filled
 22 |         for (int y = ry; y < ry + rh; y++)
 23 |         {
 24 |             if (y < 0)
 25 |                 continue;
 26 | 
 27 |             if (y >= h)
 28 |                 break;
 29 | 
 30 |             unsigned char* p = pixels + stride * y;
 31 | 
 32 |             for (int x = rx; x < rx + rw; x++)
 33 |             {
 34 |                 if (x < 0)
 35 |                     continue;
 36 | 
 37 |                 if (x >= w)
 38 |                     break;
 39 | 
 40 |                 p[x] = pen_color[0];
 41 |             }
 42 |         }
 43 | 
 44 |         return;
 45 |     }
 46 | 
 47 |     const int t0 = thickness / 2;
 48 |     const int t1 = thickness - t0;
 49 | 
 50 |     // draw top
 51 |     {
 52 |         for (int y = ry - t0; y < ry + t1; y++)
 53 |         {
 54 |             if (y < 0)
 55 |                 continue;
 56 | 
 57 |             if (y >= h)
 58 |                 break;
 59 | 
 60 |             unsigned char* p = pixels + stride * y;
 61 | 
 62 |             for (int x = rx - t0; x < rx + rw + t1; x++)
 63 |             {
 64 |                 if (x < 0)
 65 |                     continue;
 66 | 
 67 |                 if (x >= w)
 68 |                     break;
 69 | 
 70 |                 p[x] = pen_color[0];
 71 |             }
 72 |         }
 73 |     }
 74 | 
 75 |     // draw bottom
 76 |     {
 77 |         for (int y = ry + rh - t0; y < ry + rh + t1; y++)
 78 |         {
 79 |             if (y < 0)
 80 |                 continue;
 81 | 
 82 |             if (y >= h)
 83 |                 break;
 84 | 
 85 |             unsigned char* p = pixels + stride * y;
 86 | 
 87 |             for (int x = rx - t0; x < rx + rw + t1; x++)
 88 |             {
 89 |                 if (x < 0)
 90 |                     continue;
 91 | 
 92 |                 if (x >= w)
 93 |                     break;
 94 | 
 95 |                 p[x] = pen_color[0];
 96 |             }
 97 |         }
 98 |     }
 99 | 
100 |     // draw left
101 |     for (int x = rx - t0; x < rx + t1; x++)
102 |     {
103 |         if (x < 0)
104 |             continue;
105 | 
106 |         if (x >= w)
107 |             break;
108 | 
109 |         for (int y = ry + t1; y < ry + rh - t0; y++)
110 |         {
111 |             if (y < 0)
112 |                 continue;
113 | 
114 |             if (y >= h)
115 |                 break;
116 | 
117 |             unsigned char* p = pixels + stride * y;
118 | 
119 |             p[x] = pen_color[0];
120 |         }
121 |     }
122 | 
123 |     // draw right
124 |     for (int x = rx + rw - t0; x < rx + rw + t1; x++)
125 |     {
126 |         if (x < 0)
127 |             continue;
128 | 
129 |         if (x >= w)
130 |             break;
131 | 
132 |         for (int y = ry + t1; y < ry + rh - t0; y++)
133 |         {
134 |             if (y < 0)
135 |                 continue;
136 | 
137 |             if (y >= h)
138 |                 break;
139 | 
140 |             unsigned char* p = pixels + stride * y;
141 | 
142 |             p[x] = pen_color[0];
143 |         }
144 |     }
145 | }
146 | 
147 | static void draw_rectangle_c2(unsigned char* pixels, int w, int h, int stride, int rx, int ry, int rw, int rh, unsigned int color, int thickness)
148 | {
149 |     const unsigned char* pen_color = (const unsigned char*)&color;
150 | 
151 |     if (thickness == -1)
152 |     {
153 |         // filled
154 |         for (int y = ry; y < ry + rh; y++)
155 |         {
156 |             if (y < 0)
157 |                 continue;
158 | 
159 |             if (y >= h)
160 |                 break;
161 | 
162 |             unsigned char* p = pixels + stride * y;
163 | 
164 |             for (int x = rx; x < rx + rw; x++)
165 |             {
166 |                 if (x < 0)
167 |                     continue;
168 | 
169 |                 if (x >= w)
170 |                     break;
171 | 
172 |                 p[x * 2 + 0] = pen_color[0];
173 |                 p[x * 2 + 1] = pen_color[1];
174 |             }
175 |         }
176 | 
177 |         return;
178 |     }
179 | 
180 |     const int t0 = thickness / 2;
181 |     const int t1 = thickness - t0;
182 | 
183 |     // draw top
184 |     {
185 |         for (int y = ry - t0; y < ry + t1; y++)
186 |         {
187 |             if (y < 0)
188 |                 continue;
189 | 
190 |             if (y >= h)
191 |                 break;
192 | 
193 |             unsigned char* p = pixels + stride * y;
194 | 
195 |             for (int x = rx - t0; x < rx + rw + t1; x++)
196 |             {
197 |                 if (x < 0)
198 |                     continue;
199 | 
200 |                 if (x >= w)
201 |                     break;
202 | 
203 |                 p[x * 2 + 0] = pen_color[0];
204 |                 p[x * 2 + 1] = pen_color[1];
205 |             }
206 |         }
207 |     }
208 | 
209 |     // draw bottom
210 |     {
211 |         for (int y = ry + rh - t0; y < ry + rh + t1; y++)
212 |         {
213 |             if (y < 0)
214 |                 continue;
215 | 
216 |             if (y >= h)
217 |                 break;
218 | 
219 |             unsigned char* p = pixels + stride * y;
220 | 
221 |             for (int x = rx - t0; x < rx + rw + t1; x++)
222 |             {
223 |                 if (x < 0)
224 |                     continue;
225 | 
226 |                 if (x >= w)
227 |                     break;
228 | 
229 |                 p[x * 2 + 0] = pen_color[0];
230 |                 p[x * 2 + 1] = pen_color[1];
231 |             }
232 |         }
233 |     }
234 | 
235 |     // draw left
236 |     for (int x = rx - t0; x < rx + t1; x++)
237 |     {
238 |         if (x < 0)
239 |             continue;
240 | 
241 |         if (x >= w)
242 |             break;
243 | 
244 |         for (int y = ry + t1; y < ry + rh - t0; y++)
245 |         {
246 |             if (y < 0)
247 |                 continue;
248 | 
249 |             if (y >= h)
250 |                 break;
251 | 
252 |             unsigned char* p = pixels + stride * y;
253 | 
254 |             p[x * 2 + 0] = pen_color[0];
255 |             p[x * 2 + 1] = pen_color[1];
256 |         }
257 |     }
258 | 
259 |     // draw right
260 |     for (int x = rx + rw - t0; x < rx + rw + t1; x++)
261 |     {
262 |         if (x < 0)
263 |             continue;
264 | 
265 |         if (x >= w)
266 |             break;
267 | 
268 |         for (int y = ry + t1; y < ry + rh - t0; y++)
269 |         {
270 |             if (y < 0)
271 |                 continue;
272 | 
273 |             if (y >= h)
274 |                 break;
275 | 
276 |             unsigned char* p = pixels + stride * y;
277 | 
278 |             p[x * 2 + 0] = pen_color[0];
279 |             p[x * 2 + 1] = pen_color[1];
280 |         }
281 |     }
282 | }
283 | 
284 | static void draw_rectangle_c3(unsigned char* pixels, int w, int h, int stride, int rx, int ry, int rw, int rh, unsigned int color, int thickness)
285 | {
286 |     const unsigned char* pen_color = (const unsigned char*)&color;
287 | 
288 |     if (thickness == -1)
289 |     {
290 |         // filled
291 |         for (int y = ry; y < ry + rh; y++)
292 |         {
293 |             if (y < 0)
294 |                 continue;
295 | 
296 |             if (y >= h)
297 |                 break;
298 | 
299 |             unsigned char* p = pixels + stride * y;
300 | 
301 |             for (int x = rx; x < rx + rw; x++)
302 |             {
303 |                 if (x < 0)
304 |                     continue;
305 | 
306 |                 if (x >= w)
307 |                     break;
308 | 
309 |                 p[x * 3 + 0] = pen_color[0];
310 |                 p[x * 3 + 1] = pen_color[1];
311 |                 p[x * 3 + 2] = pen_color[2];
312 |             }
313 |         }
314 | 
315 |         return;
316 |     }
317 | 
318 |     const int t0 = thickness / 2;
319 |     const int t1 = thickness - t0;
320 | 
321 |     // draw top
322 |     {
323 |         for (int y = ry - t0; y < ry + t1; y++)
324 |         {
325 |             if (y < 0)
326 |                 continue;
327 | 
328 |             if (y >= h)
329 |                 break;
330 | 
331 |             unsigned char* p = pixels + stride * y;
332 | 
333 |             for (int x = rx - t0; x < rx + rw + t1; x++)
334 |             {
335 |                 if (x < 0)
336 |                     continue;
337 | 
338 |                 if (x >= w)
339 |                     break;
340 | 
341 |                 p[x * 3 + 0] = pen_color[0];
342 |                 p[x * 3 + 1] = pen_color[1];
343 |                 p[x * 3 + 2] = pen_color[2];
344 |             }
345 |         }
346 |     }
347 | 
348 |     // draw bottom
349 |     {
350 |         for (int y = ry + rh - t0; y < ry + rh + t1; y++)
351 |         {
352 |             if (y < 0)
353 |                 continue;
354 | 
355 |             if (y >= h)
356 |                 break;
357 | 
358 |             unsigned char* p = pixels + stride * y;
359 | 
360 |             for (int x = rx - t0; x < rx + rw + t1; x++)
361 |             {
362 |                 if (x < 0)
363 |                     continue;
364 | 
365 |                 if (x >= w)
366 |                     break;
367 | 
368 |                 p[x * 3 + 0] = pen_color[0];
369 |                 p[x * 3 + 1] = pen_color[1];
370 |                 p[x * 3 + 2] = pen_color[2];
371 |             }
372 |         }
373 |     }
374 | 
375 |     // draw left
376 |     for (int x = rx - t0; x < rx + t1; x++)
377 |     {
378 |         if (x < 0)
379 |             continue;
380 | 
381 |         if (x >= w)
382 |             break;
383 | 
384 |         for (int y = ry + t1; y < ry + rh - t0; y++)
385 |         {
386 |             if (y < 0)
387 |                 continue;
388 | 
389 |             if (y >= h)
390 |                 break;
391 | 
392 |             unsigned char* p = pixels + stride * y;
393 | 
394 |             p[x * 3 + 0] = pen_color[0];
395 |             p[x * 3 + 1] = pen_color[1];
396 |             p[x * 3 + 2] = pen_color[2];
397 |         }
398 |     }
399 | 
400 |     // draw right
401 |     for (int x = rx + rw - t0; x < rx + rw + t1; x++)
402 |     {
403 |         if (x < 0)
404 |             continue;
405 | 
406 |         if (x >= w)
407 |             break;
408 | 
409 |         for (int y = ry + t1; y < ry + rh - t0; y++)
410 |         {
411 |             if (y < 0)
412 |                 continue;
413 | 
414 |             if (y >= h)
415 |                 break;
416 | 
417 |             unsigned char* p = pixels + stride * y;
418 | 
419 |             p[x * 3 + 0] = pen_color[0];
420 |             p[x * 3 + 1] = pen_color[1];
421 |             p[x * 3 + 2] = pen_color[2];
422 |         }
423 |     }
424 | }
425 | 
426 | static void draw_rectangle_c4(unsigned char* pixels, int w, int h, int stride, int rx, int ry, int rw, int rh, unsigned int color, int thickness)
427 | {
428 |     const unsigned char* pen_color = (const unsigned char*)&color;
429 | 
430 |     if (thickness == -1)
431 |     {
432 |         // filled
433 |         for (int y = ry; y < ry + rh; y++)
434 |         {
435 |             if (y < 0)
436 |                 continue;
437 | 
438 |             if (y >= h)
439 |                 break;
440 | 
441 |             unsigned char* p = pixels + stride * y;
442 | 
443 |             for (int x = rx; x < rx + rw; x++)
444 |             {
445 |                 if (x < 0)
446 |                     continue;
447 | 
448 |                 if (x >= w)
449 |                     break;
450 | 
451 |                 p[x * 4 + 0] = pen_color[0];
452 |                 p[x * 4 + 1] = pen_color[1];
453 |                 p[x * 4 + 2] = pen_color[2];
454 |                 p[x * 4 + 3] = pen_color[3];
455 |             }
456 |         }
457 | 
458 |         return;
459 |     }
460 | 
461 |     const int t0 = thickness / 2;
462 |     const int t1 = thickness - t0;
463 | 
464 |     // draw top
465 |     {
466 |         for (int y = ry - t0; y < ry + t1; y++)
467 |         {
468 |             if (y < 0)
469 |                 continue;
470 | 
471 |             if (y >= h)
472 |                 break;
473 | 
474 |             unsigned char* p = pixels + stride * y;
475 | 
476 |             for (int x = rx - t0; x < rx + rw + t1; x++)
477 |             {
478 |                 if (x < 0)
479 |                     continue;
480 | 
481 |                 if (x >= w)
482 |                     break;
483 | 
484 |                 p[x * 4 + 0] = pen_color[0];
485 |                 p[x * 4 + 1] = pen_color[1];
486 |                 p[x * 4 + 2] = pen_color[2];
487 |                 p[x * 4 + 3] = pen_color[3];
488 |             }
489 |         }
490 |     }
491 | 
492 |     // draw bottom
493 |     {
494 |         for (int y = ry + rh - t0; y < ry + rh + t1; y++)
495 |         {
496 |             if (y < 0)
497 |                 continue;
498 | 
499 |             if (y >= h)
500 |                 break;
501 | 
502 |             unsigned char* p = pixels + stride * y;
503 | 
504 |             for (int x = rx - t0; x < rx + rw + t1; x++)
505 |             {
506 |                 if (x < 0)
507 |                     continue;
508 | 
509 |                 if (x >= w)
510 |                     break;
511 | 
512 |                 p[x * 4 + 0] = pen_color[0];
513 |                 p[x * 4 + 1] = pen_color[1];
514 |                 p[x * 4 + 2] = pen_color[2];
515 |                 p[x * 4 + 3] = pen_color[3];
516 |             }
517 |         }
518 |     }
519 | 
520 |     // draw left
521 |     for (int x = rx - t0; x < rx + t1; x++)
522 |     {
523 |         if (x < 0)
524 |             continue;
525 | 
526 |         if (x >= w)
527 |             break;
528 | 
529 |         for (int y = ry + t1; y < ry + rh - t0; y++)
530 |         {
531 |             if (y < 0)
532 |                 continue;
533 | 
534 |             if (y >= h)
535 |                 break;
536 | 
537 |             unsigned char* p = pixels + stride * y;
538 | 
539 |             p[x * 4 + 0] = pen_color[0];
540 |             p[x * 4 + 1] = pen_color[1];
541 |             p[x * 4 + 2] = pen_color[2];
542 |             p[x * 4 + 3] = pen_color[3];
543 |         }
544 |     }
545 | 
546 |     // draw right
547 |     for (int x = rx + rw - t0; x < rx + rw + t1; x++)
548 |     {
549 |         if (x < 0)
550 |             continue;
551 | 
552 |         if (x >= w)
553 |             break;
554 | 
555 |         for (int y = ry + t1; y < ry + rh - t0; y++)
556 |         {
557 |             if (y < 0)
558 |                 continue;
559 | 
560 |             if (y >= h)
561 |                 break;
562 | 
563 |             unsigned char* p = pixels + stride * y;
564 | 
565 |             p[x * 4 + 0] = pen_color[0];
566 |             p[x * 4 + 1] = pen_color[1];
567 |             p[x * 4 + 2] = pen_color[2];
568 |             p[x * 4 + 3] = pen_color[3];
569 |         }
570 |     }
571 | }
572 | 
573 | static void draw_rectangle_yuv420sp(unsigned char* yuv420sp, int w, int h, int rx, int ry, int rw, int rh, unsigned int color, int thickness)
574 | {
575 |     // assert w % 2 == 0
576 |     // assert h % 2 == 0
577 |     // assert rx % 2 == 0
578 |     // assert ry % 2 == 0
579 |     // assert rw % 2 == 0
580 |     // assert rh % 2 == 0
581 |     // assert thickness % 2 == 0
582 | 
583 |     const unsigned char* pen_color = (const unsigned char*)&color;
584 | 
585 |     unsigned int v_y;
586 |     unsigned int v_uv;
587 |     unsigned char* pen_color_y = (unsigned char*)&v_y;
588 |     unsigned char* pen_color_uv = (unsigned char*)&v_uv;
589 |     pen_color_y[0] = pen_color[0];
590 |     pen_color_uv[0] = pen_color[1];
591 |     pen_color_uv[1] = pen_color[2];
592 | 
593 |     unsigned char* Y = yuv420sp;
594 |     draw_rectangle_c1(Y, w, h, w, rx, ry, rw, rh, v_y, thickness);
595 | 
596 |     unsigned char* UV = yuv420sp + w * h;
597 |     int thickness_uv = thickness == -1 ? thickness : max(thickness / 2, 1);
598 |     draw_rectangle_c2(UV, w / 2, h / 2, w, rx / 2, ry / 2, rw / 2, rh / 2, v_uv, thickness_uv);
599 | }
600 | 
601 | static void draw_image_yuv420sp(unsigned char* yuv420sp, int w, int h, unsigned char* draw_img, int rx, int ry, int rw, int rh) {
602 |     for (int i = 0; i < rh; i++) {
603 |         memcpy(yuv420sp + (ry+i) * w + rx,  draw_img + i * rw,  rw);
604 |     }
605 |     for (int i = 0; i < rh/2; i++) {
606 |         memcpy(yuv420sp + w*h + (ry/2+i) * w+ rx, draw_img + rw*rh + i*rw, rw);
607 |     }
608 | }
609 | 
610 | #endif //RKNN_INFER_PLUGIN_IMAGE_OP_UTILS_H
611 | 


--------------------------------------------------------------------------------
/rknn_plugins/mpp_video_decoder.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.09
  5 |  * @brief: Rockchip MPP 视频解码封装，参考 https://blog.csdn.net/viengo/article/details/121439122
  6 |  */
  7 | 
  8 | #include <string>
  9 | #include <execution>
 10 | #include "mpp_video_decoder.h"
 11 | #include "utils_log.h"
 12 | 
 13 | MppVideoDecoder::MppVideoDecoder(const std::string &video_path) {
 14 |     // 打开视频输入源
 15 |     m_in_fp = fopen(video_path.c_str(), "rb");
 16 |     CHECK_VAL(m_in_fp == nullptr, d_mpp_module_error("open %s failed!", video_path.c_str()); return;)
 17 | 
 18 |     // 初始化解码器
 19 |     int ret = init_decoder();
 20 |     CHECK_VAL(ret < 0, d_mpp_module_error("init decoder failed!"); return;)
 21 | 
 22 |     m_init_flag = true;
 23 | }
 24 | 
 25 | MppVideoDecoder::~MppVideoDecoder() {
 26 |     if(m_in_fp != nullptr){
 27 |         fclose(m_in_fp);
 28 |         m_in_buf = nullptr;
 29 |     }
 30 |     if(m_pkt){
 31 |         mpp_packet_deinit(&m_pkt);
 32 |         m_pkt = nullptr;
 33 |     }
 34 |     if(m_in_buf != nullptr){
 35 |         free(m_in_buf);
 36 |         m_in_buf = nullptr;
 37 |     }
 38 |     if(m_mpp_mpi && m_mpp_ctx){
 39 |         m_mpp_mpi->reset(m_mpp_ctx);
 40 |         m_mpp_mpi = nullptr;
 41 |     }
 42 |     if(m_mpp_ctx != nullptr){
 43 |         mpp_destroy(m_mpp_ctx);
 44 |         m_mpp_ctx = nullptr;
 45 |     }
 46 |     if(m_frame_buffer_group != nullptr){
 47 |         mpp_buffer_group_put(m_frame_buffer_group);
 48 |         m_frame_buffer_group = nullptr;
 49 |     }
 50 |     d_mpp_module_info("MppVideoDecoder release success! buffer now total: %d, unreleased frame count :%d ",
 51 |                       mpp_buffer_total_now(),
 52 |                       m_frame_count)
 53 | }
 54 | 
 55 | int MppVideoDecoder::init_decoder() {
 56 |     // 初始化解码器上下文，MppCtx MppApi
 57 |     MPP_RET ret = mpp_create(&m_mpp_ctx, &m_mpp_mpi);
 58 |     CHECK_VAL(MPP_OK != ret, d_mpp_module_error("mpp_create error"); return -1;)
 59 | 
 60 |     // 设置解码器参数
 61 |     RK_U32 need_split = 1;
 62 |     ret = m_mpp_mpi->control(m_mpp_ctx, MPP_DEC_SET_PARSER_SPLIT_MODE, (MppParam*)&need_split);
 63 |     CHECK_VAL(MPP_OK != ret, d_mpp_module_error("m_mpp_mpi->control error MPP_DEC_SET_PARSER_SPLIT_MODE"); return -1;)
 64 | 
 65 |     ret = mpp_init(m_mpp_ctx, MPP_CTX_DEC, MPP_VIDEO_CodingAVC);  // 固定为H264
 66 |     CHECK_VAL(MPP_OK != ret, d_mpp_module_error("mpp_init error"); return -1;)
 67 | 
 68 |     // 申请解码器输入缓冲区，初始化packet
 69 |     m_in_buf = (char*)malloc(MAX_READ_BUFFER_SIZE);
 70 |     CHECK_VAL(nullptr == m_in_buf, d_mpp_module_error("malloc m_in_buf error"); return -1;)
 71 |     memset(m_in_buf, 0, MAX_READ_BUFFER_SIZE);
 72 | 
 73 |     ret = mpp_packet_init(&m_pkt, m_in_buf, MAX_READ_BUFFER_SIZE);
 74 |     CHECK_VAL(MPP_OK != ret, d_mpp_module_error("mpp_packet_init error"); return -1;)
 75 |     mpp_packet_set_length(m_pkt, 0);
 76 |     return 0;
 77 | }
 78 | 
 79 | int MppVideoDecoder::get_next_frame(DecoderMppFrame &decoder_frame) {
 80 |     if(m_video_eos){
 81 |         // 已经读到最后一帧了
 82 |         d_rknn_plugin_warn("video eos!")
 83 |         return -1;
 84 |     }
 85 |     if(m_frame_count >= MAX_DECODER_FRAME_NUM){
 86 |         d_mpp_module_warn("unreleased frame count reach max: %d", MAX_DECODER_FRAME_NUM)
 87 |     }
 88 | 
 89 |     // 读取一个数据帧
 90 |     bool get_valid_frame = false;
 91 |     while (!get_valid_frame){
 92 |         MPP_RET ret = MPP_OK;
 93 |         // 读取数据包
 94 |         get_one_packet();
 95 | 
 96 |         // 设置数据包
 97 |         int pkt_len = (int)mpp_packet_get_length(m_pkt);
 98 |         if(pkt_len > 0){
 99 |             ret = m_mpp_mpi->decode_put_packet(m_mpp_ctx, m_pkt);
100 |             d_mpp_module_info("pkt send ret:%d remain:%d", ret, (int)mpp_packet_get_length(m_pkt))
101 |         }
102 | 
103 |         // 解析帧
104 |         ret = m_mpp_mpi->decode_get_frame(m_mpp_ctx, &decoder_frame.mpp_frame);
105 |         if (MPP_OK != ret || !decoder_frame.mpp_frame) {
106 |             d_mpp_module_debug("decode_get_frame failed ret:%d, frame:%p", ret, decoder_frame.mpp_frame);
107 |             // 等待一下2ms，通常1080p解码时间2ms
108 |             usleep(2000);
109 |             continue;
110 |         }
111 | 
112 |         decoder_frame.hor_stride = mpp_frame_get_hor_stride(decoder_frame.mpp_frame);
113 |         decoder_frame.ver_stride = mpp_frame_get_ver_stride(decoder_frame.mpp_frame);
114 |         decoder_frame.hor_width = mpp_frame_get_width(decoder_frame.mpp_frame);
115 |         decoder_frame.ver_height = mpp_frame_get_height(decoder_frame.mpp_frame);
116 |         decoder_frame.data_size = mpp_frame_get_buf_size(decoder_frame.mpp_frame);
117 |         d_mpp_module_debug("decoder require buffer w:h [%d:%d] stride [%d:%d] buf_size %d",
118 |                            decoder_frame.hor_width,
119 |                            decoder_frame.ver_height,
120 |                            decoder_frame.hor_stride,
121 |                            decoder_frame.ver_stride,
122 |                            decoder_frame.data_size)
123 | 
124 |         if (mpp_frame_get_info_change(decoder_frame.mpp_frame)){
125 |             d_mpp_module_warn("decode_get_frame info changed")
126 |             if(m_frame_buffer_group == nullptr){
127 |                 /* If buffer group is not set create one and limit it */
128 |                 ret = mpp_buffer_group_get_internal(&m_frame_buffer_group, MPP_BUFFER_TYPE_DRM);
129 |                 if (ret) {
130 |                     d_mpp_module_error("%p get mpp buffer group failed ret %d ", m_mpp_ctx, ret);
131 |                     mpp_frame_deinit(&decoder_frame.mpp_frame);
132 |                     return -1;
133 |                 }
134 | 
135 |                 /* Set buffer to mpp decoder */
136 |                 ret = m_mpp_mpi->control(m_mpp_ctx, MPP_DEC_SET_EXT_BUF_GROUP, m_frame_buffer_group);
137 |                 if (ret) {
138 |                     d_mpp_module_error("%p set buffer group failed ret %d ", m_mpp_ctx, ret);
139 |                     mpp_frame_deinit(&decoder_frame.mpp_frame);
140 |                     return -1;
141 |                 }
142 |             }else{
143 |                 /* If old buffer group exist clear it */
144 |                 ret = mpp_buffer_group_clear(m_frame_buffer_group);
145 |                 if (ret) {
146 |                     d_mpp_module_error("%p clear buffer group failed ret %d ", m_mpp_ctx, ret);
147 |                     mpp_frame_deinit(&decoder_frame.mpp_frame);
148 |                     return -1;
149 |                 }
150 |             }
151 | 
152 |             // Use limit config to limit buffer count
153 |             ret = mpp_buffer_group_limit_config(m_frame_buffer_group, decoder_frame.data_size, MAX_DECODER_FRAME_NUM);
154 |             if (ret) {
155 |                 d_mpp_module_error("%p limit buffer group failed ret %d ", m_mpp_ctx, ret);
156 |                 mpp_frame_deinit(&decoder_frame.mpp_frame);
157 |                 return -1;
158 |             }
159 | 
160 |             // All buffer group config done. Set info change ready to let decoder continue decoding
161 |             ret = m_mpp_mpi->control(m_mpp_ctx, MPP_DEC_SET_INFO_CHANGE_READY, nullptr);
162 |             if (ret) {
163 |                 d_mpp_module_error("%p info change ready failed ret %d ", m_mpp_ctx, ret);
164 |                 mpp_frame_deinit(&decoder_frame.mpp_frame);
165 |                 return -1;
166 |             }
167 |         }else{
168 |             d_mpp_module_debug("decode_get_frame success")
169 |             uint32_t err_info = mpp_frame_get_errinfo(decoder_frame.mpp_frame) | mpp_frame_get_discard(decoder_frame.mpp_frame);
170 |             if (err_info) {
171 |                 d_mpp_module_error("decoder_get_frame get err info:%d discard:%d",
172 |                                    mpp_frame_get_errinfo(decoder_frame.mpp_frame),
173 |                                    mpp_frame_get_discard(decoder_frame.mpp_frame))
174 |             }
175 | //            d_mpp_module_info("decoder_get_frame, format : %d", mpp_frame_get_fmt(decoder_frame.mpp_frame))
176 |             decoder_frame.mpp_frame_format = mpp_frame_get_fmt(decoder_frame.mpp_frame);
177 |             decoder_frame.data_buf = (char *) mpp_buffer_get_ptr(mpp_frame_get_buffer(decoder_frame.mpp_frame));
178 |             decoder_frame.data_fd = mpp_buffer_get_fd(mpp_frame_get_buffer(decoder_frame.mpp_frame));
179 |             get_valid_frame = true;
180 |             m_frame_count++;
181 |         }
182 | 
183 |         if (mpp_frame_get_eos(decoder_frame.mpp_frame)) {
184 |             // 最后一帧
185 |             d_mpp_module_info("mpp_frame_get_eos");
186 |             m_video_eos = true;
187 |         }
188 |     }
189 |     return 0;
190 | }
191 | 
192 | void MppVideoDecoder::release_frame(DecoderMppFrame &decoder_frame) {
193 |     m_frame_count--;
194 |     mpp_frame_deinit(&decoder_frame.mpp_frame);
195 | }
196 | 
197 | int MppVideoDecoder::get_one_packet() {
198 |     int pkt_len = (int)mpp_packet_get_length(m_pkt);
199 |     if (pkt_len > 0) {
200 |         // 数据包未使用，不需要再读取
201 |         d_mpp_module_warn("pkt get before remain:%d", pkt_len)
202 |         return 0;
203 |     }
204 |     if (feof(m_in_fp)) {
205 |         // 文件已经读完
206 |         d_mpp_module_debug("file read end")
207 |         return -1;
208 |     }
209 | 
210 |     // 读取新的数据包
211 |     uint32_t len = fread(m_in_buf, 1, MAX_READ_BUFFER_SIZE, m_in_fp);
212 |     d_mpp_module_info("read file len:%d", len)
213 |     if (len > 0) {
214 |         mpp_packet_set_data(m_pkt, m_in_buf);
215 |         mpp_packet_set_size(m_pkt, len);
216 |         mpp_packet_set_pos(m_pkt, m_in_buf);
217 |         mpp_packet_set_length(m_pkt, len);
218 |         if (feof(m_in_fp) || len < MAX_READ_BUFFER_SIZE) {
219 |             // 读到了最后一个包
220 |             mpp_packet_set_eos(m_pkt);
221 |             d_mpp_module_info("mpp_packet_set_eos")
222 |         }
223 |     }
224 |     return 0;
225 | }
226 | 


--------------------------------------------------------------------------------
/rknn_plugins/mpp_video_decoder.h:
--------------------------------------------------------------------------------
 1 | /**
 2 |  * @author: bo.liu
 3 |  * @mail: geniusrabbit@qq.com
 4 |  * @date: 2023.08.09
 5 |  * @brief: Rockchip MPP 视频解码封装
 6 |  */
 7 | 
 8 | #ifndef RKNN_INFER_PLUGIN_MPP_VIDEO_DECODER_H
 9 | #define RKNN_INFER_PLUGIN_MPP_VIDEO_DECODER_H
10 | 
11 | #include <cstdio>
12 | 
13 | #include "rk_mpi.h"
14 | #include "plugin_common.h"
15 | 
16 | #define MAX_READ_BUFFER_SIZE (5 * 1024 * 1024)
17 | #define MAX_DECODER_FRAME_NUM (200)
18 | 
19 | struct DecoderMppFrame{
20 |     uint32_t hor_stride;
21 |     uint32_t ver_stride;
22 | 
23 |     uint32_t hor_width;
24 |     uint32_t ver_height;
25 | 
26 |     int data_fd;
27 |     void *data_buf;
28 |     uint32_t data_size;
29 | 
30 |     MppFrame mpp_frame;
31 |     MppFrameFormat mpp_frame_format;
32 | };
33 | 
34 | class MppVideoDecoder {
35 | public:
36 |     // 初始化解码器
37 |     explicit MppVideoDecoder(const std::string &video_path);
38 | 
39 |     // 释放解码器
40 |     ~MppVideoDecoder();
41 | 
42 |     [[nodiscard]] bool is_init() const { return m_init_flag; };
43 | 
44 |     // 获取视频的下一帧数据
45 |     int get_next_frame(DecoderMppFrame &frame);
46 | 
47 |     void release_frame(DecoderMppFrame &frame);
48 | private:
49 |     // 初始化解码器
50 |     int init_decoder();
51 | 
52 |     // 封装一个包
53 |     int get_one_packet();
54 | 
55 | private:
56 |     bool m_init_flag = false;
57 |     uint64_t m_frame_count = 0;
58 |     // 视频输入源
59 |     FILE *m_in_fp = nullptr;
60 |     char *m_in_buf = nullptr;
61 |     MppPacket m_pkt = nullptr;
62 | 
63 |     // Mpp视频解码器上下文
64 |     MppCtx m_mpp_ctx = nullptr;
65 |     MppApi *m_mpp_mpi = nullptr;
66 | 
67 |     // 缓存
68 |     MppBufferGroup  m_frame_buffer_group = nullptr;
69 | 
70 |     // 视频解码信息
71 |     bool m_video_eos = false; // 视频解码结束标志
72 |     bool m_video_loop_decoder = true; // 视频循环解码标志
73 | };
74 | #endif //RKNN_INFER_PLUGIN_MPP_VIDEO_DECODER_H
75 | 


--------------------------------------------------------------------------------
/rknn_plugins/mpp_video_encoder.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.09
  5 |  * @brief: Rockchip MPP 视频编码封装， 来源于 https://blog.csdn.net/qq_39839546/article/details/122023991
  6 |  */
  7 | #include <cstring>
  8 | #include "mpp_video_encoder.h"
  9 | #include "mpp_video_utils.h"
 10 | #include "utils_log.h"
 11 | 
 12 | #define MPP_ALIGN(x, a)         (((x)+(a)-1)&~((a)-1))
 13 | 
 14 | MppVideoEncoder::MppVideoEncoder(const std::string &video_path,
 15 |                                  int32_t width,int32_t height,
 16 |                                  MppFrameFormat fmt,
 17 |                                  MppCodingType type,
 18 |                                  int32_t fps, int32_t gop){
 19 |     memset(&m_enc_data, 0, sizeof(m_enc_data));
 20 | 
 21 |     m_enc_data.fp_output = fopen(video_path.c_str(), "wb+");
 22 |     CHECK_VAL(m_enc_data.fp_output == nullptr, d_mpp_module_error("failed to open output file %s", video_path.c_str()); return;)
 23 | 
 24 |     //使用输入的配置初始化编码器
 25 |     if(width != 0 && height != 0){
 26 |         int ret = init_encoder(width, height, fmt, type, fps, gop);
 27 |         CHECK_VAL(ret != MPP_OK, d_mpp_module_error("init encoder failed!"); return;)
 28 |         m_mpp_init_flag = true;
 29 |         d_mpp_module_info("init encoder success!")
 30 |     }
 31 |     m_init_flag = true;
 32 | }
 33 | 
 34 | MppVideoEncoder::~MppVideoEncoder() {
 35 |     uninit_encoder();
 36 | }
 37 | 
 38 | int MppVideoEncoder::uninit_encoder(){
 39 |     if(m_enc_data.fp_output != nullptr){
 40 |         fclose(m_enc_data.fp_output);
 41 |     }
 42 |     if (m_enc_data.ctx) {
 43 |         m_enc_data.mpi->reset(m_enc_data.ctx);
 44 |         mpp_destroy(m_enc_data.ctx);
 45 |         m_enc_data.ctx = nullptr;
 46 |     }
 47 | 
 48 |     if (m_enc_data.frm_buf) {
 49 |         mpp_buffer_put(m_enc_data.frm_buf);
 50 |         m_enc_data.frm_buf = nullptr;
 51 |     }
 52 | 
 53 |     if (m_enc_data.pkt_buf) {
 54 |         mpp_buffer_put(m_enc_data.pkt_buf);
 55 |         m_enc_data.pkt_buf = nullptr;
 56 |     }
 57 | 
 58 |     if(m_enc_data.md_info){
 59 |         mpp_buffer_put(m_enc_data.md_info);
 60 |         m_enc_data.md_info = nullptr;
 61 |     }
 62 |     if (m_enc_data.buf_grp) {
 63 |         mpp_buffer_group_put(m_enc_data.buf_grp);
 64 |         m_enc_data.buf_grp = nullptr;
 65 |     }
 66 |     d_mpp_module_info("MppVideoEncoder release success! buffer now total: %d",
 67 |                       mpp_buffer_total_now())
 68 |     return 0;
 69 | }
 70 | 
 71 | int MppVideoEncoder::init_encoder(int32_t width,int32_t height,
 72 |                                   MppFrameFormat fmt,
 73 |                                   MppCodingType type,
 74 |                                   int32_t fps, int32_t gop) {
 75 |     MPP_RET ret = MPP_OK;
 76 | 
 77 |     m_enc_data.width = width;
 78 |     m_enc_data.height = height;
 79 |     m_enc_data.hor_stride = MPP_ALIGN(m_enc_data.width, 16);
 80 |     m_enc_data.ver_stride = MPP_ALIGN(m_enc_data.height, 16);
 81 |     m_enc_data.fmt = fmt;
 82 |     m_enc_data.type = type;
 83 |     m_enc_data.fps = fps;
 84 |     m_enc_data.gop = gop;
 85 | 
 86 |     m_enc_data.mdinfo_size  = (MPP_VIDEO_CodingHEVC == m_enc_data.type) ?
 87 |                          (MPP_ALIGN(m_enc_data.hor_stride, 32) >> 5) *
 88 |                          (MPP_ALIGN(m_enc_data.ver_stride, 32) >> 5) * 16 :
 89 |                          (MPP_ALIGN(m_enc_data.hor_stride, 64) >> 6) *
 90 |                          (MPP_ALIGN(m_enc_data.ver_stride, 16) >> 4) * 16;
 91 |     //不同的图像格式所占的内存大小和其长宽的关系是不同的
 92 |     //所以要根据不同的输入图像格式为编码器编码开辟不同的内存大小，
 93 |     if (m_enc_data.fmt <= MPP_FMT_YUV420SP_VU){
 94 |         m_enc_data.frame_size = m_enc_data.hor_stride * m_enc_data.ver_stride * 3/2;
 95 | 
 96 |     } else if (m_enc_data.fmt <= MPP_FMT_YUV422_UYVY) {
 97 |         m_enc_data.hor_stride *= 2;
 98 |         m_enc_data.frame_size = m_enc_data.hor_stride * m_enc_data.ver_stride;
 99 |     } else {
100 |         m_enc_data.frame_size = m_enc_data.hor_stride * m_enc_data.ver_stride * 4;
101 |     }
102 |     d_mpp_module_info("frame_size : %d", m_enc_data.frame_size)
103 | 
104 |     //开辟编码时需要的内存
105 |     ret = mpp_buffer_group_get_internal(&m_enc_data.buf_grp, MPP_BUFFER_TYPE_DRM);
106 |     if (ret) {
107 |         d_mpp_module_error("failed to get mpp buffer group ret %d", ret);
108 |         goto MPP_INIT_OUT;
109 |     }
110 |     ret = mpp_buffer_get(m_enc_data.buf_grp, &m_enc_data.frm_buf, m_enc_data.frame_size);
111 |     if (ret) {
112 |         d_mpp_module_error("failed to get buffer for input frame ret %d", ret);
113 |         goto MPP_INIT_OUT;
114 |     }
115 |     ret = mpp_buffer_get(m_enc_data.buf_grp, &m_enc_data.pkt_buf, m_enc_data.frame_size);
116 |     if (ret) {
117 |         d_mpp_module_error("failed to get buffer for output packet ret %d", ret);
118 |         goto MPP_INIT_OUT;
119 |     }
120 |     ret = mpp_buffer_get(m_enc_data.buf_grp, &m_enc_data.md_info, m_enc_data.mdinfo_size);
121 |     if (ret) {
122 |         d_mpp_module_error("failed to get buffer for motion info output packet ret %d", ret);
123 |         goto MPP_INIT_OUT;
124 |     }
125 | 
126 |     //创建 MPP context 和 MPP api 接口
127 |     ret = mpp_create(&m_enc_data.ctx, &m_enc_data.mpi);
128 |     if (ret) {
129 |         d_mpp_module_error("mpp_create failed ret %d", ret);
130 |         goto MPP_INIT_OUT;
131 |     }
132 | 
133 |     /*
134 |      * 初始化编码还是解码，以及编解码的格式
135 |      * MPP_CTX_DEC ： 解码
136 |      * MPP_CTX_ENC ： 编码
137 |      *
138 |      * MPP_VIDEO_CodingAVC ： H.264
139 |      * MPP_VIDEO_CodingHEVC :  H.265
140 |      * MPP_VIDEO_CodingVP8 :  VP8
141 |      * MPP_VIDEO_CodingVP9 :  VP9
142 |      * MPP_VIDEO_CodingMJPEG : MJPEG
143 |      */
144 |     ret = mpp_init(m_enc_data.ctx, MPP_CTX_ENC, m_enc_data.type);
145 |     if (ret) {
146 |         d_mpp_module_error("mpp_init failed ret %d", ret);
147 |         goto MPP_INIT_OUT;
148 |     }
149 | 
150 |     // 设置编码参数：宽高、对齐后宽高等参数
151 |     m_enc_data.bps = m_enc_data.width * m_enc_data.height / 8 * m_enc_data.fps;
152 |     m_enc_data.prep_cfg.change        = MPP_ENC_PREP_CFG_CHANGE_INPUT |
153 |                                         MPP_ENC_PREP_CFG_CHANGE_ROTATION |
154 |                                         MPP_ENC_PREP_CFG_CHANGE_FORMAT;
155 |     m_enc_data.prep_cfg.width         = m_enc_data.width;
156 |     m_enc_data.prep_cfg.height        = m_enc_data.height;
157 |     m_enc_data.prep_cfg.hor_stride    = m_enc_data.hor_stride;
158 |     m_enc_data.prep_cfg.ver_stride    = m_enc_data.ver_stride;
159 |     m_enc_data.prep_cfg.format        = m_enc_data.fmt;
160 |     m_enc_data.prep_cfg.rotation      = MPP_ENC_ROT_0;
161 |     ret = m_enc_data.mpi->control(m_enc_data.ctx, MPP_ENC_SET_PREP_CFG, &m_enc_data.prep_cfg);
162 |     if (ret) {
163 |         d_mpp_module_error("mpi control enc set prep cfg failed ret %d", ret);
164 |         goto MPP_INIT_OUT;
165 |     }
166 | 
167 |     /*设置编码码率、质量、定码率变码率*/
168 |     m_enc_data.rc_cfg.change  = MPP_ENC_RC_CFG_CHANGE_ALL;
169 |     m_enc_data.rc_cfg.rc_mode = MPP_ENC_RC_MODE_VBR;
170 |     m_enc_data.rc_cfg.quality = MPP_ENC_RC_QUALITY_MEDIUM;
171 |     if (m_enc_data.rc_cfg.rc_mode == MPP_ENC_RC_MODE_CBR) {
172 |         /* constant bitrate has very small bps range of 1/16 bps */
173 |         m_enc_data.rc_cfg.bps_target   = m_enc_data.bps;
174 |         m_enc_data.rc_cfg.bps_max      = m_enc_data.bps * 17 / 16;
175 |         m_enc_data.rc_cfg.bps_min      = m_enc_data.bps * 15 / 16;
176 |     } else if (m_enc_data.rc_cfg.rc_mode ==  MPP_ENC_RC_MODE_VBR) {
177 |         if (m_enc_data.rc_cfg.quality == MPP_ENC_RC_QUALITY_CQP) {
178 |             /* constant QP does not have bps */
179 |             m_enc_data.rc_cfg.bps_target   = -1;
180 |             m_enc_data.rc_cfg.bps_max      = -1;
181 |             m_enc_data.rc_cfg.bps_min      = -1;
182 |         } else {
183 |             /* variable bitrate has large bps range */
184 |             m_enc_data.rc_cfg.bps_target   = m_enc_data.bps;
185 |             m_enc_data.rc_cfg.bps_max      = m_enc_data.bps * 17 / 16;
186 |             m_enc_data.rc_cfg.bps_min      = m_enc_data.bps * 1 / 16;
187 |         }
188 |     }
189 |     /* fix input / output frame rate */
190 |     m_enc_data.rc_cfg.fps_in_flex      = 0;
191 |     m_enc_data.rc_cfg.fps_in_num       = m_enc_data.fps;
192 |     m_enc_data.rc_cfg.fps_in_denorm    = 1;
193 |     m_enc_data.rc_cfg.fps_out_flex     = 0;
194 |     m_enc_data.rc_cfg.fps_out_num      = m_enc_data.fps;
195 |     m_enc_data.rc_cfg.fps_out_denorm   = 1;
196 | 
197 |     m_enc_data.rc_cfg.gop              = m_enc_data.gop;
198 |     m_enc_data.rc_cfg.skip_cnt         = 0;
199 | 
200 |     ret = m_enc_data.mpi->control(m_enc_data.ctx, MPP_ENC_SET_RC_CFG, &m_enc_data.rc_cfg);
201 |     if (ret) {
202 |         d_mpp_module_error("mpi control enc set rc cfg failed ret %d", ret);
203 |         goto MPP_INIT_OUT;
204 |     }
205 | 
206 |     /*设置264相关的其他编码参数*/
207 |     m_enc_data.codec_cfg.coding = m_enc_data.type;
208 |     switch (m_enc_data.codec_cfg.coding) {
209 |         case MPP_VIDEO_CodingAVC : {
210 |             m_enc_data.codec_cfg.h264.change = MPP_ENC_H264_CFG_CHANGE_PROFILE |
211 |                                                  MPP_ENC_H264_CFG_CHANGE_ENTROPY |
212 |                                                  MPP_ENC_H264_CFG_CHANGE_TRANS_8x8;
213 |             /*
214 |              * H.264 profile_idc parameter
215 |              * 66  - Baseline profile
216 |              * 77  - Main profile
217 |              * 100 - High profile
218 |              */
219 |             m_enc_data.codec_cfg.h264.profile  = 77;
220 |             /*
221 |              * H.264 level_idc parameter
222 |              * 10 / 11 / 12 / 13    - qcif@15fps / cif@7.5fps / cif@15fps / cif@30fps
223 |              * 20 / 21 / 22         - cif@30fps / half-D1@@25fps / D1@12.5fps
224 |              * 30 / 31 / 32         - D1@25fps / 720p@30fps / 720p@60fps
225 |              * 40 / 41 / 42         - 1080p@30fps / 1080p@30fps / 1080p@60fps
226 |              * 50 / 51 / 52         - 4K@30fps
227 |              */
228 |             m_enc_data.codec_cfg.h264.level    = 40;
229 |             m_enc_data.codec_cfg.h264.entropy_coding_mode  = 1;
230 |             m_enc_data.codec_cfg.h264.cabac_init_idc  = 0;
231 |             m_enc_data.codec_cfg.h264.transform8x8_mode = 1;
232 |         }
233 |             break;
234 |         case MPP_VIDEO_CodingMJPEG : {
235 |             m_enc_data.codec_cfg.jpeg.change  = MPP_ENC_JPEG_CFG_CHANGE_QP;
236 |             m_enc_data.codec_cfg.jpeg.quant   = 10;
237 |         }
238 |             break;
239 |         case MPP_VIDEO_CodingVP8 :
240 |         case MPP_VIDEO_CodingHEVC :
241 |         default : {
242 |             d_mpp_module_info("support encoder coding type %d", m_enc_data.codec_cfg.coding);
243 |         }
244 |             break;
245 |     }
246 |     ret = m_enc_data.mpi->control(m_enc_data.ctx, MPP_ENC_SET_CODEC_CFG, &m_enc_data.codec_cfg);
247 |     if (ret) {
248 |         d_mpp_module_error("mpi control enc set codec cfg failed ret %d", ret);
249 |         goto MPP_INIT_OUT;
250 |     }
251 | 
252 |     /* optional */
253 |     m_enc_data.sei_mode = MPP_ENC_SEI_MODE_ONE_FRAME;
254 |     ret = m_enc_data.mpi->control(m_enc_data.ctx, MPP_ENC_SET_SEI_CFG, &m_enc_data.sei_mode);
255 |     if (ret) {
256 |         d_mpp_module_error("mpi control enc set sei cfg failed ret %d", ret);
257 |         goto MPP_INIT_OUT;
258 |     }
259 | 
260 |     if (m_enc_data.type == MPP_VIDEO_CodingAVC) {
261 |         MppPacket packet = nullptr;
262 |         /*
263 |          * Can use packet with normal malloc buffer as input not pkt_buf.
264 |          * Please refer to vpu_api_legacy.cpp for normal buffer case.
265 |          * Using pkt_buf buffer here is just for simplifing demo.
266 |          */
267 |         mpp_packet_init_with_buffer(&packet, m_enc_data.pkt_buf);
268 |         /* NOTE: It is important to clear output packet length!! */
269 |         mpp_packet_set_length(packet, 0);
270 |         ret = m_enc_data.mpi->control(m_enc_data.ctx, MPP_ENC_GET_EXTRA_INFO, &packet);
271 |         if (MPP_OK != ret) {
272 |             d_mpp_module_error("mpi control enc get extra info failed, ret : %d", ret);
273 |             goto MPP_INIT_OUT;
274 |         }
275 | 
276 |         /* get and write sps/pps for H.264 */
277 |         d_mpp_module_info("mpi control enc get extra info successful, packet : %p", packet)
278 |         if (packet) {
279 |             void *ptr = mpp_packet_get_pos(packet);
280 |             size_t len = mpp_packet_get_length(packet);
281 |             if (m_enc_data.fp_output) {
282 |                 d_mpp_module_info("write extra data %d bytes", len)
283 |                 uint32_t w_len = fwrite(ptr, 1, len, m_enc_data.fp_output);
284 |                 if (w_len != len){
285 |                     d_mpp_module_error("failed to save extra data! w_len %d len %d", w_len, len);
286 |                     goto MPP_INIT_OUT;
287 |                 }
288 |             }
289 |             mpp_packet_deinit(&packet);
290 |         }
291 |     }
292 |     return 0;
293 | 
294 | MPP_INIT_OUT:
295 |     uninit_encoder();
296 |     d_mpp_module_error("init mpp failed!");
297 |     return -1;
298 | }
299 | 
300 | /****************************************************************************
301 | MppPacket  ：   存放编码数据，例如264、265数据
302 | MppFrame  ：   存放解码的数据，例如YUV、RGB数据
303 | MppTask   :     一次编码或者解码的session
304 | 
305 | 编码就是push MppFrame，输出MppPacket；
306 | 解码就是push MppPacket，输出MppFrame；
307 | 
308 | MPI包含两套接口做编解码：
309 | 一套是简易接口， 类似 decode_put_packet / decode_get_frame 这样put/get即可
310 | 一套是高级接口， 类似 poll / enqueue/ dequeue 这样的对input output队列进行操作
311 | *****************************************************************************/
312 | bool MppVideoEncoder::process_image(uint8_t *image_data,
313 |                                     int32_t width,int32_t height,
314 |                                     MppFrameFormat fmt,
315 |                                     MppCodingType type,
316 |                                     int32_t fps, int32_t gop) {
317 |     // 初次处理图像时初始化编码器
318 |     if(!m_mpp_init_flag && width != 0) {
319 |         int ret = init_encoder(width, height, fmt, type, fps, gop);
320 |         CHECK_VAL(ret != MPP_OK, d_mpp_module_error("init encoder failed!"); return false;)
321 | 
322 |         m_mpp_init_flag = true;
323 |         d_mpp_module_info("init encoder success!")
324 |     }
325 | 
326 |     MPP_RET ret = MPP_OK;
327 |     MppFrame frame = nullptr;
328 |     MppPacket packet = nullptr;
329 |     //获得开辟的内存的首地址
330 |     void *buf = mpp_buffer_get_ptr(m_enc_data.frm_buf);
331 | 
332 |     // TODO: improve performance here?
333 |     // 从输入图像的首地址开始读取图像数据，但是读取时会考虑16位对齐，即读取的长和宽都是16的整数倍。
334 |     // 若图像一行或者一列不满16整数倍，则会用空数据补齐行和列
335 |     yuv_add_stride(image_data,
336 |                     m_enc_data.width, m_enc_data.height,
337 |                     m_enc_data.hor_stride, m_enc_data.ver_stride,
338 |                     m_enc_data.fmt,
339 |                     (uint8_t *)buf);
340 |     ret = mpp_frame_init(&frame);
341 |     if (ret) {
342 |         d_mpp_module_error("mpp_frame_init failed\n");
343 |         return true;
344 |     }
345 |     //设置编码图像的格式
346 |     mpp_frame_set_width(frame, m_enc_data.width);
347 |     mpp_frame_set_height(frame, m_enc_data.height);
348 |     mpp_frame_set_hor_stride(frame, m_enc_data.hor_stride);
349 |     mpp_frame_set_ver_stride(frame, m_enc_data.ver_stride);
350 |     mpp_frame_set_fmt(frame, m_enc_data.fmt);
351 |     mpp_frame_set_buffer(frame, m_enc_data.frm_buf);
352 | //    mpp_frame_set_buf_size(frame, m_enc_data.buf_size);
353 |     mpp_frame_set_eos(frame, m_enc_data.frm_eos);
354 | 
355 |     MppMeta meta = mpp_frame_get_meta(frame);
356 |     mpp_meta_set_packet(meta, KEY_OUTPUT_PACKET, packet);
357 |     mpp_meta_set_buffer(meta, KEY_MOTION_INFO, m_enc_data.md_info);
358 | 
359 |     mpp_packet_init_with_buffer(&packet, m_enc_data.pkt_buf);
360 |     /* NOTE: It is important to clear output packet length!! */
361 |     mpp_packet_set_length(packet, 0);
362 |     //输入图像进行编码
363 |     ret = m_enc_data.mpi->encode_put_frame(m_enc_data.ctx, frame);
364 |     mpp_frame_deinit(&frame);
365 |     if (ret) {
366 |         d_mpp_module_error("mpp encode put frame failed")
367 |         return false;
368 |     }
369 |     //获得编码后的packet
370 |     ret = m_enc_data.mpi->encode_get_packet(m_enc_data.ctx, &packet);
371 |     if (ret) {
372 |         d_mpp_module_error("mpp encode get packet failed")
373 |         return false;
374 |     }
375 |     //获得编码后的数据长度和首地址，将其写入文件
376 |     if (packet) {
377 |         // write packet to file here
378 |         void *ptr   = mpp_packet_get_pos(packet);
379 |         size_t len  = mpp_packet_get_length(packet);
380 |         m_enc_data.pkt_eos = mpp_packet_get_eos(packet);
381 | 
382 |         if (m_enc_data.fp_output) {
383 |             fwrite(ptr, 1, len, m_enc_data.fp_output);
384 |         }
385 |         mpp_packet_deinit(&packet);
386 | 
387 |         m_enc_data.stream_size += len;
388 |         m_enc_data.frame_count++;
389 | 
390 |         if (m_enc_data.pkt_eos) {
391 |             d_mpp_module_warn("found last packet")
392 |         }
393 |     }
394 | 
395 |     if (m_enc_data.num_frames && m_enc_data.frame_count >= m_enc_data.num_frames) {
396 |         d_mpp_module_error("encode max %d frames", m_enc_data.frame_count)
397 |         return false;
398 |     }
399 |     if (m_enc_data.frm_eos && m_enc_data.pkt_eos) {
400 |         d_mpp_module_error("encode frm_eos  %d pkt_eos %d ", m_enc_data.frm_eos, m_enc_data.pkt_eos)
401 |         return false;
402 |     }
403 |     return true;
404 | }
405 | 


--------------------------------------------------------------------------------
/rknn_plugins/mpp_video_encoder.h:
--------------------------------------------------------------------------------
 1 | /**
 2 |  * @author: bo.liu
 3 |  * @mail: geniusrabbit@qq.com
 4 |  * @date: 2023.08.09
 5 |  * @brief: Rockchip MPP 视频编码封装
 6 |  */
 7 | 
 8 | #ifndef RKNN_INFER_PLUGIN_MPP_VIDEO_ENCODER_H
 9 | #define RKNN_INFER_PLUGIN_MPP_VIDEO_ENCODER_H
10 | 
11 | #include <string>
12 | 
13 | #include "rk_mpi.h"
14 | 
15 | class MppVideoEncoder {
16 | public:
17 |     explicit MppVideoEncoder(
18 |             const std::string &video_path,
19 |             int32_t width = 0,int32_t height = 0,
20 |             MppFrameFormat fmt = MPP_FMT_YUV422_YUYV,
21 |             MppCodingType type = MPP_VIDEO_CodingAVC,
22 |             int32_t fps = 30, int32_t gop = 60);
23 |     ~MppVideoEncoder();
24 | 
25 |     bool process_image(uint8_t *image_data, int32_t width = 0, int32_t height = 0,
26 |                        MppFrameFormat fmt = MPP_FMT_YUV422_YUYV,
27 |                        MppCodingType type = MPP_VIDEO_CodingAVC,
28 |                        int32_t fps = 30, int32_t gop = 60);
29 | 
30 |     [[nodiscard]] bool is_init() const { return m_init_flag; };
31 | 
32 | private:
33 |     int init_encoder(int32_t width,int32_t height,
34 |                      MppFrameFormat fmt,
35 |                      MppCodingType type,
36 |                      int32_t fps, int32_t gop);
37 |     int uninit_encoder();
38 | private:
39 |     bool m_init_flag = false;
40 |     bool m_mpp_init_flag = false;
41 |     //编码所需要的数据
42 |     struct MPP_ENC_DATA {
43 |         // global flow control flag
44 |         uint32_t frm_eos;
45 |         uint32_t pkt_eos;
46 |         uint32_t frame_count;
47 |         uint64_t stream_size;
48 | 
49 |         // base flow context
50 |         MppCtx ctx;
51 |         MppApi *mpi;
52 |         MppEncPrepCfg prep_cfg;
53 |         MppEncRcCfg rc_cfg;
54 |         MppEncCodecCfg codec_cfg;
55 | 
56 |         // input / output
57 |         MppBuffer frm_buf = nullptr;
58 |         MppBuffer pkt_buf = nullptr;
59 |         MppBuffer md_info = nullptr;
60 |         MppBufferGroup buf_grp = nullptr;
61 |         MppEncSeiMode sei_mode;
62 | 
63 |         int32_t width;
64 |         int32_t height;
65 |         int32_t hor_stride;
66 |         int32_t ver_stride;
67 |         MppFrameFormat fmt = MPP_FMT_YUV422_YUYV;
68 |         MppCodingType type = MPP_VIDEO_CodingAVC;
69 |         uint32_t num_frames;
70 | 
71 |         // resources
72 |         size_t frame_size;
73 |         size_t mdinfo_size;
74 | 
75 |         int32_t gop = 60;
76 |         int32_t fps = 30;
77 |         int32_t bps;
78 | 
79 |         FILE *fp_output;
80 |     };
81 |     MPP_ENC_DATA m_enc_data{};
82 | };
83 | 
84 | #endif //RKNN_INFER_PLUGIN_MPP_VIDEO_ENCODER_H
85 | 


--------------------------------------------------------------------------------
/rknn_plugins/mpp_video_utils.h:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.09
  5 |  * @brief: Rockchip MPP 视频编码解码函数封装
  6 |  */
  7 | #ifndef RKNN_INFER_PLUGIN__MPP_VIDEO_UTILS_H
  8 | #define RKNN_INFER_PLUGIN__MPP_VIDEO_UTILS_H
  9 | #include "rk_mpi.h"
 10 | #include "opencv2/opencv.hpp"
 11 | #include "utils_log.h"
 12 | 
 13 | /**
 14 |  * @brief 将 MPP 帧数据中的 padding 部分去掉得到原始图像，参考 https://github.com/MUZLATAN/ffmpeg_rtsp_mpp/blob/master/MppDecode.cpp#L214
 15 |  */
 16 | static int yuv_del_stride(
 17 |         const uint8_t *data_buf,
 18 |         uint32_t width, uint32_t height,
 19 |         uint32_t h_stride, uint32_t v_stride,
 20 |         MppFrameFormat fmt,
 21 |         uint8_t *image) {
 22 |     switch (fmt) {
 23 |         case MPP_FMT_YUV420SP : {
 24 |             const uint8_t *base_y = data_buf;
 25 |             const uint8_t *base_c = data_buf + h_stride * v_stride;
 26 |             for (uint32_t row = 0; row < height; row++) {
 27 |                 memcpy(image, base_y + row * h_stride, width);
 28 |                 image += width;
 29 |             }
 30 |             for (uint32_t row = 0; row < height / 2; row++) {
 31 |                 memcpy(image, base_c + row * h_stride, width);
 32 |                 image += width;
 33 |             }
 34 |         } break;
 35 |         case MPP_FMT_YUV420P : {
 36 |             const uint8_t *base_y = data_buf;
 37 |             const uint8_t *base_u = base_y + h_stride * v_stride;
 38 |             const uint8_t *base_v = base_u + h_stride * v_stride / 4; // NOTE: diff from gen_yuv_image
 39 |             for (uint32_t row = 0; row < height; row++) {
 40 |                 memcpy(image, base_y + row * h_stride, width);
 41 |                 image += width;
 42 |             }
 43 |             for (uint32_t row = 0; row < height / 2; row++) {
 44 |                 memcpy(image, base_u + row * h_stride / 2, width / 2);
 45 |                 image += width / 2;
 46 |             }
 47 |             for (uint32_t row = 0; row < height / 2; row++) {
 48 |                 memcpy(image, base_v + row * h_stride / 2, width / 2);
 49 |                 image += width / 2;
 50 |             }
 51 |         } break;
 52 |         case MPP_FMT_ARGB8888 : {
 53 |             const uint8_t *base_y = data_buf;
 54 |             for (uint32_t row = 0; row < height; row++) {
 55 |                 memcpy(image, base_y + row * h_stride * 4, width * 4);
 56 |                 image += width * 4;
 57 |             }
 58 |         } break;
 59 |         case MPP_FMT_YUV422_YUYV :
 60 |         case MPP_FMT_YUV422_UYVY : {
 61 |             const uint8_t *base_y = data_buf;
 62 |             for (uint32_t row = 0; row < height; row++) {
 63 |                 memcpy(image, base_y + row * h_stride, width * 2);
 64 |                 image += width * 2;
 65 |             }
 66 |         } break;
 67 |         default : {
 68 |             d_mpp_module_error("read image do not support fmt %d", fmt)
 69 |             return -1;
 70 |         }
 71 |     }
 72 |     return 0;
 73 | }
 74 | 
 75 | /**
 76 |  * @brief 给原始图像添加 padding，封装成 MPP 帧数据，来源于 https://blog.csdn.net/qq_39839546/article/details/122023991
 77 |  */
 78 | static int yuv_add_stride(
 79 |         const uint8_t *image,
 80 |         uint32_t width, uint32_t height,
 81 |         uint32_t hor_stride, uint32_t ver_stride,
 82 |         MppFrameFormat fmt,
 83 |         uint8_t *buf) {
 84 |     // YUV格式的图像都是将YUV三个通道分为三部分存取，YYYYYY*****UUUUU*****VVVVV,所以在将其按照16位对齐copy时先将YUV三个通道的起始地址指定好。
 85 |     uint8_t *buf_y = buf;
 86 |     uint8_t *buf_u = buf_y + hor_stride * ver_stride; // NOTE: diff from gen_yuv_image
 87 |     uint8_t *buf_v = buf_u + hor_stride * ver_stride / 4; // NOTE: diff from gen_yuv_image
 88 | 
 89 |     // 然后按照不同的格式，按照16位对齐copy图像数据到buf下，不同格式读取方式不同。
 90 |     switch (fmt) {
 91 |         case MPP_FMT_YUV420SP : {
 92 |             for (uint32_t row = 0; row < height; row++) {
 93 |                 memcpy(buf_y + row * hor_stride, image,width);
 94 |                 image += width;
 95 |             }
 96 |             for (uint32_t row = 0; row < height / 2; row++) {
 97 |                 memcpy(buf_u + row * hor_stride, image, width);
 98 |                 image += width;
 99 |             }
100 |         } break;
101 |         case MPP_FMT_YUV420P : {
102 |             for (uint32_t row = 0; row < height; row++) {
103 |                 memcpy(buf_y + row * hor_stride, image, width);
104 |                 image += width;
105 |             }
106 |             for (uint32_t row = 0; row < height / 2; row++) {
107 |                 memcpy(buf_u + row * hor_stride/2, image, width/2);
108 |                 image += width/2;
109 |             }
110 |             for (uint32_t row = 0; row < height / 2; row++) {
111 |                 memcpy(buf_v + row * hor_stride/2, image, width/2);
112 |                 image += width/2;
113 |             }
114 |         } break;
115 |         case MPP_FMT_ARGB8888 : {
116 |             for (uint32_t row = 0; row < height; row++) {
117 |                 memcpy(buf_y + row * hor_stride*4, image, width*4);
118 |                 image += width*4;
119 |             }
120 |         } break;
121 |         case MPP_FMT_YUV422_YUYV :
122 |         case MPP_FMT_YUV422_UYVY : {
123 |             for (uint32_t row = 0; row < height; row++) {
124 |                 memcpy(buf_y + row * hor_stride, image, width*2);
125 |                 image += width*2;
126 |             }
127 |         } break;
128 |         default : {
129 |             d_mpp_module_error("read image do not support fmt %d", fmt)
130 |             return -1;
131 |         } break;
132 |     }
133 |     return 0;
134 | }
135 | 
136 | /**
137 |  * @brief 将 MPP 帧数据转换为 OpenCV Mat
138 |  */
139 | static int frame_data_to_mat(
140 |         const uint8_t *data_buf,
141 |         uint32_t width, uint32_t height,
142 |         uint32_t hor_stride, uint32_t ver_stride,
143 |         MppFrameFormat fmt,
144 |         cv::Mat &rgb_img) {
145 |     cv::Mat yuv_img;
146 |     switch (fmt) {
147 |         case MPP_FMT_YUV420SP :{
148 |             yuv_img.create(height * 3 / 2, width, CV_8UC1);
149 |             yuv_del_stride(data_buf, width, height, hor_stride, ver_stride, fmt, yuv_img.data);
150 |             cv::cvtColor(yuv_img, rgb_img, cv::COLOR_YUV420sp2RGB);
151 |         } break;
152 |         case MPP_FMT_YUV420P : {
153 |             yuv_img.create(height * 3 / 2, width, CV_8UC1);
154 |             yuv_del_stride(data_buf, width, height, hor_stride, ver_stride, fmt, yuv_img.data);
155 |             cv::cvtColor(yuv_img, rgb_img, cv::COLOR_YUV420p2RGB);
156 |         } break;
157 | //        case MPP_FMT_ARGB8888 : {
158 | //            yuv_img.create(height, width * 4, CV_8UC1);
159 | //            yuv_del_stride(data_buf, width, height, hor_stride, ver_stride, fmt, yuv_img.data);
160 | //            cv::cvtColor(yuv_img, rgb_img, cv::COLOR_YUV420p2RGB);
161 | //        } break;
162 | //        case MPP_FMT_YUV422_YUYV :
163 | //        case MPP_FMT_YUV422_UYVY : {
164 | //            yuv_img.create(height, width * 2, CV_8UC1);
165 | //            yuv_del_stride(data_buf, width, height, hor_stride, ver_stride, fmt, yuv_img.data);
166 | //            cv::cvtColor(yuv_img, rgb_img, cv::COLOR_YUV2RGBA_Y422);
167 | //        } break;
168 |         default : {
169 |             d_mpp_module_error("read image do not support fmt %d", fmt)
170 |             return -1;
171 |         }
172 |     }
173 |     return 0;
174 | }
175 | #endif //RKNN_INFER_PLUGIN__MPP_VIDEO_UTILS_H
176 | 


--------------------------------------------------------------------------------
/rknn_plugins/plugin_common.h:
--------------------------------------------------------------------------------
 1 | /**
 2 |  * @author: bo.liu
 3 |  * @mail: geniusrabbit@qq.com
 4 |  * @date: 2023.08.06
 5 |  * @brief: 插件公用函数
 6 |  */
 7 | #ifndef RKNN_INFER_PLUGIN_COMMON_H
 8 | #define RKNN_INFER_PLUGIN_COMMON_H
 9 | 
10 | #include <cstring>
11 | 
12 | static int rknn_plugin_get_top (
13 |         const float* p_prob,
14 |         float* p_max_prob,
15 |         uint32_t* p_max_class,
16 |         uint32_t output_count,
17 |         uint32_t top_num){
18 |     uint32_t i, j;
19 |     memset(p_max_prob, 0, sizeof(float) * top_num);
20 |     memset(p_max_class, 0xff, sizeof(float) * top_num);
21 | 
22 |     for (j = 0; j < top_num; j++) {
23 |         for (i = 0; i < output_count; i++) {
24 |             if ((i == *(p_max_class + 0)) || (i == *(p_max_class + 1)) || (i == *(p_max_class + 2)) || (i == *(p_max_class + 3)) ||
25 |                 (i == *(p_max_class + 4))) {
26 |                 continue;
27 |             }
28 | 
29 |             if (p_prob[i] > *(p_max_prob + j)) {
30 |                 *(p_max_prob + j) = p_prob[i];
31 |                 *(p_max_class + j) = i;
32 |             }
33 |         }
34 |     }
35 |     return 1;
36 | }
37 | 
38 | 
39 | #endif //RKNN_INFER_PLUGIN_COMMON_H
40 | 


--------------------------------------------------------------------------------
/rknn_plugins/rknn_mobilenet/rknn_mobilenet.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.03
  5 |  * @brief: RKNN 图像推理示例插件，参考：https://github.com/rockchip-linux/rknpu2/tree/master/examples/rknn_mobilenet_demo 实现
  6 |  */
  7 | 
  8 | /* 包含定义插件必须的头文件 */
  9 | #include <string>
 10 | #include <cstring>
 11 | 
 12 | #include "opencv2/core/core.hpp"
 13 | #include "opencv2/imgcodecs.hpp"
 14 | #include "opencv2/imgproc.hpp"
 15 | 
 16 | #include "rknn_infer_api.h"
 17 | #include "plugin_common.h"
 18 | #include "utils_log.h"
 19 | 
 20 | // 插件全局配置信息，由调度程序给插件传来的信息
 21 | PluginConfigSet g_plugin_config_set;
 22 | 
 23 | // 插件私有变量定义
 24 | struct PluginInputData {
 25 |     // 输入源
 26 |     std::string image_path;
 27 | };
 28 | 
 29 | const int MODEL_IN_WIDTH    = 224;
 30 | const int MODEL_IN_HEIGHT   = 224;
 31 | const int MODEL_IN_CHANNELS = 3;
 32 | 
 33 | static int get_config(PluginConfigGet *plugin_config){
 34 |     // 输入线程个数
 35 |     plugin_config->input_thread_nums = 2;
 36 |     // 输出线程个数
 37 |     plugin_config->output_thread_nums = 2;
 38 |     // 是否需要输出float类型的输出结果
 39 |     plugin_config->output_want_float = true;
 40 |     return 0;
 41 | }
 42 | 
 43 | static int set_config(PluginConfigSet *plugin_config){
 44 |     // 注意拷贝构造函数
 45 |     memcpy(&g_plugin_config_set, plugin_config, sizeof(PluginConfigSet));
 46 |     d_rknn_plugin_info("plugin config set success")
 47 |     return 0;
 48 | }
 49 | 
 50 | static int rknn_plugin_init(struct ThreadData *td) {
 51 |     if(td->thread_type == THREAD_TYPE_INPUT) {
 52 |         // 数据读取线程初始化
 53 |         td->plugin_private_data = new PluginInputData();
 54 |         auto *pri_data = (PluginInputData *)td->plugin_private_data;
 55 |         if(td->thread_id == 0){
 56 |             pri_data->image_path = "model/dog_224x224.jpg";
 57 |         }else if(td->thread_id == 1) {
 58 |             pri_data->image_path = "model/cat_224x224.jpg";
 59 |         }
 60 |     }
 61 |     return 0;
 62 | }
 63 | 
 64 | static int rknn_plugin_uninit(struct ThreadData *td) {
 65 |     if(td->thread_type == THREAD_TYPE_INPUT) {
 66 |         // 数据读取线程反初始化
 67 |         auto *pri_data = (PluginInputData *)td->plugin_private_data;
 68 |         delete pri_data;
 69 |     }
 70 |     return 0;
 71 | }
 72 | 
 73 | static int rknn_plugin_input(struct ThreadData *td, struct InputUnit *input_unit) {
 74 |     // 根据td来收集输入源（多线程收集）
 75 |     d_rknn_plugin_debug("plugin read input data")
 76 |     auto *pri_data = (PluginInputData *)td->plugin_private_data;
 77 | 
 78 |     // Load image
 79 |     cv::Mat orig_img = imread(pri_data->image_path, cv::IMREAD_COLOR);
 80 |     if (!orig_img.data) {
 81 |         d_rknn_plugin_error("cv::imread %s fail!", pri_data->image_path.c_str());
 82 |         return -1;
 83 |     }
 84 | 
 85 |     cv::Mat orig_img_rgb;
 86 |     cv::cvtColor(orig_img, orig_img_rgb, cv::COLOR_BGR2RGB);
 87 | 
 88 |     cv::Mat img = orig_img_rgb.clone();
 89 |     if (orig_img.cols != MODEL_IN_WIDTH || orig_img.rows != MODEL_IN_HEIGHT) {
 90 |         d_rknn_plugin_warn("resize %d %d to %d %d", orig_img.cols, orig_img.rows, MODEL_IN_WIDTH, MODEL_IN_HEIGHT);
 91 |         cv::resize(orig_img, img, cv::Size(MODEL_IN_WIDTH, MODEL_IN_HEIGHT), 0, 0, cv::INTER_LINEAR);
 92 |     }
 93 | 
 94 |     input_unit->n_inputs = g_plugin_config_set.io_num.n_input;
 95 |     input_unit->inputs = (rknn_input*)malloc(input_unit->n_inputs * sizeof(rknn_input));
 96 |     memset(input_unit->inputs, 0, input_unit->n_inputs * sizeof(rknn_input));
 97 | 
 98 |     input_unit->inputs[0].index = 0;
 99 |     input_unit->inputs[0].type  = RKNN_TENSOR_UINT8;
100 |     input_unit->inputs[0].size  = img.cols * img.rows * img.channels() * sizeof(uint8_t);
101 |     input_unit->inputs[0].fmt   = RKNN_TENSOR_NHWC;
102 |     input_unit->inputs[0].buf = new uint8_t[img.cols * img.rows * img.channels()];
103 |     memcpy(input_unit->inputs[0].buf, img.data, input_unit->inputs[0].size);
104 |     return 0;
105 | }
106 | 
107 | static int rknn_plugin_input_release(struct ThreadData *td, struct InputUnit *input_unit) {
108 |     // 释放输入源
109 |     for(uint32_t idx = 0; idx < input_unit->n_inputs; idx++){
110 |         delete[] (uint8_t*)input_unit->inputs[0].buf;
111 |         input_unit->inputs[0].buf = nullptr;
112 |     }
113 |     free(input_unit->inputs);
114 |     return 0;
115 | }
116 | 
117 | static int rknn_plugin_output(struct ThreadData *td, struct OutputUnit *output_unit) {
118 |     d_rknn_plugin_info("plugin print output data, thread_id: %d", td->thread_id)
119 |     // 收集测试结果
120 |     for (int i = 0; i < output_unit->n_outputs; i++) {
121 |         uint32_t max_class[5];
122 |         float max_prob[5];
123 |         rknn_plugin_get_top(
124 |                 (float *) output_unit->outputs[i].buf,
125 |                 max_prob,
126 |                 max_class,
127 |                 output_unit->outputs[i].size / 4,
128 |                 5);
129 | 
130 |         d_rknn_plugin_info(" --- Top5 ---");
131 |         for (int j = 0; j < 5; j++) {
132 |             d_rknn_plugin_info("%3d: %8.6f", max_class[j], max_prob[j]);
133 |         }
134 |     }
135 |     return 0;
136 | }
137 | 
138 | // 注册所有本插件的相关函数到插件结构体中
139 | // 注意：插件名称和结构体定义名称必须和插件动态库名称一致
140 | static struct PluginStruct rknn_mobilenet = {
141 |         .plugin_name 		= "rknn_mobilenet",
142 |         .plugin_version 	= 1,
143 |         .get_config         = get_config,
144 |         .set_config         = set_config,
145 |         .init				= rknn_plugin_init,
146 |         .uninit 			= rknn_plugin_uninit,
147 |         .rknn_input 		= rknn_plugin_input,
148 |         .rknn_input_release = rknn_plugin_input_release,
149 |         .rknn_output		= rknn_plugin_output,
150 | };
151 | 
152 | // 插件动态库在加载时会自动调用该函数
153 | static void plugin_init plugin_auto_register(){
154 |     d_rknn_plugin_info("auto register plugin %p, name: %s", &rknn_mobilenet, rknn_mobilenet.plugin_name)
155 |     plugin_register(&rknn_mobilenet);
156 | }
157 | 
158 | // 插件动态库在关闭时会自动调用该函数
159 | static void plugin_exit plugin_auto_unregister(){
160 |     d_rknn_plugin_info("auto unregister plugin %p, name: %s", &rknn_mobilenet, rknn_mobilenet.plugin_name)
161 |     plugin_unregister(&rknn_mobilenet);
162 | }
163 | 


--------------------------------------------------------------------------------
/rknn_plugins/rknn_plugin_template/rknn_plugin_template.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.03
  5 |  * @brief: RKNN 推理插件模板
  6 |  */
  7 | 
  8 | /* 包含定义插件必须的头文件 */
  9 | #include <string>
 10 | #include <cstring>
 11 | 
 12 | #include "rknn_infer_api.h"
 13 | #include "plugin_common.h"
 14 | #include "utils_log.h"
 15 | 
 16 | // 插件全局配置信息，由调度程序给插件传来的信息
 17 | PluginConfigSet g_plugin_config_set;
 18 | 
 19 | // 输入线程私有数据
 20 | struct PluginInputData {
 21 |     // 每个线程定制输入源
 22 | };
 23 | 
 24 | // 输出线程私有数据
 25 | struct PluginOutputData {
 26 |     // 每个线程定制输出
 27 | };
 28 | 
 29 | // 插件输入输出线程同步数据
 30 | struct PluginSyncData {
 31 | };
 32 | 
 33 | static int get_config(PluginConfigGet *plugin_config){
 34 |     // 输入线程个数
 35 |     plugin_config->input_thread_nums = 2;
 36 |     // 输出线程个数
 37 |     plugin_config->output_thread_nums = 2;
 38 |     // 是否需要输出float类型的输出结果
 39 |     plugin_config->output_want_float = true;
 40 |     return 0;
 41 | }
 42 | 
 43 | static int set_config(PluginConfigSet *plugin_config){
 44 |     // 注意拷贝构造函数
 45 |     memcpy(&g_plugin_config_set, plugin_config, sizeof(PluginConfigSet));
 46 |     d_rknn_plugin_info("plugin config set success")
 47 |     return 0;
 48 | }
 49 | 
 50 | static int rknn_plugin_init(struct ThreadData *td) {
 51 |     // 设置输入线程的输入源
 52 |     return 0;
 53 | }
 54 | 
 55 | static int rknn_plugin_uninit(struct ThreadData *td) {
 56 |     // 释放输入线程的输入源
 57 |     return 0;
 58 | }
 59 | 
 60 | static int rknn_plugin_input(struct ThreadData *td, struct InputUnit *input_unit) {
 61 |     // 根据数据源采集数据，使用 动态 的内存做封装
 62 |     return 0;
 63 | }
 64 | 
 65 | static int rknn_plugin_input_release(struct ThreadData *td, struct InputUnit *input_unit) {
 66 |     // 释放 rknn_plugin_input 中申请的内存
 67 |      return 0;
 68 | }
 69 | 
 70 | static int rknn_plugin_output(struct ThreadData *td, struct OutputUnit *output_unit) {
 71 |     // 处理输出数据
 72 |     d_rknn_plugin_info("plugin print output data, thread_id: %d", td->thread_id)
 73 |     return 0;
 74 | }
 75 | 
 76 | // 注册所有本插件的相关函数到插件结构体中
 77 | // 注意：插件名称和结构体定义名称必须和插件动态库名称一致
 78 | static struct PluginStruct rknn_plugin_name = {
 79 |         .plugin_name 		= "rknn_plugin_name",
 80 |         .plugin_version 	= 1,
 81 |         .get_config         = get_config,
 82 |         .set_config         = set_config,
 83 |         .init				= rknn_plugin_init,
 84 |         .uninit 			= rknn_plugin_uninit,
 85 |         .rknn_input 		= rknn_plugin_input,
 86 |         .rknn_input_release = rknn_plugin_input_release,
 87 |         .rknn_output		= rknn_plugin_output,
 88 | };
 89 | 
 90 | // 插件动态库在加载时会自动调用该函数
 91 | static void plugin_init plugin_auto_register(){
 92 |     d_rknn_plugin_info("auto register plugin %p, name: %s", &rknn_plugin_name, rknn_plugin_name.plugin_name)
 93 |     plugin_register(&rknn_plugin_name);
 94 | }
 95 | 
 96 | // 插件动态库在关闭时会自动调用该函数
 97 | static void plugin_exit plugin_auto_unregister(){
 98 |     d_rknn_plugin_info("auto unregister plugin %p, name: %s", &rknn_plugin_name, rknn_plugin_name.plugin_name)
 99 |     plugin_unregister(&rknn_plugin_name);
100 | }
101 | 


--------------------------------------------------------------------------------
/rknn_plugins/rknn_yolo_v5/drm_func.h:
--------------------------------------------------------------------------------
 1 | #ifndef __DRM_FUNC_H__
 2 | #define __DRM_FUNC_H__
 3 | #include <stdint.h>
 4 | #include <stdio.h>
 5 | #include <stdlib.h>
 6 | #include <string.h>
 7 | #include <sys/fcntl.h>// open function
 8 | #include <unistd.h> // close function
 9 | #include <errno.h>
10 | #include <sys/mman.h>
11 | 
12 | 
13 | #include <linux/input.h>
14 | #include "libdrm/drm_fourcc.h"
15 | #include "xf86drm.h"
16 | 
17 | #ifdef __cplusplus
18 | extern "C" {
19 | #endif
20 | 
21 | typedef int (* FUNC_DRM_IOCTL)(int fd, unsigned long request, void *arg);
22 | 
23 | typedef struct _drm_context{
24 |     void *drm_handle;
25 |     FUNC_DRM_IOCTL io_func;
26 | } drm_context;
27 | 
28 | /* memory type definitions. */
29 | enum drm_rockchip_gem_mem_type
30 | {
31 |     /* Physically Continuous memory and used as default. */
32 |     ROCKCHIP_BO_CONTIG = 1 << 0,
33 |     /* cachable mapping. */
34 |     ROCKCHIP_BO_CACHABLE = 1 << 1,
35 |     /* write-combine mapping. */
36 |     ROCKCHIP_BO_WC = 1 << 2,
37 |     ROCKCHIP_BO_SECURE = 1 << 3,
38 |     ROCKCHIP_BO_MASK = ROCKCHIP_BO_CONTIG | ROCKCHIP_BO_CACHABLE |
39 |                        ROCKCHIP_BO_WC | ROCKCHIP_BO_SECURE
40 | };
41 | 
42 | int drm_init(drm_context *drm_ctx);
43 | 
44 | void* drm_buf_alloc(drm_context *drm_ctx,int drm_fd, int TexWidth, int TexHeight,int bpp,int *fd,unsigned int *handle,size_t *actual_size);
45 | 
46 | int drm_buf_destroy(drm_context *drm_ctx,int drm_fd,int buf_fd, int handle,void *drm_buf,size_t size);
47 | 
48 | void drm_deinit(drm_context *drm_ctx, int drm_fd);
49 | 
50 | #ifdef __cplusplus
51 | }
52 | #endif
53 | #endif /*__DRM_FUNC_H__*/


--------------------------------------------------------------------------------
/rknn_plugins/rknn_yolo_v5/postprocess.cc:
--------------------------------------------------------------------------------
  1 | // Copyright (c) 2021 by Rockchip Electronics Co., Ltd. All Rights Reserved.
  2 | //
  3 | // Licensed under the Apache License, Version 2.0 (the "License");
  4 | // you may not use this file except in compliance with the License.
  5 | // You may obtain a copy of the License at
  6 | //
  7 | //     http://www.apache.org/licenses/LICENSE-2.0
  8 | //
  9 | // Unless required by applicable law or agreed to in writing, software
 10 | // distributed under the License is distributed on an "AS IS" BASIS,
 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12 | // See the License for the specific language governing permissions and
 13 | // limitations under the License.
 14 | 
 15 | #include "postprocess.h"
 16 | 
 17 | #include <math.h>
 18 | #include <stdint.h>
 19 | #include <stdio.h>
 20 | #include <stdlib.h>
 21 | #include <string.h>
 22 | #include <sys/time.h>
 23 | 
 24 | #include <set>
 25 | #include <vector>
 26 | #define LABEL_NALE_TXT_PATH "./model/coco_80_labels_list.txt"
 27 | 
 28 | static char* labels[OBJ_CLASS_NUM];
 29 | 
 30 | const int anchor0[6] = {10, 13, 16, 30, 33, 23};
 31 | const int anchor1[6] = {30, 61, 62, 45, 59, 119};
 32 | const int anchor2[6] = {116, 90, 156, 198, 373, 326};
 33 | 
 34 | inline static int clamp(float val, int min, int max) { return val > min ? (val < max ? val : max) : min; }
 35 | 
 36 | char* readLine(FILE* fp, char* buffer, int* len)
 37 | {
 38 |   int    ch;
 39 |   int    i        = 0;
 40 |   size_t buff_len = 0;
 41 | 
 42 |   buffer = (char*)malloc(buff_len + 1);
 43 |   if (!buffer)
 44 |     return NULL; // Out of memory
 45 | 
 46 |   while ((ch = fgetc(fp)) != '\n' && ch != EOF) {
 47 |     buff_len++;
 48 |     void* tmp = realloc(buffer, buff_len + 1);
 49 |     if (tmp == NULL) {
 50 |       free(buffer);
 51 |       return NULL; // Out of memory
 52 |     }
 53 |     buffer = (char*)tmp;
 54 | 
 55 |     buffer[i] = (char)ch;
 56 |     i++;
 57 |   }
 58 |   buffer[i] = '\0';
 59 | 
 60 |   *len = buff_len;
 61 | 
 62 |   // Detect end
 63 |   if (ch == EOF && (i == 0 || ferror(fp))) {
 64 |     free(buffer);
 65 |     return NULL;
 66 |   }
 67 |   return buffer;
 68 | }
 69 | 
 70 | int readLines(const char* fileName, char* lines[], int max_line)
 71 | {
 72 |   FILE* file = fopen(fileName, "r");
 73 |   char* s;
 74 |   int   i = 0;
 75 |   int   n = 0;
 76 | 
 77 |   if (file == NULL) {
 78 |     printf("Open %s fail!\n", fileName);
 79 |     return -1;
 80 |   }
 81 | 
 82 |   while ((s = readLine(file, s, &n)) != NULL) {
 83 |     lines[i++] = s;
 84 |     if (i >= max_line)
 85 |       break;
 86 |   }
 87 |   fclose(file);
 88 |   return i;
 89 | }
 90 | 
 91 | int loadLabelName(const char* locationFilename, char* label[])
 92 | {
 93 |   printf("loadLabelName %s\n", locationFilename);
 94 |   readLines(locationFilename, label, OBJ_CLASS_NUM);
 95 |   return 0;
 96 | }
 97 | 
 98 | static float CalculateOverlap(float xmin0, float ymin0, float xmax0, float ymax0, float xmin1, float ymin1, float xmax1,
 99 |                               float ymax1)
100 | {
101 |   float w = fmax(0.f, fmin(xmax0, xmax1) - fmax(xmin0, xmin1) + 1.0);
102 |   float h = fmax(0.f, fmin(ymax0, ymax1) - fmax(ymin0, ymin1) + 1.0);
103 |   float i = w * h;
104 |   float u = (xmax0 - xmin0 + 1.0) * (ymax0 - ymin0 + 1.0) + (xmax1 - xmin1 + 1.0) * (ymax1 - ymin1 + 1.0) - i;
105 |   return u <= 0.f ? 0.f : (i / u);
106 | }
107 | 
108 | static int nms(int validCount, std::vector<float>& outputLocations, std::vector<int> classIds, std::vector<int>& order,
109 |                int filterId, float threshold)
110 | {
111 |   for (int i = 0; i < validCount; ++i) {
112 |     if (order[i] == -1 || classIds[i] != filterId) {
113 |       continue;
114 |     }
115 |     int n = order[i];
116 |     for (int j = i + 1; j < validCount; ++j) {
117 |       int m = order[j];
118 |       if (m == -1 || classIds[i] != filterId) {
119 |         continue;
120 |       }
121 |       float xmin0 = outputLocations[n * 4 + 0];
122 |       float ymin0 = outputLocations[n * 4 + 1];
123 |       float xmax0 = outputLocations[n * 4 + 0] + outputLocations[n * 4 + 2];
124 |       float ymax0 = outputLocations[n * 4 + 1] + outputLocations[n * 4 + 3];
125 | 
126 |       float xmin1 = outputLocations[m * 4 + 0];
127 |       float ymin1 = outputLocations[m * 4 + 1];
128 |       float xmax1 = outputLocations[m * 4 + 0] + outputLocations[m * 4 + 2];
129 |       float ymax1 = outputLocations[m * 4 + 1] + outputLocations[m * 4 + 3];
130 | 
131 |       float iou = CalculateOverlap(xmin0, ymin0, xmax0, ymax0, xmin1, ymin1, xmax1, ymax1);
132 | 
133 |       if (iou > threshold) {
134 |         order[j] = -1;
135 |       }
136 |     }
137 |   }
138 |   return 0;
139 | }
140 | 
141 | static int quick_sort_indice_inverse(std::vector<float>& input, int left, int right, std::vector<int>& indices)
142 | {
143 |   float key;
144 |   int   key_index;
145 |   int   low  = left;
146 |   int   high = right;
147 |   if (left < right) {
148 |     key_index = indices[left];
149 |     key       = input[left];
150 |     while (low < high) {
151 |       while (low < high && input[high] <= key) {
152 |         high--;
153 |       }
154 |       input[low]   = input[high];
155 |       indices[low] = indices[high];
156 |       while (low < high && input[low] >= key) {
157 |         low++;
158 |       }
159 |       input[high]   = input[low];
160 |       indices[high] = indices[low];
161 |     }
162 |     input[low]   = key;
163 |     indices[low] = key_index;
164 |     quick_sort_indice_inverse(input, left, low - 1, indices);
165 |     quick_sort_indice_inverse(input, low + 1, right, indices);
166 |   }
167 |   return low;
168 | }
169 | 
170 | static float sigmoid(float x) { return 1.0 / (1.0 + expf(-x)); }
171 | 
172 | static float unsigmoid(float y) { return -1.0 * logf((1.0 / y) - 1.0); }
173 | 
174 | inline static int32_t __clip(float val, float min, float max)
175 | {
176 |   float f = val <= min ? min : (val >= max ? max : val);
177 |   return f;
178 | }
179 | 
180 | static int8_t qnt_f32_to_affine(float f32, int32_t zp, float scale)
181 | {
182 |   float  dst_val = (f32 / scale) + zp;
183 |   int8_t res     = (int8_t)__clip(dst_val, -128, 127);
184 |   return res;
185 | }
186 | 
187 | static float deqnt_affine_to_f32(int8_t qnt, int32_t zp, float scale) { return ((float)qnt - (float)zp) * scale; }
188 | 
189 | static int process(int8_t* input, int* anchor, int grid_h, int grid_w, int height, int width, int stride,
190 |                    std::vector<float>& boxes, std::vector<float>& objProbs, std::vector<int>& classId, float threshold,
191 |                    int32_t zp, float scale)
192 | {
193 |   int    validCount = 0;
194 |   int    grid_len   = grid_h * grid_w;
195 |   float  thres      = unsigmoid(threshold);
196 |   int8_t thres_i8   = qnt_f32_to_affine(thres, zp, scale);
197 |   for (int a = 0; a < 3; a++) {
198 |     for (int i = 0; i < grid_h; i++) {
199 |       for (int j = 0; j < grid_w; j++) {
200 |         int8_t box_confidence = input[(PROP_BOX_SIZE * a + 4) * grid_len + i * grid_w + j];
201 |         if (box_confidence >= thres_i8) {
202 |           int     offset = (PROP_BOX_SIZE * a) * grid_len + i * grid_w + j;
203 |           int8_t* in_ptr = input + offset;
204 |           float   box_x  = sigmoid(deqnt_affine_to_f32(*in_ptr, zp, scale)) * 2.0 - 0.5;
205 |           float   box_y  = sigmoid(deqnt_affine_to_f32(in_ptr[grid_len], zp, scale)) * 2.0 - 0.5;
206 |           float   box_w  = sigmoid(deqnt_affine_to_f32(in_ptr[2 * grid_len], zp, scale)) * 2.0;
207 |           float   box_h  = sigmoid(deqnt_affine_to_f32(in_ptr[3 * grid_len], zp, scale)) * 2.0;
208 |           box_x          = (box_x + j) * (float)stride;
209 |           box_y          = (box_y + i) * (float)stride;
210 |           box_w          = box_w * box_w * (float)anchor[a * 2];
211 |           box_h          = box_h * box_h * (float)anchor[a * 2 + 1];
212 |           box_x -= (box_w / 2.0);
213 |           box_y -= (box_h / 2.0);
214 | 
215 |           int8_t maxClassProbs = in_ptr[5 * grid_len];
216 |           int    maxClassId    = 0;
217 |           for (int k = 1; k < OBJ_CLASS_NUM; ++k) {
218 |             int8_t prob = in_ptr[(5 + k) * grid_len];
219 |             if (prob > maxClassProbs) {
220 |               maxClassId    = k;
221 |               maxClassProbs = prob;
222 |             }
223 |           }
224 |           if (maxClassProbs>thres_i8){
225 |             objProbs.push_back(sigmoid(deqnt_affine_to_f32(maxClassProbs, zp, scale))* sigmoid(deqnt_affine_to_f32(box_confidence, zp, scale)));
226 |             classId.push_back(maxClassId);
227 |             validCount++;
228 |             boxes.push_back(box_x);
229 |             boxes.push_back(box_y);
230 |             boxes.push_back(box_w);
231 |             boxes.push_back(box_h);
232 |           }
233 |         }
234 |       }
235 |     }
236 |   }
237 |   return validCount;
238 | }
239 | 
240 | int post_process(int8_t* input0, int8_t* input1, int8_t* input2, int model_in_h, int model_in_w, float conf_threshold,
241 |                  float nms_threshold, float scale_w, float scale_h, std::vector<int32_t>& qnt_zps,
242 |                  std::vector<float>& qnt_scales, detect_result_group_t* group)
243 | {
244 |   static int init = -1;
245 |   if (init == -1) {
246 |     int ret = 0;
247 |     ret     = loadLabelName(LABEL_NALE_TXT_PATH, labels);
248 |     if (ret < 0) {
249 |       return -1;
250 |     }
251 | 
252 |     init = 0;
253 |   }
254 |   memset(group, 0, sizeof(detect_result_group_t));
255 | 
256 |   std::vector<float> filterBoxes;
257 |   std::vector<float> objProbs;
258 |   std::vector<int>   classId;
259 | 
260 |   // stride 8
261 |   int stride0     = 8;
262 |   int grid_h0     = model_in_h / stride0;
263 |   int grid_w0     = model_in_w / stride0;
264 |   int validCount0 = 0;
265 |   validCount0 = process(input0, (int*)anchor0, grid_h0, grid_w0, model_in_h, model_in_w, stride0, filterBoxes, objProbs,
266 |                         classId, conf_threshold, qnt_zps[0], qnt_scales[0]);
267 | 
268 |   // stride 16
269 |   int stride1     = 16;
270 |   int grid_h1     = model_in_h / stride1;
271 |   int grid_w1     = model_in_w / stride1;
272 |   int validCount1 = 0;
273 |   validCount1 = process(input1, (int*)anchor1, grid_h1, grid_w1, model_in_h, model_in_w, stride1, filterBoxes, objProbs,
274 |                         classId, conf_threshold, qnt_zps[1], qnt_scales[1]);
275 | 
276 |   // stride 32
277 |   int stride2     = 32;
278 |   int grid_h2     = model_in_h / stride2;
279 |   int grid_w2     = model_in_w / stride2;
280 |   int validCount2 = 0;
281 |   validCount2 = process(input2, (int*)anchor2, grid_h2, grid_w2, model_in_h, model_in_w, stride2, filterBoxes, objProbs,
282 |                         classId, conf_threshold, qnt_zps[2], qnt_scales[2]);
283 | 
284 |   int validCount = validCount0 + validCount1 + validCount2;
285 |   // no object detect
286 |   if (validCount <= 0) {
287 |     return 0;
288 |   }
289 | 
290 |   std::vector<int> indexArray;
291 |   for (int i = 0; i < validCount; ++i) {
292 |     indexArray.push_back(i);
293 |   }
294 | 
295 |   quick_sort_indice_inverse(objProbs, 0, validCount - 1, indexArray);
296 | 
297 |   std::set<int> class_set(std::begin(classId), std::end(classId));
298 | 
299 |   for (auto c : class_set) {
300 |     nms(validCount, filterBoxes, classId, indexArray, c, nms_threshold);
301 |   }
302 | 
303 |   int last_count = 0;
304 |   group->count   = 0;
305 |   /* box valid detect target */
306 |   for (int i = 0; i < validCount; ++i) {
307 |     if (indexArray[i] == -1 || last_count >= OBJ_NUMB_MAX_SIZE) {
308 |       continue;
309 |     }
310 |     int n = indexArray[i];
311 | 
312 |     float x1       = filterBoxes[n * 4 + 0];
313 |     float y1       = filterBoxes[n * 4 + 1];
314 |     float x2       = x1 + filterBoxes[n * 4 + 2];
315 |     float y2       = y1 + filterBoxes[n * 4 + 3];
316 |     int   id       = classId[n];
317 |     float obj_conf = objProbs[i];
318 | 
319 |     group->results[last_count].box.left   = (int)(clamp(x1, 0, model_in_w) / scale_w);
320 |     group->results[last_count].box.top    = (int)(clamp(y1, 0, model_in_h) / scale_h);
321 |     group->results[last_count].box.right  = (int)(clamp(x2, 0, model_in_w) / scale_w);
322 |     group->results[last_count].box.bottom = (int)(clamp(y2, 0, model_in_h) / scale_h);
323 |     group->results[last_count].prop       = obj_conf;
324 |     char* label                           = labels[id];
325 |     strncpy(group->results[last_count].name, label, OBJ_NAME_MAX_SIZE);
326 | 
327 |     // printf("result %2d: (%4d, %4d, %4d, %4d), %s\n", i, group->results[last_count].box.left,
328 |     // group->results[last_count].box.top,
329 |     //        group->results[last_count].box.right, group->results[last_count].box.bottom, label);
330 |     last_count++;
331 |   }
332 |   group->count = last_count;
333 | 
334 |   return 0;
335 | }
336 | 
337 | void deinitPostProcess()
338 | {
339 |   for (int i = 0; i < OBJ_CLASS_NUM; i++) {
340 |     if (labels[i] != nullptr) {
341 |       free(labels[i]);
342 |       labels[i] = nullptr;
343 |     }
344 |   }
345 | }
346 | 


--------------------------------------------------------------------------------
/rknn_plugins/rknn_yolo_v5/postprocess.h:
--------------------------------------------------------------------------------
 1 | #ifndef _RKNN_ZERO_COPY_DEMO_POSTPROCESS_H_
 2 | #define _RKNN_ZERO_COPY_DEMO_POSTPROCESS_H_
 3 | 
 4 | #include <stdint.h>
 5 | #include <vector>
 6 | 
 7 | #define OBJ_NAME_MAX_SIZE 16
 8 | #define OBJ_NUMB_MAX_SIZE 64
 9 | #define OBJ_CLASS_NUM     80
10 | #define NMS_THRESH        0.45
11 | #define BOX_THRESH        0.25
12 | #define PROP_BOX_SIZE     (5+OBJ_CLASS_NUM)
13 | 
14 | typedef struct _BOX_RECT
15 | {
16 |     int left;
17 |     int right;
18 |     int top;
19 |     int bottom;
20 | } BOX_RECT;
21 | 
22 | typedef struct __detect_result_t
23 | {
24 |     char name[OBJ_NAME_MAX_SIZE];
25 |     BOX_RECT box;
26 |     float prop;
27 | } detect_result_t;
28 | 
29 | typedef struct _detect_result_group_t
30 | {
31 |     int id;
32 |     int count;
33 |     detect_result_t results[OBJ_NUMB_MAX_SIZE];
34 | } detect_result_group_t;
35 | 
36 | int post_process(int8_t *input0, int8_t *input1, int8_t *input2, int model_in_h, int model_in_w,
37 |                  float conf_threshold, float nms_threshold, float scale_w, float scale_h,
38 |                  std::vector<int32_t> &qnt_zps, std::vector<float> &qnt_scales,
39 |                  detect_result_group_t *group);
40 | 
41 | void deinitPostProcess();
42 | #endif //_RKNN_ZERO_COPY_DEMO_POSTPROCESS_H_
43 | 


--------------------------------------------------------------------------------
/rknn_plugins/rknn_yolo_v5/rga_func.h:
--------------------------------------------------------------------------------
 1 | #ifndef __RGA_FUNC_H__
 2 | #define __RGA_FUNC_H__
 3 | 
 4 | #include <dlfcn.h> 
 5 | #include "RgaApi.h"
 6 | 
 7 | #ifdef __cplusplus
 8 | extern "C" {
 9 | #endif
10 | 
11 | typedef int(* FUNC_RGA_INIT)();
12 | typedef void(* FUNC_RGA_DEINIT)();
13 | typedef int(* FUNC_RGA_BLIT)(rga_info_t *, rga_info_t *, rga_info_t *);
14 | 
15 | typedef struct _rga_context{
16 |     void *rga_handle;
17 |     FUNC_RGA_INIT init_func;
18 |     FUNC_RGA_DEINIT deinit_func;
19 |     FUNC_RGA_BLIT blit_func;
20 | } rga_context;
21 | 
22 | int RGA_init(rga_context* rga_ctx);
23 | 
24 | void img_resize_fast(rga_context *rga_ctx, int src_fd, int src_w, int src_h, uint64_t dst_phys, int dst_w, int dst_h);
25 | 
26 | void img_resize_slow(rga_context *rga_ctx, void *src_virt, int src_w, int src_h, void *dst_virt, int dst_w, int dst_h);
27 | 
28 | int RGA_deinit(rga_context* rga_ctx);
29 | 
30 | #ifdef __cplusplus
31 | }
32 | #endif
33 | #endif/*__RGA_FUNC_H__*/
34 | 


--------------------------------------------------------------------------------
/rknn_plugins/rknn_yolo_v5/rknn_yolo_v5.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.03
  5 |  * @brief: RKNN 图像推理示例插件，参考：https://github.com/rockchip-linux/rknpu2/tree/master/examples/rknn_yolov5_demo 实现
  6 |  */
  7 | 
  8 | /* 包含定义插件必须的头文件 */
  9 | #include <string>
 10 | #include <cstring>
 11 | 
 12 | #include "opencv2/core/core.hpp"
 13 | #include "opencv2/imgcodecs.hpp"
 14 | #include "opencv2/imgproc.hpp"
 15 | 
 16 | #include "RgaUtils.h"
 17 | #include "im2d.h"
 18 | #include "rga.h"
 19 | 
 20 | #include "rknn_infer_api.h"
 21 | #include "utils_log.h"
 22 | 
 23 | #include "postprocess.h"
 24 | 
 25 | const float    nms_threshold      = NMS_THRESH;
 26 | const float    box_conf_threshold = BOX_THRESH;
 27 | 
 28 | // 插件全局配置信息，由调度程序给插件传来的信息
 29 | PluginConfigSet g_plugin_config_set;
 30 | 
 31 | // 输入线程私有数据
 32 | struct PluginInputData {
 33 |     // 每个线程定制输入源
 34 |     std::string image_path;
 35 | };
 36 | 
 37 | // 输出线程私有数据
 38 | struct PluginOutputData {
 39 |     // 每个线程定制输出
 40 |     std::string image_path;
 41 | };
 42 | 
 43 | // 插件输入输出线程同步数据
 44 | struct PluginSyncData {
 45 |     // 线程输入图像缓存
 46 |     cv::Mat orig_img;
 47 |     // 模型出入结构
 48 |     uint32_t input_channel = 3;
 49 |     uint32_t input_width   = 0;
 50 |     uint32_t input_height  = 0;
 51 | };
 52 | 
 53 | static int get_config(PluginConfigGet *plugin_config){
 54 |     // 输入线程个数
 55 |     plugin_config->input_thread_nums = 2;
 56 |     // 输出线程个数
 57 |     plugin_config->output_thread_nums = 2;
 58 |     // 是否需要输出float类型的输出结果
 59 |     plugin_config->output_want_float = false;
 60 | 
 61 |     d_rknn_plugin_info("post process config: box_conf_threshold = %.2f, nms_threshold = %.2f", box_conf_threshold, nms_threshold);
 62 |     return 0;
 63 | }
 64 | 
 65 | static int set_config(PluginConfigSet *plugin_config){
 66 |     // 注意拷贝构造函数
 67 |     memcpy(&g_plugin_config_set, plugin_config, sizeof(PluginConfigSet));
 68 |     d_rknn_plugin_info("plugin config set success")
 69 |     return 0;
 70 | }
 71 | 
 72 | static int rknn_plugin_init(struct ThreadData *td) {
 73 |     // 设置输入线程的输入源
 74 |     if(td->thread_type == THREAD_TYPE_INPUT) {
 75 |         td->plugin_private_data = new PluginInputData();
 76 |         auto *pri_data = (PluginInputData *)td->plugin_private_data;
 77 |         if(td->thread_id == 0){
 78 |             pri_data->image_path = "model/bus.jpg";
 79 |         }else if(td->thread_id == 1) {
 80 |             pri_data->image_path = "model/bus.jpg";
 81 |         }
 82 |     }else{
 83 |         // 设置输出线程的输出源
 84 |         td->plugin_private_data = new PluginOutputData();
 85 |         auto *pri_data = (PluginOutputData *)td->plugin_private_data;
 86 |         if(td->thread_id == 0) {
 87 |             pri_data->image_path = "thread_1_out.jpg";
 88 |         }else if(td->thread_id == 1) {
 89 |             pri_data->image_path = "thread_2_out.jpg";
 90 |         }
 91 |     }
 92 |     return 0;
 93 | }
 94 | 
 95 | static int rknn_plugin_uninit(struct ThreadData *td) {
 96 |     // 释放输入线程的输入源
 97 |     if(td->thread_type == THREAD_TYPE_INPUT) {
 98 |         auto *pri_data = (PluginInputData *)td->plugin_private_data;
 99 |         delete pri_data;
100 |     }
101 |     return 0;
102 | }
103 | 
104 | static int rknn_plugin_input(struct ThreadData *td, struct InputUnit *input_unit) {
105 |     // 根据数据源采集数据，使用 动态 的内存做封装
106 |     auto *pri_data = (PluginInputData *)td->plugin_private_data;
107 |     td->plugin_sync_data = new PluginSyncData();
108 |     auto *sync_data = (PluginSyncData *)td->plugin_sync_data;
109 | 
110 |     // Load image
111 |     sync_data->orig_img = imread(pri_data->image_path, cv::IMREAD_COLOR);
112 |     if (!sync_data->orig_img.data) {
113 |         d_rknn_plugin_error("cv::imread %s fail!", pri_data->image_path.c_str());
114 |         return -1;
115 |     }
116 | 
117 |     cv::Mat img;
118 |     cv::cvtColor(sync_data->orig_img, img, cv::COLOR_BGR2RGB);
119 |     d_rknn_plugin_info("img input_width = %d, img input_height = %d", sync_data->orig_img.cols, sync_data->orig_img.rows);
120 | 
121 |     if (g_plugin_config_set.input_attr[0].fmt == RKNN_TENSOR_NCHW) {
122 |         d_rknn_plugin_info("model is NCHW input fmt");
123 |         sync_data->input_channel = g_plugin_config_set.input_attr[0].dims[1];
124 |         sync_data->input_height  = g_plugin_config_set.input_attr[0].dims[2];
125 |         sync_data->input_width   = g_plugin_config_set.input_attr[0].dims[3];
126 |     } else {
127 |         d_rknn_plugin_info("model is NHWC input fmt");
128 |         sync_data->input_height  = g_plugin_config_set.input_attr[0].dims[1];
129 |         sync_data->input_width   = g_plugin_config_set.input_attr[0].dims[2];
130 |         sync_data->input_channel = g_plugin_config_set.input_attr[0].dims[3];
131 |     }
132 |     d_rknn_plugin_info("model input input_height=%d, input_width=%d, input_channel=%d", sync_data->input_height, sync_data->input_width, sync_data->input_channel);
133 | 
134 |     input_unit->n_inputs = g_plugin_config_set.io_num.n_input;
135 |     input_unit->inputs = (rknn_input*)malloc(input_unit->n_inputs * sizeof(rknn_input));
136 |     memset(input_unit->inputs, 0, input_unit->n_inputs * sizeof(rknn_input));
137 |     input_unit->inputs[0].index        = 0;
138 |     input_unit->inputs[0].type         = RKNN_TENSOR_UINT8;
139 |     input_unit->inputs[0].size         = sync_data->input_width * sync_data->input_height * sync_data->input_channel;
140 |     input_unit->inputs[0].fmt          = RKNN_TENSOR_NHWC;
141 |     input_unit->inputs[0].pass_through = 0;
142 |     input_unit->inputs[0].buf = new uint8_t[sync_data->input_width * sync_data->input_height * sync_data->input_channel];
143 | 
144 |     if (sync_data->orig_img.cols != sync_data->input_width || sync_data->orig_img.rows != sync_data->input_height) {
145 |         d_rknn_plugin_info("resize with RGA!");
146 | 
147 |         // init rga context
148 |         rga_buffer_t src;
149 |         rga_buffer_t dst;
150 |         im_rect      src_rect;
151 |         im_rect      dst_rect;
152 |         memset(&src_rect, 0, sizeof(src_rect));
153 |         memset(&dst_rect, 0, sizeof(dst_rect));
154 |         memset(&src, 0, sizeof(src));
155 |         memset(&dst, 0, sizeof(dst));
156 | 
157 |         void* resize_buf = input_unit->inputs[0].buf;
158 |         memset(resize_buf, 0x00, sync_data->input_width * sync_data->input_height * sync_data->input_channel);
159 | 
160 |         src = wrapbuffer_virtualaddr((void*)img.data, sync_data->orig_img.cols, sync_data->orig_img.rows, RK_FORMAT_RGB_888);
161 |         dst = wrapbuffer_virtualaddr((void*)resize_buf, sync_data->input_width, sync_data->input_height, RK_FORMAT_RGB_888);
162 |         int ret = imcheck(src, dst, src_rect, dst_rect);
163 |         if (IM_STATUS_NOERROR != ret) {
164 |             d_rknn_plugin_info("%d, check error! %s", __LINE__, imStrError((IM_STATUS)ret));
165 |             return -1;
166 |         }
167 |         IM_STATUS STATUS = imresize(src, dst);
168 |         if (IM_STATUS_NOERROR != STATUS) {
169 |             d_rknn_plugin_info("%d, resize error! %s", __LINE__, imStrError(STATUS));
170 |             return -1;
171 |         }
172 |     } else {
173 |         memcpy(input_unit->inputs[0].buf, img.data, input_unit->inputs[0].size);
174 |     }
175 |     return 0;
176 | }
177 | 
178 | static int rknn_plugin_input_release(struct ThreadData *td, struct InputUnit *input_unit) {
179 |     // 释放 rknn_plugin_input 中申请的内存
180 |     for(uint32_t idx = 0; idx < input_unit->n_inputs; idx++){
181 |         delete[] (uint8_t*)input_unit->inputs[0].buf;
182 |         input_unit->inputs[0].buf = nullptr;
183 |     }
184 |     free(input_unit->inputs);
185 |     return 0;
186 | }
187 | 
188 | static int rknn_plugin_output(struct ThreadData *td, struct OutputUnit *output_unit) {
189 |     // 处理输出数据
190 |     d_rknn_plugin_info("plugin print output data, thread_id: %d", td->thread_id)
191 |     auto *pri_data = (PluginOutputData *)td->plugin_private_data;
192 |     auto *sync_data = (PluginSyncData *)td->plugin_sync_data;
193 | 
194 |     // post process
195 |     float scale_w = (float)sync_data->input_width / (float)sync_data->orig_img.cols;
196 |     float scale_h = (float)sync_data->input_height / (float)sync_data->orig_img.rows;
197 |     d_rknn_plugin_info("scale_w=%f, scale_h=%f", scale_w, scale_h);
198 | 
199 |     detect_result_group_t detect_result_group;
200 |     std::vector<float>    out_scales;
201 |     std::vector<int32_t>  out_zps;
202 |     for (int i = 0; i < output_unit->n_outputs; ++i) {
203 |         out_scales.push_back(g_plugin_config_set.output_attr[i].scale);
204 |         out_zps.push_back(g_plugin_config_set.output_attr[i].zp);
205 |     }
206 |     post_process(
207 |             (int8_t*)output_unit->outputs[0].buf,
208 |             (int8_t*)output_unit->outputs[1].buf,
209 |             (int8_t*)output_unit->outputs[2].buf,
210 |             (int)sync_data->input_height, (int)sync_data->input_width,
211 |             box_conf_threshold, nms_threshold,
212 |             scale_w, scale_h,
213 |             out_zps, out_scales, &detect_result_group);
214 | 
215 |     // Draw Objects
216 |     char text[256];
217 |     for (int i = 0; i < detect_result_group.count; i++) {
218 |         detect_result_t* det_result = &(detect_result_group.results[i]);
219 |         sprintf(text, "%s %.1f%%", det_result->name, det_result->prop * 100);
220 |         d_rknn_plugin_info("%s @ (%d %d %d %d) %f", det_result->name, det_result->box.left, det_result->box.top,
221 |                det_result->box.right, det_result->box.bottom, det_result->prop);
222 |         int x1 = det_result->box.left;
223 |         int y1 = det_result->box.top;
224 |         int x2 = det_result->box.right;
225 |         int y2 = det_result->box.bottom;
226 |         rectangle(sync_data->orig_img, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(255, 0, 0, 255), 3);
227 |         putText(sync_data->orig_img, text, cv::Point(x1, y1 + 12), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));
228 |     }
229 | 
230 |     imwrite(pri_data->image_path, sync_data->orig_img);
231 | 
232 |     // 释放输出
233 |     delete sync_data;
234 |     td->plugin_sync_data = nullptr;
235 |     return 0;
236 | }
237 | 
238 | // 注册所有本插件的相关函数到插件结构体中
239 | // 注意：插件名称和结构体定义名称必须和插件动态库名称一致
240 | static struct PluginStruct rknn_yolo_v5 = {
241 |         .plugin_name 		= "rknn_yolo_v5",
242 |         .plugin_version 	= 1,
243 |         .get_config         = get_config,
244 |         .set_config         = set_config,
245 |         .init				= rknn_plugin_init,
246 |         .uninit 			= rknn_plugin_uninit,
247 |         .rknn_input 		= rknn_plugin_input,
248 |         .rknn_input_release = rknn_plugin_input_release,
249 |         .rknn_output		= rknn_plugin_output,
250 | };
251 | 
252 | // 插件动态库在加载时会自动调用该函数
253 | static void plugin_init plugin_auto_register(){
254 |     d_rknn_plugin_info("auto register plugin %p, name: %s", &rknn_yolo_v5, rknn_yolo_v5.plugin_name)
255 |     plugin_register(&rknn_yolo_v5);
256 | }
257 | 
258 | // 插件动态库在关闭时会自动调用该函数
259 | static void plugin_exit plugin_auto_unregister(){
260 |     d_rknn_plugin_info("auto unregister plugin %p, name: %s", &rknn_yolo_v5, rknn_yolo_v5.plugin_name)
261 |     deinitPostProcess();
262 |     plugin_unregister(&rknn_yolo_v5);
263 | }
264 | 


--------------------------------------------------------------------------------
/rknn_plugins/rknn_yolo_v5/rknn_yolo_v5_video.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.03
  5 |  * @brief: RKNN 图像推理示例插件，参考：https://github.com/rockchip-linux/rknpu2/tree/master/examples/rknn_yolov5_demo 实现
  6 |  */
  7 | 
  8 | /* 包含定义插件必须的头文件 */
  9 | #include <string>
 10 | #include <cstring>
 11 | 
 12 | #include "opencv2/core/core.hpp"
 13 | #include "opencv2/imgcodecs.hpp"
 14 | #include "opencv2/imgproc.hpp"
 15 | 
 16 | #include "RgaUtils.h"
 17 | #include "im2d.h"
 18 | #include "rga.h"
 19 | 
 20 | #include "rknn_infer_api.h"
 21 | #include "utils_log.h"
 22 | 
 23 | #include "postprocess.h"
 24 | #include "mpp_video_decoder.h"
 25 | #include "mpp_video_encoder.h"
 26 | 
 27 | const float    nms_threshold      = NMS_THRESH;
 28 | const float    box_conf_threshold = BOX_THRESH;
 29 | 
 30 | // 插件全局配置信息，由调度程序给插件传来的信息
 31 | PluginConfigSet g_plugin_config_set;
 32 | 
 33 | std::vector<MppVideoEncoder *> g_mpp_video_encoders;
 34 | 
 35 | // 输入线程私有数据
 36 | struct PluginInputData {
 37 |     // 每个线程定制解码器
 38 |     MppVideoDecoder *mpp_video_decoder;
 39 | };
 40 | 
 41 | // 输出线程私有数据
 42 | struct PluginOutputData {
 43 |     // 每个线程定制输出
 44 |     std::string image_path;
 45 | };
 46 | 
 47 | // 插件输入输出线程同步数据
 48 | struct PluginSyncData {
 49 |     // 帧缓存
 50 |     DecoderMppFrame frame = {};
 51 |     MppVideoDecoder *mpp_video_decoder = nullptr;
 52 | 
 53 |     int input_thread = -1;
 54 | 
 55 |     // 模型出入结构
 56 |     uint32_t input_channel = 3;
 57 |     uint32_t input_width   = 0;
 58 |     uint32_t input_height  = 0;
 59 | };
 60 | 
 61 | static int get_config(PluginConfigGet *plugin_config){
 62 |     // 输入线程个数
 63 |     plugin_config->input_thread_nums = 2;
 64 |     // 输出线程个数
 65 |     plugin_config->output_thread_nums = 2;
 66 |     // 是否需要输出float类型的输出结果
 67 |     plugin_config->output_want_float = false;
 68 | 
 69 |     d_rknn_plugin_info("post process config: box_conf_threshold = %.2f, nms_threshold = %.2f", box_conf_threshold, nms_threshold);
 70 | 
 71 |     // 模型输出编码器定义
 72 |     auto *mpp_video_encoder_1 = new MppVideoEncoder("out_1.h264");
 73 |     auto *mpp_video_encoder_2 = new MppVideoEncoder("out_2.h264");
 74 |     g_mpp_video_encoders.push_back(mpp_video_encoder_1);
 75 |     g_mpp_video_encoders.push_back(mpp_video_encoder_2);
 76 |     return 0;
 77 | }
 78 | 
 79 | static int set_config(PluginConfigSet *plugin_config){
 80 |     // 注意拷贝构造函数
 81 |     memcpy(&g_plugin_config_set, plugin_config, sizeof(PluginConfigSet));
 82 |     d_rknn_plugin_info("plugin config set success")
 83 |     return 0;
 84 | }
 85 | 
 86 | static int rknn_plugin_init(struct ThreadData *td) {
 87 |     // 设置输入线程的输入源
 88 |     if(td->thread_type == THREAD_TYPE_INPUT) {
 89 |         td->plugin_private_data = new PluginInputData();
 90 |         auto *pri_data = (PluginInputData *)td->plugin_private_data;
 91 |         if(td->thread_id == 0){
 92 |             pri_data->mpp_video_decoder = new MppVideoDecoder("1080p.264");
 93 |         }else if(td->thread_id == 1) {
 94 |             pri_data->mpp_video_decoder = new MppVideoDecoder("1080p.264");
 95 |         }
 96 |     }else{
 97 |         // 设置输出线程的输出源
 98 |         td->plugin_private_data = new PluginOutputData();
 99 |         auto *pri_data = (PluginOutputData *)td->plugin_private_data;
100 |         if(td->thread_id == 0) {
101 |             pri_data->image_path = "thread_1_out.jpg";
102 |         }else if(td->thread_id == 1) {
103 |             pri_data->image_path = "thread_2_out.jpg";
104 |         }
105 |     }
106 |     return 0;
107 | }
108 | 
109 | static int rknn_plugin_uninit(struct ThreadData *td) {
110 |     // 释放输入线程的输入源
111 |     if(td->thread_type == THREAD_TYPE_INPUT) {
112 |         auto *pri_data = (PluginInputData *)td->plugin_private_data;
113 |         delete pri_data;
114 |     }
115 |     return 0;
116 | }
117 | 
118 | static int rknn_plugin_input(struct ThreadData *td, struct InputUnit *input_unit) {
119 |     // 根据数据源采集数据，使用 动态 的内存做封装
120 |     auto *pri_data = (PluginInputData *)td->plugin_private_data;
121 |     td->plugin_sync_data = new PluginSyncData();
122 |     auto *sync_data = (PluginSyncData *)td->plugin_sync_data;
123 |     sync_data->mpp_video_decoder = pri_data->mpp_video_decoder;
124 |     sync_data->input_thread = td->thread_id;
125 | 
126 |     // Load frame
127 |     if (pri_data->mpp_video_decoder->get_next_frame(sync_data->frame) < 0) {
128 |         d_rknn_plugin_error("get_frame fail!");
129 |         return -1;
130 |     }
131 | 
132 |     if (g_plugin_config_set.input_attr[0].fmt == RKNN_TENSOR_NCHW) {
133 |         d_rknn_plugin_info("model is NCHW input fmt");
134 |         sync_data->input_channel = g_plugin_config_set.input_attr[0].dims[1];
135 |         sync_data->input_height  = g_plugin_config_set.input_attr[0].dims[2];
136 |         sync_data->input_width   = g_plugin_config_set.input_attr[0].dims[3];
137 |     } else {
138 |         d_rknn_plugin_info("model is NHWC input fmt");
139 |         sync_data->input_height  = g_plugin_config_set.input_attr[0].dims[1];
140 |         sync_data->input_width   = g_plugin_config_set.input_attr[0].dims[2];
141 |         sync_data->input_channel = g_plugin_config_set.input_attr[0].dims[3];
142 |     }
143 |     d_rknn_plugin_info("model input input_height=%d, input_width=%d, input_channel=%d", sync_data->input_height, sync_data->input_width, sync_data->input_channel);
144 | 
145 |     input_unit->n_inputs = g_plugin_config_set.io_num.n_input;
146 |     input_unit->inputs = (rknn_input*)malloc(input_unit->n_inputs * sizeof(rknn_input));
147 |     memset(input_unit->inputs, 0, input_unit->n_inputs * sizeof(rknn_input));
148 |     input_unit->inputs[0].index        = 0;
149 |     input_unit->inputs[0].type         = RKNN_TENSOR_UINT8;
150 |     input_unit->inputs[0].size         = sync_data->input_width * sync_data->input_height * sync_data->input_channel;
151 |     input_unit->inputs[0].fmt          = RKNN_TENSOR_NHWC;
152 |     input_unit->inputs[0].pass_through = 0;
153 |     input_unit->inputs[0].buf = new uint8_t[sync_data->input_width * sync_data->input_height * sync_data->input_channel];
154 | 
155 |     d_rknn_plugin_info("resize with RGA!");
156 | 
157 |     // init rga context
158 |     rga_buffer_t src;
159 |     rga_buffer_t dst;
160 |     im_rect      src_rect;
161 |     im_rect      dst_rect;
162 |     memset(&src_rect, 0, sizeof(src_rect));
163 |     memset(&dst_rect, 0, sizeof(dst_rect));
164 |     memset(&src, 0, sizeof(src));
165 |     memset(&dst, 0, sizeof(dst));
166 | 
167 |     void* resize_buf = input_unit->inputs[0].buf;
168 |     memset(resize_buf, 0x00, sync_data->input_width * sync_data->input_height * sync_data->input_channel);
169 | 
170 |     src = wrapbuffer_virtualaddr((void*)sync_data->frame.data_buf,
171 |                                  sync_data->frame.hor_width, sync_data->frame.ver_height,
172 |                                  sync_data->frame.mpp_frame_format,
173 |                                  (int)sync_data->frame.hor_stride, (int)sync_data->frame.ver_stride);
174 |     dst = wrapbuffer_virtualaddr((void*)resize_buf, sync_data->input_width, sync_data->input_height, RK_FORMAT_RGB_888);
175 |     int ret = imcheck(src, dst, src_rect, dst_rect);
176 |     if (IM_STATUS_NOERROR != ret) {
177 |         d_rknn_plugin_info("%d, check error! %s", __LINE__, imStrError((IM_STATUS)ret));
178 |         return -1;
179 |     }
180 |     IM_STATUS STATUS = imresize(src, dst);
181 |     if (IM_STATUS_NOERROR != STATUS) {
182 |         d_rknn_plugin_info("%d, resize error! %s", __LINE__, imStrError(STATUS));
183 |         return -1;
184 |     }
185 |     return 0;
186 | }
187 | 
188 | static int rknn_plugin_input_release(struct ThreadData *td, struct InputUnit *input_unit) {
189 |     // 释放 rknn_plugin_input 中申请的内存
190 |     for(uint32_t idx = 0; idx < input_unit->n_inputs; idx++){
191 |         delete[] (uint8_t*)input_unit->inputs[0].buf;
192 |         input_unit->inputs[0].buf = nullptr;
193 |     }
194 |     free(input_unit->inputs);
195 |     return 0;
196 | }
197 | 
198 | static int rknn_plugin_output(struct ThreadData *td, struct OutputUnit *output_unit) {
199 |     // 处理输出数据
200 |     d_rknn_plugin_info("plugin print output data, thread_id: %d", td->thread_id)
201 |     auto *pri_data = (PluginOutputData *)td->plugin_private_data;
202 |     auto *sync_data = (PluginSyncData *)td->plugin_sync_data;
203 | 
204 |     // post process
205 |     float scale_w = (float)sync_data->input_width / (float)sync_data->frame.hor_width;
206 |     float scale_h = (float)sync_data->input_height / (float)sync_data->frame.ver_height;
207 |     d_rknn_plugin_info("scale_w=%f, scale_h=%f", scale_w, scale_h);
208 | 
209 |     detect_result_group_t detect_result_group;
210 |     std::vector<float>    out_scales;
211 |     std::vector<int32_t>  out_zps;
212 |     for (int i = 0; i < output_unit->n_outputs; ++i) {
213 |         out_scales.push_back(g_plugin_config_set.output_attr[i].scale);
214 |         out_zps.push_back(g_plugin_config_set.output_attr[i].zp);
215 |     }
216 |     post_process(
217 |             (int8_t*)output_unit->outputs[0].buf,
218 |             (int8_t*)output_unit->outputs[1].buf,
219 |             (int8_t*)output_unit->outputs[2].buf,
220 |             (int)sync_data->input_height, (int)sync_data->input_width,
221 |             box_conf_threshold, nms_threshold,
222 |             scale_w, scale_h,
223 |             out_zps, out_scales, &detect_result_group);
224 | 
225 | //    // Copy To another buffer avoid to modify mpp decoder buffer
226 | //    rga_buffer_t frame_buffer;
227 | //    rga_buffer_t draw_buffer;
228 | //    uint8_t out_buf = new uint8_t[sync_data->input_width * sync_data->input_height * sync_data->input_channel];
229 | //    frame_buffer = wrapbuffer_virtualaddr((void*)sync_data->frame.data_buf, sync_data->orig_img.cols, sync_data->orig_img.rows, RK_FORMAT_RGB_888);
230 | //    draw_buffer = wrapbuffer_virtualaddr((void*)resize_buf, sync_data->input_width, sync_data->input_height, RK_FORMAT_RGB_888);
231 | //    imcopy(origin, src);
232 | //
233 | //    // Draw objects
234 | //    for (int i = 0; i < detect_result.count; i++) {
235 | //        detect_result_t* det_result = &(detect_result.results[i]);
236 | //        printf("%s @ (%d %d %d %d) %f\n", det_result->name, det_result->box.left, det_result->box.top,
237 | //               det_result->box.right, det_result->box.bottom, det_result->prop);
238 | //        int x1 = det_result->box.left;
239 | //        int y1 = det_result->box.top;
240 | //        int x2 = det_result->box.right;
241 | //        int y2 = det_result->box.bottom;
242 | //        draw_rectangle_yuv420sp((unsigned char*)mpp_frame_addr, width_stride, height_stride, x1, y1, x2-x1+1, y2-y1+1, 0x00FF0000, 4);
243 | //    }
244 |     // 编码器输出
245 | 
246 |     // 释放输出
247 | 
248 |     sync_data->mpp_video_decoder->release_frame(sync_data->frame);
249 |     sync_data->mpp_video_decoder = nullptr;
250 |     delete sync_data;
251 |     td->plugin_sync_data = nullptr;
252 |     return 0;
253 | }
254 | 
255 | // 注册所有本插件的相关函数到插件结构体中
256 | // 注意：插件名称和结构体定义名称必须和插件动态库名称一致
257 | static struct PluginStruct rknn_yolo_v5 = {
258 |         .plugin_name 		= "rknn_yolo_v5",
259 |         .plugin_version 	= 1,
260 |         .get_config         = get_config,
261 |         .set_config         = set_config,
262 |         .init				= rknn_plugin_init,
263 |         .uninit 			= rknn_plugin_uninit,
264 |         .rknn_input 		= rknn_plugin_input,
265 |         .rknn_input_release = rknn_plugin_input_release,
266 |         .rknn_output		= rknn_plugin_output,
267 | };
268 | 
269 | // 插件动态库在加载时会自动调用该函数
270 | static void plugin_init plugin_auto_register(){
271 |     d_rknn_plugin_info("auto register plugin %p, name: %s", &rknn_yolo_v5, rknn_yolo_v5.plugin_name)
272 |     plugin_register(&rknn_yolo_v5);
273 | }
274 | 
275 | // 插件动态库在关闭时会自动调用该函数
276 | static void plugin_exit plugin_auto_unregister(){
277 |     d_rknn_plugin_info("auto unregister plugin %p, name: %s", &rknn_yolo_v5, rknn_yolo_v5.plugin_name)
278 |     deinitPostProcess();
279 | 
280 |     // 清除编码器
281 |     for(auto &it : g_mpp_video_encoders){
282 |         delete it;
283 |     }
284 |     plugin_unregister(&rknn_yolo_v5);
285 | }
286 | 


--------------------------------------------------------------------------------
/rknn_plugins/rknn_yolo_v5/src/main_video.cc:
--------------------------------------------------------------------------------
  1 | // Copyright (c) 2023 by Rockchip Electronics Co., Ltd. All Rights Reserved.
  2 | //
  3 | // Licensed under the Apache License, Version 2.0 (the "License");
  4 | // you may not use this file except in compliance with the License.
  5 | // You may obtain a copy of the License at
  6 | //
  7 | //     http://www.apache.org/licenses/LICENSE-2.0
  8 | //
  9 | // Unless required by applicable law or agreed to in writing, software
 10 | // distributed under the License is distributed on an "AS IS" BASIS,
 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12 | // See the License for the specific language governing permissions and
 13 | // limitations under the License.
 14 | 
 15 | /*-------------------------------------------
 16 |                 Includes
 17 | -------------------------------------------*/
 18 | #include <dlfcn.h>
 19 | #include <stdio.h>
 20 | #include <stdlib.h>
 21 | #include <string.h>
 22 | #include <unistd.h>
 23 | #include <sys/time.h>
 24 | 
 25 | #include "im2d.h"
 26 | #include "rga.h"
 27 | #include "RgaUtils.h"
 28 | 
 29 | #include "rknn_api.h"
 30 | #include "postprocess.h"
 31 | 
 32 | #include "utils/mpp_decoder.h"
 33 | #include "utils/mpp_encoder.h"
 34 | #include "utils/drawing.h"
 35 | #if defined(BUILD_VIDEO_RTSP)
 36 | #include "mk_mediakit.h"
 37 | #endif
 38 | 
 39 | #define OUT_VIDEO_PATH "out.h264"
 40 | 
 41 | typedef struct {
 42 |   rknn_context rknn_ctx;
 43 |   rknn_input_output_num io_num;
 44 |   rknn_tensor_attr* input_attrs;
 45 |   rknn_tensor_attr* output_attrs;
 46 |   int model_channel;
 47 |   int model_width;
 48 |   int model_height;
 49 |   FILE* out_fp;
 50 |   MppDecoder* decoder;
 51 |   MppEncoder* encoder;
 52 | } rknn_app_context_t;
 53 | 
 54 | typedef struct {
 55 |   int width;
 56 |   int height;
 57 |   int width_stride;
 58 |   int height_stride;
 59 |   int format;
 60 |   char* virt_addr;
 61 |   int fd;
 62 | } image_frame_t;
 63 | 
 64 | /*-------------------------------------------
 65 |                   Functions
 66 | -------------------------------------------*/
 67 | 
 68 | static void dump_tensor_attr(rknn_tensor_attr* attr)
 69 | {
 70 |   printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
 71 |          "zp=%d, scale=%f\n",
 72 |          attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
 73 |          attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
 74 |          get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
 75 | }
 76 | 
 77 | double __get_us(struct timeval t) { return (t.tv_sec * 1000000 + t.tv_usec); }
 78 | 
 79 | static unsigned char* load_data(FILE* fp, size_t ofst, size_t sz)
 80 | {
 81 |   unsigned char* data;
 82 |   int ret;
 83 | 
 84 |   data = NULL;
 85 | 
 86 |   if (NULL == fp) {
 87 |     return NULL;
 88 |   }
 89 | 
 90 |   ret = fseek(fp, ofst, SEEK_SET);
 91 |   if (ret != 0) {
 92 |     printf("blob seek failure.\n");
 93 |     return NULL;
 94 |   }
 95 | 
 96 |   data = (unsigned char*)malloc(sz);
 97 |   if (data == NULL) {
 98 |     printf("buffer malloc failure.\n");
 99 |     return NULL;
100 |   }
101 |   ret = fread(data, 1, sz, fp);
102 |   return data;
103 | }
104 | 
105 | static unsigned char* read_file_data(const char* filename, int* model_size)
106 | {
107 |   FILE* fp;
108 |   unsigned char* data;
109 | 
110 |   fp = fopen(filename, "rb");
111 |   if (NULL == fp) {
112 |     printf("Open file %s failed.\n", filename);
113 |     return NULL;
114 |   }
115 | 
116 |   fseek(fp, 0, SEEK_END);
117 |   int size = ftell(fp);
118 | 
119 |   data = load_data(fp, 0, size);
120 | 
121 |   fclose(fp);
122 | 
123 |   *model_size = size;
124 |   return data;
125 | }
126 | 
127 | static int write_data_to_file(const char *path, char *data, unsigned int size) {
128 |   FILE *fp;
129 | 
130 |   fp = fopen(path, "w");
131 |   if(fp == NULL) {
132 |     printf("open error: %s", path);
133 |     return -1;
134 |   }
135 | 
136 |   fwrite(data, 1, size, fp);
137 |   fflush(fp);
138 | 
139 |   fclose(fp);
140 |   return 0;
141 | }
142 | 
143 | static int init_model(const char* model_path, rknn_app_context_t* app_ctx) {
144 |   int ret;
145 |   rknn_context ctx;
146 | 
147 |   /* Create the neural network */
148 |   printf("Loading mode...\n");
149 |   int model_data_size = 0;
150 |   unsigned char* model_data = read_file_data(model_path, &model_data_size);
151 |   if (model_data == NULL) {
152 |     return -1;
153 |   }
154 | 
155 |   ret = rknn_init(&ctx, model_data, model_data_size, 0, NULL);
156 |   if (ret < 0) {
157 |     printf("rknn_init error ret=%d\n", ret);
158 |     return -1;
159 |   }
160 | 
161 |   if (model_data) {
162 |     free(model_data);
163 |   }
164 | 
165 |   rknn_sdk_version version;
166 |   ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version));
167 |   if (ret < 0) {
168 |     printf("rknn_query RKNN_QUERY_SDK_VERSION error ret=%d\n", ret);
169 |     return -1;
170 |   }
171 |   printf("sdk version: %s driver version: %s\n", version.api_version, version.drv_version);
172 | 
173 |   ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &app_ctx->io_num, sizeof(rknn_input_output_num));
174 |   if (ret < 0) {
175 |     printf("rknn_query RKNN_QUERY_IN_OUT_NUM error ret=%d\n", ret);
176 |     return -1;
177 |   }
178 |   printf("model input num: %d, output num: %d\n", app_ctx->io_num.n_input, app_ctx->io_num.n_output);
179 | 
180 |   rknn_tensor_attr* input_attrs = (rknn_tensor_attr*)malloc(app_ctx->io_num.n_input * sizeof(rknn_tensor_attr));
181 |   memset(input_attrs, 0, sizeof(input_attrs));
182 |   for (int i = 0; i < app_ctx->io_num.n_input; i++) {
183 |     input_attrs[i].index = i;
184 |     ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
185 |     if (ret < 0) {
186 |       printf("rknn_query RKNN_QUERY_INPUT_ATTR error ret=%d\n", ret);
187 |       return -1;
188 |     }
189 |     dump_tensor_attr(&(input_attrs[i]));
190 |   }
191 | 
192 |   rknn_tensor_attr* output_attrs = (rknn_tensor_attr*)malloc(app_ctx->io_num.n_output * sizeof(rknn_tensor_attr));
193 |   memset(output_attrs, 0, sizeof(output_attrs));
194 |   for (int i = 0; i < app_ctx->io_num.n_output; i++) {
195 |     output_attrs[i].index = i;
196 |     ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
197 |     if (ret < 0) {
198 |       printf("rknn_query RKNN_QUERY_OUTPUT_ATTR error ret=%d\n", ret);
199 |       return -1;
200 |     }
201 |     dump_tensor_attr(&(output_attrs[i]));
202 |   }
203 | 
204 |   app_ctx->input_attrs = input_attrs;
205 |   app_ctx->output_attrs = output_attrs;
206 |   app_ctx->rknn_ctx = ctx;
207 | 
208 |   if (input_attrs[0].fmt == RKNN_TENSOR_NCHW) {
209 |     printf("model is NCHW input fmt\n");
210 |     app_ctx->model_channel = input_attrs[0].dims[1];
211 |     app_ctx->model_height  = input_attrs[0].dims[2];
212 |     app_ctx->model_width   = input_attrs[0].dims[3];
213 |   } else {
214 |     printf("model is NHWC input fmt\n");
215 |     app_ctx->model_height  = input_attrs[0].dims[1];
216 |     app_ctx->model_width   = input_attrs[0].dims[2];
217 |     app_ctx->model_channel = input_attrs[0].dims[3];
218 |   }
219 |   printf("model input height=%d, width=%d, channel=%d\n", app_ctx->model_height, app_ctx->model_width, app_ctx->model_channel);
220 | 
221 |   return 0;
222 | }
223 | 
224 | static int release_model(rknn_app_context_t* app_ctx) {
225 |   if (app_ctx->rknn_ctx != NULL) {
226 |     rknn_destroy(app_ctx->rknn_ctx);
227 |   }
228 |   free(app_ctx->input_attrs);
229 |   free(app_ctx->output_attrs);
230 |   deinitPostProcess();
231 |   return 0;
232 | }
233 | 
234 | static int inference_model(rknn_app_context_t* app_ctx, image_frame_t* img, detect_result_group_t* detect_result) {
235 |   int ret;
236 |   rknn_context ctx = app_ctx->rknn_ctx;
237 |   int model_width = app_ctx->model_width;
238 |   int model_height = app_ctx->model_height;
239 |   int model_channel = app_ctx->model_channel;
240 | 
241 |   struct timeval start_time, stop_time;
242 |   const float    nms_threshold      = NMS_THRESH;
243 |   const float    box_conf_threshold = BOX_THRESH;
244 |   // You may not need resize when src resulotion equals to dst resulotion
245 |   void* resize_buf = nullptr;
246 |   // init rga context
247 |   rga_buffer_t src;
248 |   rga_buffer_t dst;
249 |   im_rect      src_rect;
250 |   im_rect      dst_rect;
251 |   memset(&src_rect, 0, sizeof(src_rect));
252 |   memset(&dst_rect, 0, sizeof(dst_rect));
253 |   memset(&src, 0, sizeof(src));
254 |   memset(&dst, 0, sizeof(dst));
255 | 
256 |   printf("input image %dx%d stride %dx%d format=%d\n", img->width, img->height, img->width_stride, img->height_stride, img->format);
257 | 
258 |   float scale_w = (float)model_width / img->width;
259 |   float scale_h = (float)model_height / img->height;
260 | 
261 |   rknn_input inputs[1];
262 |   memset(inputs, 0, sizeof(inputs));
263 |   inputs[0].index        = 0;
264 |   inputs[0].type         = RKNN_TENSOR_UINT8;
265 |   inputs[0].size         = model_width * model_height * model_channel;
266 |   inputs[0].fmt          = RKNN_TENSOR_NHWC;
267 |   inputs[0].pass_through = 0;
268 | 
269 |   printf("resize with RGA!\n");
270 |   resize_buf = malloc(model_width * model_height * model_channel);
271 |   memset(resize_buf, 0, model_width * model_height * model_channel);
272 | 
273 |   src = wrapbuffer_virtualaddr((void*)img->virt_addr, img->width, img->height, img->format, img->width_stride, img->height_stride);
274 |   dst = wrapbuffer_virtualaddr((void*)resize_buf, model_width, model_height, RK_FORMAT_RGB_888);
275 |   ret = imcheck(src, dst, src_rect, dst_rect);
276 |   if (IM_STATUS_NOERROR != ret) {
277 |     printf("%d, check error! %s", __LINE__, imStrError((IM_STATUS)ret));
278 |     return -1;
279 |   }
280 |   IM_STATUS STATUS = imresize(src, dst);
281 | 
282 |   inputs[0].buf = resize_buf;
283 | 
284 |   gettimeofday(&start_time, NULL);
285 |   rknn_inputs_set(ctx, app_ctx->io_num.n_input, inputs);
286 | 
287 |   rknn_output outputs[app_ctx->io_num.n_output];
288 |   memset(outputs, 0, sizeof(outputs));
289 |   for (int i = 0; i < app_ctx->io_num.n_output; i++) {
290 |     outputs[i].want_float = 0;
291 |   }
292 | 
293 |   ret = rknn_run(ctx, NULL);
294 |   ret = rknn_outputs_get(ctx, app_ctx->io_num.n_output, outputs, NULL);
295 |   gettimeofday(&stop_time, NULL);
296 |   printf("once run use %f ms\n", (__get_us(stop_time) - __get_us(start_time)) / 1000);
297 | 
298 |   printf("post process config: box_conf_threshold = %.2f, nms_threshold = %.2f\n", box_conf_threshold, nms_threshold);
299 | 
300 |   std::vector<float> out_scales;
301 |   std::vector<int32_t> out_zps;
302 |   for (int i = 0; i < app_ctx->io_num.n_output; ++i) {
303 |     out_scales.push_back(app_ctx->output_attrs[i].scale);
304 |     out_zps.push_back(app_ctx->output_attrs[i].zp);
305 |   }
306 | 
307 |   post_process((int8_t*)outputs[0].buf, (int8_t*)outputs[1].buf, (int8_t*)outputs[2].buf, model_height, model_width,
308 |                box_conf_threshold, nms_threshold, scale_w, scale_h, out_zps, out_scales, detect_result);
309 |   ret = rknn_outputs_release(ctx, app_ctx->io_num.n_output, outputs);
310 | 
311 |   if (resize_buf) {
312 |     free(resize_buf);
313 |   }
314 |   return 0;
315 | }
316 | 
317 | void mpp_decoder_frame_callback(void* userdata, int width_stride, int height_stride, int width, int height, int format, int fd, void* data) {
318 | 
319 |   rknn_app_context_t* ctx = (rknn_app_context_t*)userdata;
320 | 
321 |   int ret = 0;
322 |   static int frame_index = 0;
323 |   frame_index++;
324 | 
325 |   void* mpp_frame = NULL;
326 |   int mpp_frame_fd = 0;
327 |   void* mpp_frame_addr = NULL;
328 |   int enc_data_size;
329 | 
330 |   rga_buffer_t origin;
331 |   rga_buffer_t src;
332 | 
333 |   if (ctx->encoder == NULL) {
334 |     MppEncoder* mpp_encoder = new MppEncoder();
335 |     MppEncoderParams enc_params;
336 |     memset(&enc_params, 0, sizeof(MppEncoderParams));
337 |     enc_params.width = width;
338 |     enc_params.height = height;
339 |     enc_params.hor_stride = width_stride;
340 |     enc_params.ver_stride = height_stride;
341 |     enc_params.fmt = MPP_FMT_YUV420SP;
342 |     //enc_params.type = MPP_VIDEO_CodingHEVC;
343 |     //Note: rk3562只能支持h264格式的视频流
344 |     enc_params.type = MPP_VIDEO_CodingAVC;
345 |     mpp_encoder->Init(enc_params, NULL);
346 | 
347 |     ctx->encoder = mpp_encoder;
348 |   }
349 | 
350 |   int enc_buf_size = ctx->encoder->GetFrameSize();
351 |   char* enc_data = (char*)malloc(enc_buf_size);
352 | 
353 |   image_frame_t img;
354 |   img.width = width;
355 |   img.height = height;
356 |   img.width_stride = width_stride;
357 |   img.height_stride = height_stride;
358 |   img.fd = fd;
359 |   img.virt_addr = (char*)data;
360 |   img.format = RK_FORMAT_YCbCr_420_SP;
361 |   detect_result_group_t detect_result;
362 |   memset(&detect_result, 0, sizeof(detect_result_group_t));
363 | 
364 |   ret = inference_model(ctx, &img, &detect_result);
365 |   if (ret != 0) {
366 |     printf("inference model fail\n");
367 |     goto RET;
368 |   }
369 | 
370 |   mpp_frame = ctx->encoder->GetInputFrameBuffer();
371 |   mpp_frame_fd = ctx->encoder->GetInputFrameBufferFd(mpp_frame);
372 |   mpp_frame_addr = ctx->encoder->GetInputFrameBufferAddr(mpp_frame);
373 | 
374 |   // Copy To another buffer avoid to modify mpp decoder buffer
375 |   origin = wrapbuffer_fd(fd, width, height, RK_FORMAT_YCbCr_420_SP, width_stride, height_stride);
376 |   src = wrapbuffer_fd(mpp_frame_fd, width, height, RK_FORMAT_YCbCr_420_SP, width_stride, height_stride);
377 |   imcopy(origin, src);
378 | 
379 |   // Draw objects
380 |   for (int i = 0; i < detect_result.count; i++) {
381 |     detect_result_t* det_result = &(detect_result.results[i]);
382 |     printf("%s @ (%d %d %d %d) %f\n", det_result->name, det_result->box.left, det_result->box.top,
383 |            det_result->box.right, det_result->box.bottom, det_result->prop);
384 |     int x1 = det_result->box.left;
385 |     int y1 = det_result->box.top;
386 |     int x2 = det_result->box.right;
387 |     int y2 = det_result->box.bottom;
388 |     draw_rectangle_yuv420sp((unsigned char*)mpp_frame_addr, width_stride, height_stride, x1, y1, x2-x1+1, y2-y1+1, 0x00FF0000, 4);
389 |   }
390 | 
391 |   // Encode to file
392 |   // Write header on first frame
393 |   if (frame_index == 1) {
394 |     enc_data_size = ctx->encoder->GetHeader(enc_data, enc_buf_size);
395 |     fwrite(enc_data, 1, enc_data_size, ctx->out_fp);
396 |   }
397 |   memset(enc_data, 0, enc_buf_size);
398 |   enc_data_size = ctx->encoder->Encode(mpp_frame, enc_data, enc_buf_size);
399 |   fwrite(enc_data, 1, enc_data_size, ctx->out_fp);
400 | 
401 | RET:
402 |   if (enc_data != nullptr) {
403 |     free(enc_data);
404 |   }
405 | }
406 | 
407 | int process_video_file(rknn_app_context_t* ctx, const char* path)
408 | {
409 |   int video_size;
410 |   char* video_data = (char*)read_file_data(path, &video_size);
411 |   char* video_data_end = video_data + video_size;
412 |   printf("read video size=%d\n", video_size);
413 | 
414 |   const int SIZE = 8192;
415 |   char* video_data_ptr = video_data;
416 | 
417 |   do {
418 |       int pkt_eos = 0;
419 |       int size = SIZE;
420 |       if (video_data_ptr + size >= video_data_end) {
421 |           pkt_eos = 1;
422 |           size = video_data_end - video_data_ptr;
423 |       }
424 | 
425 |       ctx->decoder->Decode((uint8_t*)video_data_ptr, size, pkt_eos);
426 | 
427 |       video_data_ptr += size;
428 | 
429 |       if (video_data_ptr >= video_data_end) {
430 |           printf("reset decoder\n");
431 |           break;
432 |       }
433 | 
434 |       // LOGD("video_data_ptr=%p video_data_end=%p", video_data_ptr, video_data_end);
435 |       // usleep(10*1000);
436 |   } while (1);
437 | 
438 |   return 0;
439 | }
440 | 
441 | #if defined(BUILD_VIDEO_RTSP)
442 | void API_CALL on_track_frame_out(void *user_data, mk_frame frame) {
443 |   rknn_app_context_t *ctx = (rknn_app_context_t *) user_data;
444 |   printf("on_track_frame_out ctx=%p\n", ctx);
445 |   const char* data = mk_frame_get_data(frame);
446 |   size_t size = mk_frame_get_data_size(frame);
447 |   printf("decoder=%p\n", ctx->decoder);
448 |   ctx->decoder->Decode((uint8_t*)data, size, 0);
449 | }
450 | 
451 | void API_CALL on_mk_play_event_func(void *user_data, int err_code, const char *err_msg, mk_track tracks[],
452 |                                     int track_count) {
453 |   rknn_app_context_t *ctx = (rknn_app_context_t *) user_data;
454 |   if (err_code == 0) {
455 |       //success
456 |       printf("play success!");
457 |       int i;
458 |       for (i = 0; i < track_count; ++i) {
459 |           if (mk_track_is_video(tracks[i])) {
460 |               log_info("got video track: %s", mk_track_codec_name(tracks[i]));
461 |               //监听track数据回调
462 |               mk_track_add_delegate(tracks[i], on_track_frame_out, user_data);
463 |           }
464 |       }
465 |   } else {
466 |       printf("play failed: %d %s", err_code, err_msg);
467 |   }
468 | }
469 | 
470 | void API_CALL on_mk_shutdown_func(void *user_data, int err_code, const char *err_msg, mk_track tracks[], int track_count) {
471 |   printf("play interrupted: %d %s", err_code, err_msg);
472 | }
473 | 
474 | int process_video_rtsp(rknn_app_context_t* ctx, const char* url)
475 | {
476 |   mk_config config;
477 |   memset(&config, 0, sizeof(mk_config));
478 |   config.log_mask = LOG_CONSOLE;
479 |   mk_env_init(&config);
480 |   mk_player player = mk_player_create();
481 |   mk_player_set_on_result(player, on_mk_play_event_func, ctx);
482 |   mk_player_set_on_shutdown(player, on_mk_shutdown_func, ctx);
483 |   mk_player_play(player, url);
484 | 
485 |   printf("enter any key to exit\n");
486 |   getchar();
487 | 
488 |   if (player) {
489 |       mk_player_release(player);
490 |   }
491 |   return 0;
492 | }
493 | #endif
494 | 
495 | /*-------------------------------------------
496 |                   Main Functions
497 | -------------------------------------------*/
498 | int main(int argc, char** argv)
499 | {
500 |   int status = 0;
501 |   int ret;
502 | 
503 |   if (argc != 4) {
504 |     printf("Usage: %s <rknn_model> <video_path> <video_type 264/265> \n", argv[0]);
505 |     return -1;
506 |   }
507 | 
508 |   char* model_name = (char*)argv[1];
509 |   char* video_name = argv[2];
510 |   int video_type = atoi(argv[3]);
511 | 
512 |   rknn_app_context_t app_ctx;
513 |   memset(&app_ctx, 0, sizeof(rknn_app_context_t));
514 | 
515 |   ret = init_model(model_name, &app_ctx);
516 |   if (ret != 0) {
517 |     printf("init model fail\n");
518 |     return -1;
519 |   }
520 | 
521 |   if (app_ctx.decoder == NULL) {
522 |     MppDecoder* decoder = new MppDecoder();
523 |     decoder->Init(video_type, 30, &app_ctx);
524 |     decoder->SetCallback(mpp_decoder_frame_callback);
525 |     app_ctx.decoder = decoder;
526 |   }
527 | 
528 |   if (app_ctx.out_fp == NULL) {
529 |     FILE* fp = fopen(OUT_VIDEO_PATH, "w");
530 |     if(fp == NULL) {
531 |         printf("open %s error\n", OUT_VIDEO_PATH);
532 |         return -1;
533 |     }
534 |     app_ctx.out_fp = fp;
535 |   }
536 | 
537 |   printf("app_ctx=%p decoder=%p\n", &app_ctx, app_ctx.decoder);
538 | 
539 |   if (strncmp(video_name, "rtsp", 4) == 0) {
540 | #if defined(BUILD_VIDEO_RTSP)
541 |         process_video_rtsp(&app_ctx, video_name);
542 | #else
543 |         printf("rtsp no support\n");
544 | #endif
545 |   } else {
546 |     process_video_file(&app_ctx, video_name);
547 |   }
548 | 
549 |   printf("waiting finish\n");
550 |   usleep(3*1000*1000);
551 | 
552 |   // release
553 |   fflush(app_ctx.out_fp);
554 |   fclose(app_ctx.out_fp);
555 | 
556 |   if (app_ctx.decoder != nullptr) {
557 |     delete(app_ctx.decoder);
558 |     app_ctx.decoder = nullptr;
559 |   }
560 |   if (app_ctx.encoder != nullptr) {
561 |     delete(app_ctx.encoder);
562 |     app_ctx.encoder = nullptr;
563 |   }
564 | 
565 |   release_model(&app_ctx);
566 | 
567 |   return 0;
568 | }
569 | 


--------------------------------------------------------------------------------
/unit_test/test_mpp_video_decoder.cpp:
--------------------------------------------------------------------------------
 1 | /**
 2 |  * @author: bo.liu
 3 |  * @mail: geniusrabbit@qq.com
 4 |  * @date: 2023.08.09
 5 |  * @brief: Rockchip MPP 视频解码测试
 6 |  */
 7 | #include <string>
 8 | #include "mpp_video_decoder.h"
 9 | #include "utils_log.h"
10 | #include "utils.h"
11 | #include "mpp_video_encoder.h"
12 | 
13 | #include "mpp_video_utils.h"
14 | 
15 | int test_decoder(){
16 |     // 初始化解码器
17 |     std::string path = "1080p_ffmpeg.h264";
18 |     MppVideoDecoder video_decoder = MppVideoDecoder(path);
19 |     if(!video_decoder.is_init()){
20 |         d_unit_test_error("MppVideoDecoder init failed!")
21 |         return -1;
22 |     }
23 | 
24 |     // 获取视频的下一帧数据
25 |     uint64_t time_start = get_time_of_ms();
26 |     uint32_t frame_count = 0;
27 |     DecoderMppFrame frame{};
28 |     while(video_decoder.get_next_frame(frame) == 0){
29 |         frame_count++;
30 |         d_unit_test_info("frame count: %d", frame_count)
31 | 
32 | //        cv::Mat image;
33 | //        frame_data_to_mat(
34 | //                (uint8_t *) frame.data_buf,
35 | //                frame.hor_width, frame.ver_height,
36 | //                frame.hor_stride, frame.ver_stride,
37 | //                frame.mpp_frame_format,
38 | //                image);
39 | //        cv::imwrite("1080p_ffmpeg.jpg", image);
40 |         video_decoder.release_frame(frame);
41 | //        break;
42 |     }
43 |     uint64_t time_end = get_time_of_ms();
44 |     d_unit_test_warn("frame count: %d", frame_count)
45 |     d_unit_test_warn("time cost: %d ms", time_end - time_start)
46 | 
47 |     return 0;
48 | }
49 | 
50 | int test_encoder(){
51 |     // 初始化解码器
52 |     std::string path = "1080p_ffmpeg.h264";
53 |     MppVideoDecoder video_decoder = MppVideoDecoder(path);
54 |     if(!video_decoder.is_init()){
55 |         d_unit_test_error("MppVideoDecoder init failed!")
56 |         return -1;
57 |     }
58 |     std::string output_path = "1080p_ffmpeg_encoder.h264";
59 |     MppVideoEncoder video_encoder = MppVideoEncoder(output_path);
60 |     if(!video_encoder.is_init()){
61 |         d_unit_test_error("MppVideoEncoder init failed!")
62 |         return -1;
63 |     }
64 | 
65 |     // 获取视频的下一帧数据
66 |     uint64_t time_start = get_time_of_ms();
67 |     uint32_t frame_count = 0;
68 |     DecoderMppFrame frame{};
69 |     while(video_decoder.get_next_frame(frame) == 0){
70 |         frame_count++;
71 |         d_unit_test_info("frame count: %d", frame_count)
72 | 
73 |         uint8_t image_data[frame.data_size];
74 |         yuv_del_stride(
75 |                 (uint8_t *) frame.data_buf,
76 |                 frame.hor_width, frame.ver_height,
77 |                 frame.hor_stride, frame.ver_stride,
78 |                 frame.mpp_frame_format,
79 |                 image_data);
80 |         video_encoder.process_image(image_data, frame.hor_width, frame.ver_height, frame.mpp_frame_format);
81 | 
82 |         video_decoder.release_frame(frame);
83 |     }
84 |     uint64_t time_end = get_time_of_ms();
85 |     d_unit_test_warn("frame count: %d", frame_count)
86 |     d_unit_test_warn("time cost: %d ms", time_end - time_start)
87 |     return 0;
88 | }
89 | 
90 | int main(){
91 | //    test_decoder();
92 | 
93 |     test_encoder();
94 |     return 0;
95 | }


--------------------------------------------------------------------------------
/unit_test/test_rknn_model.cpp:
--------------------------------------------------------------------------------
  1 | /**
  2 |  * @author: bo.liu
  3 |  * @mail: geniusrabbit@qq.com
  4 |  * @date: 2023.08.04
  5 |  * @brief: RKNN model 的测试，参考样例 https://github.com/rockchip-linux/rknpu2/tree/master/examples/rknn_mobilenet_demo
  6 |  */
  7 | #include <string>
  8 | #include "rknn_model.h"
  9 | #include "utils_log.h"
 10 | #include "rknn_infer_api.h"
 11 | 
 12 | #include "opencv2/core/core.hpp"
 13 | #include "opencv2/imgcodecs.hpp"
 14 | #include "opencv2/imgproc.hpp"
 15 | 
 16 | const int MODEL_IN_WIDTH    = 224;
 17 | const int MODEL_IN_HEIGHT   = 224;
 18 | const int MODEL_IN_CHANNELS = 3;
 19 | 
 20 | static unsigned char* load_model(const char* filename, int* model_size)
 21 | {
 22 |     FILE* fp = fopen(filename, "rb");
 23 |     if (fp == nullptr) {
 24 |         printf("fopen %s fail!\n", filename);
 25 |         return NULL;
 26 |     }
 27 |     fseek(fp, 0, SEEK_END);
 28 |     int            model_len = ftell(fp);
 29 |     unsigned char* model     = (unsigned char*)malloc(model_len);
 30 |     fseek(fp, 0, SEEK_SET);
 31 |     if (model_len != fread(model, 1, model_len, fp)) {
 32 |         printf("fread %s fail!\n", filename);
 33 |         free(model);
 34 |         return NULL;
 35 |     }
 36 |     *model_size = model_len;
 37 |     if (fp) {
 38 |         fclose(fp);
 39 |     }
 40 |     return model;
 41 | }
 42 | 
 43 | static void dump_tensor_attr(rknn_tensor_attr* attr){
 44 |     printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
 45 |            "zp=%d, scale=%f\n",
 46 |            attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
 47 |            attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
 48 |            get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
 49 | }
 50 | 
 51 | static int rknn_plugin_GetTop(const float* pfProb, float* pfMaxProb, uint32_t* pMaxClass, uint32_t outputCount, uint32_t topNum){
 52 |     uint32_t i, j;
 53 | 
 54 | #define MAX_TOP_NUM 20
 55 |     if (topNum > MAX_TOP_NUM)
 56 |         return 0;
 57 | 
 58 |     memset(pfMaxProb, 0, sizeof(float) * topNum);
 59 |     memset(pMaxClass, 0xff, sizeof(float) * topNum);
 60 | 
 61 |     for (j = 0; j < topNum; j++) {
 62 |         for (i = 0; i < outputCount; i++) {
 63 |             if ((i == *(pMaxClass + 0)) || (i == *(pMaxClass + 1)) || (i == *(pMaxClass + 2)) || (i == *(pMaxClass + 3)) ||
 64 |                 (i == *(pMaxClass + 4))) {
 65 |                 continue;
 66 |             }
 67 | 
 68 |             if (pfProb[i] > *(pfMaxProb + j)) {
 69 |                 *(pfMaxProb + j) = pfProb[i];
 70 |                 *(pMaxClass + j) = i;
 71 |             }
 72 |         }
 73 |     }
 74 |     return 1;
 75 | }
 76 | 
 77 | int test_rknn_model(){
 78 |     std::string model_path = "model/RK3566_RK3568/mobilenet_v1.rknn";
 79 |     std::string image_path   = "model/dog_224x224.jpg";
 80 | 
 81 |     // Load image
 82 |     cv::Mat orig_img = imread(image_path, cv::IMREAD_COLOR);
 83 |     if (!orig_img.data) {
 84 |         printf("cv::imread %s fail!\n", image_path.c_str());
 85 |         return -1;
 86 |     }
 87 | 
 88 |     cv::Mat orig_img_rgb;
 89 |     cv::cvtColor(orig_img, orig_img_rgb, cv::COLOR_BGR2RGB);
 90 | 
 91 |     cv::Mat img = orig_img_rgb.clone();
 92 |     if (orig_img.cols != MODEL_IN_WIDTH || orig_img.rows != MODEL_IN_HEIGHT) {
 93 |         printf("resize %d %d to %d %d\n", orig_img.cols, orig_img.rows, MODEL_IN_WIDTH, MODEL_IN_HEIGHT);
 94 |         cv::resize(orig_img, img, cv::Size(MODEL_IN_WIDTH, MODEL_IN_HEIGHT), 0, 0, cv::INTER_LINEAR);
 95 |     }
 96 | 
 97 |     PluginConfigSet m_plugin_set_config{};
 98 |     RknnModel model(model_path, m_plugin_set_config, true);
 99 | 
100 |     auto *input_unit = new InputUnit();
101 |     input_unit->n_inputs = 1;
102 |     input_unit->inputs = (rknn_input*)malloc(input_unit->n_inputs * sizeof(rknn_input));
103 |     memset(input_unit->inputs, 0, input_unit->n_inputs * sizeof(rknn_input));
104 | 
105 |     input_unit->inputs[0].index = 0;
106 |     input_unit->inputs[0].type  = RKNN_TENSOR_UINT8;
107 |     input_unit->inputs[0].size  = img.cols * img.rows * img.channels() * sizeof(uint8_t);
108 |     input_unit->inputs[0].fmt   = RKNN_TENSOR_NHWC;
109 |     input_unit->inputs[0].buf = new uint8_t[img.cols * img.rows * img.channels()];
110 |     memcpy(input_unit->inputs[0].buf, img.data, input_unit->inputs[0].size);
111 | 
112 |     auto *output_unit = new OutputUnit();
113 |     output_unit->n_outputs = 1;
114 |     output_unit->outputs = (rknn_output*)malloc(output_unit->n_outputs * sizeof(rknn_output));
115 |     memset(output_unit->outputs, 0, output_unit->n_outputs * sizeof(rknn_output));
116 |     output_unit->outputs[0].want_float = 1;
117 | 
118 |     model.model_infer_sync(input_unit->n_inputs, input_unit->inputs, output_unit->n_outputs, output_unit->outputs);
119 | 
120 |     for (int i = 0; i < 1; i++) {
121 |         uint32_t MaxClass[5];
122 |         float    fMaxProb[5];
123 |         auto*   buffer = (float*)output_unit->outputs[i].buf;
124 |         uint32_t sz     = output_unit->outputs[i].size / 4;
125 | 
126 |         rknn_plugin_GetTop(buffer, fMaxProb, MaxClass, sz, 5);
127 | 
128 |         printf(" --- Top5 ---\n");
129 |         for (int j = 0; j < 5; j++) {
130 |             printf("%3d: %8.6f\n", MaxClass[j], fMaxProb[j]);
131 |         }
132 |     }
133 |     model.model_infer_release(output_unit->n_outputs, output_unit->outputs);
134 |     return 0;
135 | }
136 | 
137 | int main(){
138 |     test_rknn_model();
139 |     return 0;
140 | }


--------------------------------------------------------------------------------