├── CMakeLists.txt
├── README.md
├── RKNNModel.cpp
├── RKNNModel.h
├── T_model_backbone.rknn
├── X_model_backbone.rknn
├── main.cpp
├── model_head.rknn
├── nanotrack.cpp
└── nanotrack.hpp


/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | cmake_minimum_required(VERSION 3.10)
 2 | 
 3 | project(nanotrack)
 4 | 
 5 | # 指定交叉编译工具链
 6 | set(CMAKE_SYSTEM_NAME Linux)
 7 | set(CMAKE_SYSTEM_PROCESSOR aarch64)
 8 | 
 9 | set(CMAKE_C_COMPILER aarch64-linux-gnu-gcc)
10 | set(CMAKE_CXX_COMPILER aarch64-linux-gnu-g++)
11 | 
12 | # 设置 OpenCV 路径
13 | set(OpenCV_DIR "/home/xcc/workcode/NanoTrack/opencv/build/install/lib/cmake/opencv4")
14 | 
15 | # 查找 OpenCV
16 | find_package(OpenCV REQUIRED)
17 | if (NOT OpenCV_FOUND)
18 |     message(FATAL_ERROR "OpenCV not found!")
19 | endif()
20 | 
21 | if (ENABLE_ASAN)
22 |     message(STATUS "BUILD WITH ADDRESS SANITIZER")
23 |     set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address")
24 |     set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address")
25 |     set(CMAKE_LINKER_FLAGS_DEBUG "${CMAKE_LINKER_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address")
26 | endif()
27 | 
28 | set(rknpu_tracker_file nanotrack.cpp)
29 | 
30 | if (TARGET_SOC STREQUAL "rv1106" OR TARGET_SOC STREQUAL "rv1103")
31 |     add_definitions(-DRV1106_1103)
32 | endif()
33 | 
34 | if(TARGET_SOC STREQUAL "rk1808" OR TARGET_SOC STREQUAL "rv1109" OR TARGET_SOC STREQUAL "rv1126")
35 |     add_definitions(-DRKNPU1)
36 | endif()
37 | 
38 | # 设置RKNN包含路径和库路径
39 | set(RKNN_INCLUDE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/rknpu2/include")
40 | set(RKNN_LIB_PATH "${CMAKE_CURRENT_SOURCE_DIR}/rknpu2/Linux/aarch64")
41 | 
42 | # 设置RGA包含路径
43 | set(RGA_INCLUDE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/librga/include")
44 | 
45 | # 设置CMake的安装路径
46 | set(CMAKE_INSTALL_RPATH "$ORIGIN/../lib")
47 | 
48 | # 查找所有的源文件
49 | file(GLOB SRCS ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp)
50 | 
51 | # 添加可执行文件
52 | add_executable(${PROJECT_NAME}
53 |     main.cpp
54 |     RKNNModel.cpp
55 |     ${rknpu_tracker_file}
56 | )
57 | 
58 | # 链接必要的库文件
59 | target_link_libraries(${PROJECT_NAME}
60 |     ${OpenCV_LIBS}
61 |     ${RKNN_LIB_PATH}/librknnrt.so
62 |     dl
63 | )
64 | 
65 | if (CMAKE_SYSTEM_NAME STREQUAL "Android")
66 |     target_link_libraries(${PROJECT_NAME}
67 |         log
68 |     )
69 | endif()
70 | 
71 | message(STATUS "!!!!!!!!!!!CMAKE_SYSTEM_NAME: ${CMAKE_SYSTEM_NAME}")
72 | if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
73 |     set(THREADS_PREFER_PTHREAD_FLAG ON)
74 |     find_package(Threads REQUIRED)
75 |     target_link_libraries(${PROJECT_NAME} Threads::Threads)
76 | endif()
77 | 
78 | # 设置包含目录
79 | target_include_directories(${PROJECT_NAME} PRIVATE
80 |     ${CMAKE_CURRENT_SOURCE_DIR}
81 |     ${RKNN_INCLUDE_PATH}
82 |     ${RGA_INCLUDE_PATH}  # 添加 RGA 头文件目录
83 |     ${OpenCV_INCLUDE_DIRS}
84 | )
85 | 
86 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # nanotrack-c-rk3588
2 | nanotrack c++版本部署rk3588
3 | 可在rk3588上进行目标跟踪，速度达125fps
4 | 


--------------------------------------------------------------------------------
/RKNNModel.cpp:
--------------------------------------------------------------------------------
  1 | #include "RKNNModel.h"
  2 | #include <iostream>
  3 | using namespace std;
  4 | 
  5 | 
  6 | RKNNModel::RKNNModel():pModel(nullptr),ctx(0) {}
  7 | RKNNModel::~RKNNModel()
  8 | {
  9 |     if (this->pModel)
 10 |     {
 11 |         free(this->pModel);
 12 |         this->pModel = nullptr;
 13 |     }
 14 |     if (this->ctx > 0)
 15 |     {
 16 |         rknn_destroy(ctx);
 17 |         ctx = 0;
 18 |     }
 19 | }
 20 | 
 21 | int RKNNModel::runRKNN(vector<vector<float>> &output, void *input_data, uint32_t input_size, rknn_tensor_type input_type, bool pass_through)
 22 | {
 23 |     string logMsg;
 24 | 
 25 |     rknn_input rknnInputs[1];
 26 |     rknnInputs[0].index = 0;
 27 |     rknnInputs[0].buf = input_data;
 28 |     rknnInputs[0].size = input_size;
 29 |     rknnInputs[0].pass_through = pass_through;
 30 |     rknnInputs[0].fmt = RKNN_TENSOR_NHWC;
 31 |     // rknn_tensor_type: RKNN_TENSOR_UINT8 / RKNN_TENSOR_FLOAT32
 32 |     rknnInputs[0].type = input_type;
 33 | 
 34 | 
 35 | 
 36 |     // input
 37 |     int ret = rknn_inputs_set(this->ctx, 1, rknnInputs);
 38 | 
 39 |     if (ret < 0)
 40 |     {
 41 |         logMsg = "rknn_input_set failed! ret=" + to_string(ret);
 42 |         
 43 |         return -1;
 44 |     }
 45 | 
 46 |     // run
 47 |     ret = rknn_run(this->ctx, nullptr);
 48 |     if (ret < 0)
 49 |     {
 50 |         logMsg = "rknn_run failed! ret=" + to_string(ret);
 51 |         
 52 |         return -1;
 53 |     }
 54 | 
 55 | 
 56 |     // infer output length
 57 |     int outputLength = this->outputsAttr.size();
 58 |     if (outputLength < 1)
 59 |     {
 60 |         logMsg = "outputsAttr is empty!";
 61 |         
 62 |         return -1;
 63 |     }
 64 | 
 65 |     // get output
 66 |     rknn_output *rknnOutputs = new rknn_output[outputLength];
 67 |     memset(rknnOutputs, 0, sizeof(rknn_output)*outputLength);
 68 | 
 69 |     // rknn_output rknnOutputs[1];
 70 |     for (int out_i = 0; out_i < outputLength; out_i++)
 71 |     {
 72 |         rknnOutputs[out_i].want_float = true;
 73 |         rknnOutputs[out_i].is_prealloc = false;
 74 |     }
 75 | 
 76 |     ret = rknn_outputs_get(this->ctx, outputLength, rknnOutputs, nullptr);
 77 |     if (ret < 0)
 78 |     {
 79 |         logMsg = "rknn_outputs_get failed! ret=" + to_string(ret);
 80 |         
 81 |         rknn_outputs_release(this->ctx, outputLength, rknnOutputs);
 82 |         delete rknnOutputs;
 83 |         return -1;
 84 |     }
 85 | 
 86 |     // set output
 87 |     output.resize(outputLength);
 88 |     for (int out_i = 0; out_i < outputLength; out_i++)
 89 |     {
 90 |         // cout << "n_elems=" << this->outputsAttr[out_i].n_elems << ", size=" << this->outputsAttr[out_i].size << endl;
 91 |         if (rknnOutputs[out_i].size == this->outputsAttr[out_i].n_elems * sizeof(float))
 92 |         {
 93 |             float *out_arr = (float *)rknnOutputs[out_i].buf;
 94 |             output[out_i] = vector<float>(out_arr, out_arr + this->outputsAttr[out_i].n_elems);
 95 |         }
 96 |         else
 97 |         {
 98 |             output.clear();
 99 |             logMsg = "rknn_outputs_get #" + to_string(out_i) + " of " + to_string(outputLength) + " failed! get_outputs_size=" + to_string(this->outputsAttr[out_i].size) + ", but expect " + to_string(this->outputsAttr[out_i].n_elems * sizeof(float));
100 |             
101 |             rknn_outputs_release(this->ctx, outputLength, rknnOutputs);
102 |             delete rknnOutputs;
103 |             return -1;
104 |         }
105 |     }
106 | 
107 |     // release resources
108 |     rknn_outputs_release(this->ctx, outputLength, rknnOutputs);
109 |     delete rknnOutputs;
110 |     
111 |     return 0;
112 | }
113 | 
114 | 
115 | int RKNNModel::runRKNN(vector<vector<float>> &output, void *input_data1, uint32_t input_size1, void *input_data2, uint32_t input_size2,rknn_tensor_type input_type, bool pass_through)
116 | {
117 |     string logMsg;
118 |     string funcName = "runRKNN:" + this->modelName;
119 | 
120 |     rknn_input rknnInputs[2];
121 |     rknnInputs[0].index = 0;
122 |     rknnInputs[0].buf = input_data1;
123 |     rknnInputs[0].size = input_size1;
124 |     rknnInputs[0].pass_through = pass_through;
125 |     rknnInputs[0].fmt = RKNN_TENSOR_NHWC;
126 |     // rknn_tensor_type: RKNN_TENSOR_UINT8 / RKNN_TENSOR_FLOAT32
127 |     rknnInputs[0].type = input_type;
128 | 
129 | 	rknnInputs[1].index = 1;
130 |     rknnInputs[1].buf = input_data2;
131 |     rknnInputs[1].size = input_size2;
132 |     rknnInputs[1].pass_through = pass_through;
133 |     rknnInputs[1].fmt = RKNN_TENSOR_NHWC;
134 |     // rknn_tensor_type: RKNN_TENSOR_UINT8 / RKNN_TENSOR_FLOAT32
135 |     rknnInputs[1].type = input_type;
136 | 
137 |     // input
138 |     int ret = rknn_inputs_set(this->ctx, 2, rknnInputs);
139 |     if (ret < 0)
140 |     {
141 |         logMsg = "rknn_input_set failed! ret=" + to_string(ret);
142 |         // 
143 |         return -1;
144 |     }
145 | 
146 |     // run
147 |     ret = rknn_run(this->ctx, nullptr);
148 |     if (ret < 0)
149 |     {
150 |         logMsg = "rknn_run failed! ret=" + to_string(ret);
151 |         
152 |         return -1;
153 |     }
154 | 
155 |     // infer output length
156 |     int outputLength = this->outputsAttr.size();
157 |     if (outputLength < 1)
158 |     {
159 |         logMsg = "outputsAttr is empty!";
160 |         
161 |         return -1;
162 |     }
163 | 
164 |     // get output
165 |     rknn_output *rknnOutputs = new rknn_output[outputLength];
166 | 
167 |     memset(rknnOutputs, 0, sizeof(rknn_output)*outputLength);
168 | 
169 |     for (int out_i = 0; out_i < outputLength; out_i++)
170 |     {
171 |         rknnOutputs[out_i].want_float = true;
172 |         rknnOutputs[out_i].is_prealloc = false;
173 |     }
174 |     ret = rknn_outputs_get(this->ctx, outputLength, rknnOutputs, nullptr);
175 |     if (ret < 0)
176 |     {
177 |         logMsg = "rknn_outputs_get failed! ret=" + to_string(ret);
178 |         
179 |         rknn_outputs_release(this->ctx, outputLength, rknnOutputs);
180 |         delete rknnOutputs;
181 |         return -1;
182 |     }
183 | 
184 |     // set output
185 |     output.resize(outputLength);
186 |     for (int out_i = 0; out_i < outputLength; out_i++)
187 |     {
188 |         // cout << "n_elems=" << this->outputsAttr[out_i].n_elems << ", size=" << this->outputsAttr[out_i].size << endl;
189 |         // for (size_t i = 0; i < this->outputsAttr[out_i].n_dims; i++)
190 |         // {
191 |         //    cout << "dims[" << i << "]=" << this->outputsAttr[out_i].dims[i] << endl;
192 |         // }
193 |         
194 |         if (rknnOutputs[out_i].size == this->outputsAttr[out_i].n_elems * sizeof(float))
195 |         {
196 |             float *out_arr = (float *)rknnOutputs[out_i].buf;
197 |             output[out_i] = vector<float>(out_arr, out_arr + this->outputsAttr[out_i].n_elems);
198 |         }
199 |         else
200 |         {
201 |             output.clear();
202 |             logMsg = "rknn_outputs_get #" + to_string(out_i) + " of " + to_string(outputLength) + " failed! get_outputs_size=" + to_string(this->outputsAttr[out_i].size) + ", but expect " + to_string(this->outputsAttr[out_i].n_elems * sizeof(float));
203 |             
204 |             rknn_outputs_release(this->ctx, outputLength, rknnOutputs);
205 |             delete rknnOutputs;
206 |             return -1;
207 |         }
208 |     }
209 | 
210 |     // release resources
211 |     rknn_outputs_release(this->ctx, outputLength, rknnOutputs);
212 |     delete rknnOutputs;
213 |     
214 |     return 0;
215 | }
216 | 
217 | 
218 | int RKNNModel::loadRKNN(string modelPath, int outputLength, string modelName)
219 | {
220 |     if (modelName != "")
221 |     {
222 |         this->modelName = modelName;
223 |     }
224 |     string logMsg;
225 | 
226 |     int modelLength = -1;
227 |     try
228 |     {
229 |         // if (!check_exist(modelPath))
230 |         // {
231 |         //     logMsg = "modelPath not exist,  " + modelPath;
232 |             
233 |         //     return RKNN_FILE_INVALID;
234 |         // }
235 |         FILE *modelFP = fopen(modelPath.c_str(), "rb");
236 |         if (modelFP == NULL)
237 |         {
238 |             logMsg = "fopen fail! " + modelPath;
239 |             
240 |             this->releaseRKNN();
241 |             return -1;
242 |         }
243 |         fseek(modelFP, 0, SEEK_END);
244 |         modelLength = ftell(modelFP);
245 |         this->pModel = malloc(modelLength);
246 |         fseek(modelFP, 0, SEEK_SET);
247 |         if (modelLength != fread(this->pModel, 1, modelLength, modelFP))
248 |         {
249 |             logMsg = "fread fail! " + modelPath;
250 |             
251 |             fclose(modelFP);
252 |             this->releaseRKNN();
253 |             return -1;
254 |         }
255 |         fclose(modelFP);
256 |     }
257 |     catch (...)
258 |     {
259 |         logMsg = "load rknn fail! exception caught! " + modelPath;
260 |         
261 |         return -1;
262 |     }
263 | 
264 |     int ret = rknn_init(&(this->ctx), this->pModel, modelLength, 0, nullptr); //| RKNN_FLAG_COLLECT_PERF_MASK
265 |     ret |= rknn_set_core_mask(this->ctx, RKNN_NPU_CORE_0_1_2);
266 |     // printf("this->ctx: %ld\n",this->ctx);
267 |     if (ret < 0)
268 |     {
269 |         logMsg = "rknn_init fail! ret=" + to_string(ret);
270 |         
271 |         this->releaseRKNN();
272 |         return -1;
273 |     }
274 | 
275 |     // output attribute setting
276 |     for (int iOut = 0; iOut < outputLength; iOut++)
277 |     {
278 |         rknn_tensor_attr modelOutput;
279 |         memset(&modelOutput, 0, sizeof(rknn_tensor_attr));
280 |         modelOutput.index = iOut;
281 |         ret = rknn_query(this->ctx, RKNN_QUERY_OUTPUT_ATTR, &modelOutput, sizeof(rknn_tensor_attr));
282 |         if (ret < 0)
283 |         {
284 |             logMsg = "rknn_query failed #" + to_string(iOut) + " of " + to_string(outputLength) + ", ret=" + to_string(ret);
285 |             
286 |             this->releaseRKNN();
287 |             return -1;
288 |         }
289 |         this->outputsAttr.push_back(modelOutput);
290 | 
291 |         rknn_tensor_attr *attr = &modelOutput;
292 |         std::string shape_str = attr->n_dims < 1 ? "" : std::to_string(attr->dims[0]);
293 |         for (int i = 1; i < attr->n_dims; ++i)
294 |         {
295 |             shape_str += ", " + std::to_string(attr->dims[i]);
296 |         }
297 |     }
298 | 
299 |     rknn_input_output_num io_num;
300 |     ret = rknn_query(this->ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
301 |     if (ret < 0)
302 |     {
303 |         printf("rknn_init error ret=%d\n", ret);
304 |         return -1;
305 |     }
306 |     // printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);
307 | 
308 |     rknn_tensor_attr input_attrs[io_num.n_input];
309 |     memset(input_attrs, 0, sizeof(input_attrs));
310 |     for (int i = 0; i < io_num.n_input; i++)
311 |     {
312 |         input_attrs[i].index = i;
313 |         ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
314 |         if (ret < 0)
315 |         {
316 |         printf("rknn_init error ret=%d\n", ret);
317 |         return -1;
318 |         }
319 |     }
320 | 
321 |     
322 | 
323 |     // get sdk version
324 |     rknn_sdk_version version;
325 |     ret = rknn_query(this->ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version));
326 |     if (ret < 0)
327 |     {
328 |         logMsg = "rknn_query sdk version failed, , ret=" + to_string(ret);
329 |         
330 |         this->releaseRKNN();
331 |         return -1;
332 |     }
333 |     string api_version(version.api_version);
334 |     string drv_version(version.drv_version);
335 |     cout << modelPath << endl;
336 |     cout << logMsg << endl;
337 | 
338 |     return 0;
339 | }
340 | 
341 | int RKNNModel::releaseRKNN()
342 | {
343 |     if (this->ctx > 0)
344 |     {
345 |         int ret = rknn_destroy(this->ctx);
346 |         if (ret < 0)
347 |         {
348 |             return ret;
349 |         }
350 |         this->ctx = 0;
351 |     }
352 |     if (this->pModel)
353 |     {
354 |         free(this->pModel);
355 |         this->pModel = nullptr;
356 |     }
357 |     if (!this->outputsAttr.empty())
358 |         this->outputsAttr.clear();
359 |     this->modelName = "UndefinedModel";
360 |     return 0;
361 | }
362 | 


--------------------------------------------------------------------------------
/RKNNModel.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | #include <string.h>
 3 | #include <string>
 4 | #include <vector>
 5 | #include "rknn_api.h"
 6 | 
 7 | class RKNNModel
 8 | {
 9 | public:
10 |     RKNNModel();
11 |     ~RKNNModel();
12 | 
13 |     std::vector<rknn_tensor_attr> outputsAttr;
14 | 
15 |     int runRKNN(std::vector<std::vector<float>> &output, void *input_data, uint32_t input_size, rknn_tensor_type input_type, bool pass_through = false);
16 |     int runRKNN(std::vector<std::vector<float>> &output, void *input_data1, uint32_t input_size1, void *input_data2, uint32_t input_size2, rknn_tensor_type input_type, bool pass_through);
17 | 
18 |     int loadRKNN(std::string modelPath, int outputLength, std::string modelName="");
19 | 
20 |     int releaseRKNN();
21 | 
22 | private:
23 |     void *pModel;
24 |     rknn_context ctx;
25 |     std::string modelName;  // 添加这个成员变量
26 | };
27 | 


--------------------------------------------------------------------------------
/T_model_backbone.rknn:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/cab0625/nanotrack-c-rk3588/47b2706655e889b5b15e282df4c24fed4a5935b6/T_model_backbone.rknn


--------------------------------------------------------------------------------
/X_model_backbone.rknn:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/cab0625/nanotrack-c-rk3588/47b2706655e889b5b15e282df4c24fed4a5935b6/X_model_backbone.rknn


--------------------------------------------------------------------------------
/main.cpp:
--------------------------------------------------------------------------------
 1 | #include "nanotrack.hpp"
 2 | #include "RKNNModel.h"
 3 | #include <opencv2/opencv.hpp>
 4 | #include <iostream>
 5 | 
 6 | using namespace cv;
 7 | using namespace std;
 8 | 
 9 | int main(int argc, char** argv) {
10 |     // 检查命令行参数
11 |     if (argc < 2) {
12 |         cout << "Usage: " << argv[0] << " <video_file_path>" << endl;
13 |         return -1;
14 |     }
15 | 
16 |     string video_name = argv[1];
17 | 
18 |     // 加载模型
19 |     NanoTrack tracker;
20 |     tracker.load_model("models/model_T.rknn", "models/model_X.rknn", "models/model_head.rknn");
21 | 
22 |     // 打开视频文件
23 |     VideoCapture cap(video_name);
24 |     if (!cap.isOpened()) {
25 |         cerr << "Error: Unable to open video file " << video_name << endl;
26 |         return -1;
27 |     }
28 | 
29 |     // 获取视频帧率和帧尺寸
30 |     double fps = cap.get(CAP_PROP_FPS);
31 |     int width = static_cast<int>(cap.get(CAP_PROP_FRAME_WIDTH));
32 |     int height = static_cast<int>(cap.get(CAP_PROP_FRAME_HEIGHT));
33 | 
34 |     // 创建视频写入器
35 |     VideoWriter video_writer("output_video.avi", VideoWriter::fourcc('X', 'V', 'I', 'D'), fps, Size(width, height));
36 | 
37 |     // 读取第一帧并初始化
38 |     Mat frame;
39 |     cap >> frame;
40 |     if (frame.empty()) {
41 |         cerr << "Error: Unable to read first frame from video file " << video_name << endl;
42 |         return -1;
43 |     }
44 | 
45 |     Rect bbox(1104, 373, 60, 41); // 初始边界框
46 |     tracker.init(frame, bbox);
47 | 
48 |     // 追踪
49 |     while (true) {
50 |         cap >> frame;
51 |         if (frame.empty()) {
52 |             break;
53 |         }
54 | 
55 |         double t1 = getTickCount();
56 |         float score = tracker.track(frame);
57 |         double t2 = getTickCount();
58 |         double process_time_ms = (t2 - t1) * 1000 / getTickFrequency();
59 |         double fps_value = getTickFrequency() / (t2 - t1);
60 |         cout << "每帧处理时间: " << process_time_ms << " ms, FPS: " << fps_value << endl;
61 | 
62 |         // 绘制边界框
63 |         rectangle(frame, tracker.state.target_pos, tracker.state.target_sz, Scalar(0, 255, 0), 2);
64 | 
65 |         // 写入视频
66 |         video_writer.write(frame);
67 | 
68 |         // 显示追踪结果
69 |         imshow("Tracking", frame);
70 |         if (waitKey(30) == 27) { // 按下Esc键退出
71 |             break;
72 |         }
73 |     }
74 | 
75 |     // 释放资源
76 |     video_writer.release();
77 |     cap.release();
78 |     destroyAllWindows();
79 | 
80 |     return 0;
81 | }
82 | 
83 | 


--------------------------------------------------------------------------------
/model_head.rknn:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/cab0625/nanotrack-c-rk3588/47b2706655e889b5b15e282df4c24fed4a5935b6/model_head.rknn


--------------------------------------------------------------------------------
/nanotrack.cpp:
--------------------------------------------------------------------------------
  1 | #include <iostream>
  2 | #include <cstdlib>
  3 | #include <string>
  4 | #include "nanotrack.hpp"
  5 | 
  6 | using namespace std;
  7 | 
  8 | std::vector<float> convert_score(const std::vector<float> &input){
  9 |     std::vector<float> exp_values(input.size());
 10 |     for (size_t i = 0; i < input.size(); ++i) {
 11 |         exp_values[i] = std::exp(input[i]);
 12 |     }
 13 |     std::vector<float> output(256);
 14 |     for (size_t i = 256; i < input.size(); ++i) {
 15 |         output[i-256] = exp_values[i] / (exp_values[i] + exp_values[i-256]);
 16 |     }
 17 |     return output;
 18 | }
 19 | 
 20 | inline float fast_exp(float x)
 21 | {
 22 |     union {
 23 |         uint32_t i;
 24 |         float f;
 25 |     } v{};
 26 |     v.i = (1 << 23) * (1.4426950409 * x + 126.93490512f);
 27 |     return v.f;
 28 | }
 29 | 
 30 | inline float sigmoid(float x)
 31 | {
 32 |     return 1.0f / (1.0f + fast_exp(-x));
 33 | }
 34 | 
 35 | static float sz_whFun(cv::Point2f wh)
 36 | {
 37 |     float pad = (wh.x + wh.y) * 0.5f;
 38 |     float sz2 = (wh.x + pad) * (wh.y + pad);
 39 |     return std::sqrt(sz2);
 40 | }
 41 | 
 42 | static std::vector<float> sz_change_fun(std::vector<float> w, std::vector<float> h,float sz)
 43 | {
 44 |     int rows = int(std::sqrt(w.size()));
 45 |     int cols = int(std::sqrt(w.size()));
 46 |     std::vector<float> pad(rows * cols, 0);
 47 |     std::vector<float> sz2;
 48 |     for (int i = 0; i < cols; i++)
 49 |     {
 50 |         for (int j = 0; j < rows; j++)
 51 |         {
 52 |             pad[i*cols+j] = (w[i * cols + j] + h[i * cols + j]) * 0.5f;
 53 |         }
 54 |     }
 55 |     for (int i = 0; i < cols; i++)
 56 |     {
 57 |         for (int j = 0; j < rows; j++)
 58 |         {
 59 |             float t = std::sqrt((w[i*cols+j] + pad[i*cols+j]) * (h[i*cols+j] + pad[i*cols+j])) / sz;
 60 |             sz2.push_back(std::max(t,(float)1.0/t) );
 61 |         }
 62 |     }
 63 |     return sz2;
 64 | }
 65 | 
 66 | static std::vector<float> ratio_change_fun(std::vector<float> w, std::vector<float> h, cv::Point2f target_sz)
 67 | {
 68 |     int rows = int(std::sqrt(w.size()));
 69 |     int cols = int(std::sqrt(w.size()));
 70 |     float ratio = target_sz.x / target_sz.y;
 71 |     std::vector<float> sz2; 
 72 |     for (int i = 0; i < rows; i++)
 73 |     {
 74 |         for (int j = 0; j < cols; j++)
 75 |         {
 76 |             float t = ratio / (w[i * rows + j] / h[i * rows + j]);
 77 |             sz2.push_back(std::max(t, (float)1.0 / t));
 78 |         }
 79 |     }
 80 | 
 81 |     return sz2; 
 82 | }
 83 | 
 84 | NanoTrack::NanoTrack()
 85 | {   
 86 |     
 87 |     
 88 | }
 89 | 
 90 | NanoTrack::~NanoTrack()
 91 | {
 92 |     
 93 | }
 94 | 
 95 | void NanoTrack::init(cv::Mat img, cv::Rect bbox) 
 96 | {
 97 |     create_window(); 
 98 |     create_grids(); 
 99 |     cv::Point2f target_pos ={0.f, 0.f}; // cx, cy
100 |     cv::Point2f target_sz = {0.f, 0.f}; //w,h
101 | 
102 |     target_pos.x = bbox.x + (bbox.width - 1) / 2; 
103 |     target_pos.y = bbox.y + (bbox.height -1) / 2;
104 |     target_sz.x=bbox.width;
105 |     target_sz.y=bbox.height;
106 |     
107 |     float wc_z = target_sz.x + cfg.context_amount * (target_sz.x + target_sz.y);
108 |     float hc_z = target_sz.y + cfg.context_amount * (target_sz.x + target_sz.y);
109 |     float s_z = round(sqrt(wc_z * hc_z));  
110 | 
111 |     cv::Scalar avg_chans = cv::mean(img);
112 |     cv::Mat z_crop;
113 |     
114 |     z_crop  = get_subwindow_tracking(img, target_pos, cfg.exemplar_size, int(s_z),avg_chans); //cv::Mat BGR order 
115 | 
116 |     vector<vector<float>> rknnOutputs;
117 |     int ret = module_T127.runRKNN(rknnOutputs, (void*)z_crop.data, 127 * 127 * 3, RKNN_TENSOR_UINT8, false);
118 | 
119 |     this->result_T = rknnOutputs[0];
120 |     
121 |     this->state.channel_ave=avg_chans;
122 |     this->state.im_h=img.rows;
123 |     this->state.im_w=img.cols;
124 |     this->state.target_pos=target_pos;
125 |     this->state.target_sz= target_sz; 
126 | 
127 | }
128 | 
129 | void NanoTrack::update(const cv::Mat &x_crops, cv::Point &target_pos, cv::Point2f &target_sz,  float scale_z, float &cls_score_max)
130 | {
131 | 
132 | 
133 |     vector<vector<float>> rknnOutputs;
134 |     int ret = this->module_X255.runRKNN(rknnOutputs, (void*)x_crops.data, 255 * 255 * 3 , RKNN_TENSOR_UINT8, false);
135 |     this->result_X = rknnOutputs[0];
136 | 
137 | 
138 |     std::vector<float> result_T_transposedVec(48 * 8 * 8);
139 |     // 原始 shape 为 (48, 8, 8)
140 |     // 目标 shape 为 (8, 8, 48)
141 |     for (int i = 0; i < 48; ++i) {
142 |         for (int j = 0; j < 8; ++j) {
143 |             for (int k = 0; k < 8; ++k) {
144 |                 result_T_transposedVec[j * 8 * 48 + k * 48 + i] = result_T[i * 8 * 8 + j * 8 + k];
145 |             }
146 |         }
147 |     }
148 | 
149 |     std::vector<float> result_X_transposedVec(48 * 16 * 16);
150 |     // 原始 shape 为 (48, 16, 16)
151 |     // 目标 shape 为 (16, 16, 48)
152 |     for (int i = 0; i < 48; ++i) {
153 |         for (int j = 0; j < 16; ++j) {
154 |             for (int k = 0; k < 16; ++k) {
155 |                 result_X_transposedVec[j * 16 * 48 + k * 48 + i] = result_X[i * 16 * 16 + j * 16 + k];
156 |             }
157 |         }
158 |     }
159 | 
160 |     vector<vector<float>> rknnOutputs_2;
161 |     net_head.runRKNN(rknnOutputs_2, (void*)result_T_transposedVec.data(), 48 * 8 * 8*4, (void*)result_X_transposedVec.data(), 48*16*16*4, RKNN_TENSOR_FLOAT32, false);
162 | 
163 |     vector<float> cls_score_result = rknnOutputs_2[0];
164 | 
165 |     vector<float> bbox_pred_result = rknnOutputs_2[1];
166 | 
167 |     float* cls_score_data = (float*) cls_score_result.data();
168 | 
169 |     int cols = 16; 
170 |     int rows = 16; 
171 |     vector<float> cls_scores = convert_score(cls_score_result);
172 |     std::vector<float> pred_x1(cols*rows, 0), pred_y1(cols*rows, 0), pred_x2(cols*rows, 0), pred_y2(cols*rows, 0);
173 | 
174 |     float* bbox_pred_data = (float*) bbox_pred_result.data();
175 | 
176 |     for (int i=0; i<rows; i++)
177 |     {
178 |         for (int j=0; j<cols; j++)
179 |         {
180 |             pred_x1[i*cols + j] = this->grid_to_search_x[i*cols + j] - bbox_pred_data[i*cols + j];
181 |             pred_y1[i*cols + j] = this->grid_to_search_y[i*cols + j] - bbox_pred_data[i*cols + j + 16*16*1];
182 |             pred_x2[i*cols + j] = this->grid_to_search_x[i*cols + j] + bbox_pred_data[i*cols + j + 16*16*2];
183 |             pred_y2[i*cols + j] = this->grid_to_search_y[i*cols + j] + bbox_pred_data[i*cols + j + 16*16*3];
184 |         }
185 |     }
186 | 
187 |     // size penalty  
188 |     std::vector<float> w(cols*rows, 0), h(cols*rows, 0); 
189 |     for (int i=0; i<rows; i++)
190 |     {
191 |         for (int j=0; j<cols; j++) 
192 |         {
193 |             w[i*cols + j] = pred_x2[i*cols + j] - pred_x1[i*cols + j];
194 |             h[i*rows + j] = pred_y2[i*rows + j] - pred_y1[i*cols + j];
195 |         }
196 |     }
197 | 
198 |     float sz_wh = sz_whFun(target_sz);
199 |     std::vector<float> s_c = sz_change_fun(w, h, sz_wh);
200 |     std::vector<float> r_c = ratio_change_fun(w, h, target_sz);
201 | 
202 |     std::vector<float> penalty(rows*cols,0);
203 |     for (int i = 0; i < rows * cols; i++)
204 |     {
205 |         penalty[i] = std::exp(-1 * (s_c[i] * r_c[i]-1) * cfg.penalty_k);
206 |     }
207 | 
208 |     // window penalty
209 |     std::vector<float> pscore(rows*cols,0);
210 |     // int r_max = 0, c_max = 0; 
211 |     float maxScore = 0; 
212 | 
213 |     int max_idx = 0;
214 |     for (int i = 0; i < rows * cols; i++)
215 |     {
216 |         pscore[i] = (penalty[i] * cls_scores[i]) * (1 - cfg.window_influence) + this->window[i] * cfg.window_influence; 
217 |         if (pscore[i] > maxScore) 
218 |         {
219 |             // get max 
220 |             maxScore = pscore[i]; 
221 |             max_idx = i;
222 | 
223 |         }
224 |     }
225 |     
226 |     // to real size
227 |     float pred_x1_real = pred_x1[max_idx]; 
228 |     float pred_y1_real = pred_y1[max_idx];
229 |     float pred_x2_real = pred_x2[max_idx];
230 |     float pred_y2_real = pred_y2[max_idx];
231 | 
232 |     float pred_xs = (pred_x1_real + pred_x2_real) / 2;
233 |     float pred_ys = (pred_y1_real + pred_y2_real) / 2;
234 |     float pred_w = pred_x2_real - pred_x1_real;
235 |     float pred_h = pred_y2_real - pred_y1_real;
236 |     float diff_xs = pred_xs ;
237 |     float diff_ys = pred_ys ;
238 | 
239 |     diff_xs /= scale_z; 
240 |     diff_ys /= scale_z;
241 |     pred_w /= scale_z;
242 |     pred_h /= scale_z;
243 | 
244 |     target_sz.x = target_sz.x / scale_z;
245 |     target_sz.y = target_sz.y / scale_z;
246 | 
247 |     // size learning rate
248 |     float lr = penalty[max_idx] * cls_scores[max_idx] * cfg.lr;
249 | 
250 |     // size rate
251 |     auto res_xs = float (target_pos.x + diff_xs);
252 |     auto res_ys = float (target_pos.y + diff_ys);
253 |     float res_w = pred_w * lr + (1 - lr) * target_sz.x;
254 |     float res_h = pred_h * lr + (1 - lr) * target_sz.y;
255 | 
256 |     target_pos.x = res_xs;
257 |     target_pos.y = res_ys;
258 |     target_sz.x = res_w;
259 |     target_sz.y = res_h;
260 |     cls_score_max = cls_scores[max_idx];
261 | }
262 | 
263 | float NanoTrack::track(cv::Mat& im)
264 | {
265 |     cv::Point target_pos = this->state.target_pos;
266 |     cv::Point2f target_sz = this->state.target_sz;
267 |     
268 |     float hc_z = target_sz.y + cfg.context_amount * (target_sz.x + target_sz.y);
269 |     float wc_z = target_sz.x + cfg.context_amount * (target_sz.x + target_sz.y);
270 |     float s_z = sqrt(wc_z * hc_z);  
271 |     float scale_z = cfg.exemplar_size / s_z;  
272 | 
273 |     float d_search = (cfg.instance_size - cfg.exemplar_size) / 2; 
274 |     float pad = d_search / scale_z; 
275 |     float s_x = s_z + 2*pad;
276 | 
277 |     cv::Mat x_crop;  
278 |     x_crop  = get_subwindow_tracking(im, target_pos, cfg.instance_size, std::round(s_x),state.channel_ave);
279 | 
280 |     // update
281 |     target_sz.x = target_sz.x * scale_z;
282 |     target_sz.y = target_sz.y * scale_z;
283 | 
284 |     float cls_score_max;
285 | 
286 |     this->update(x_crop, target_pos, target_sz, scale_z, cls_score_max);
287 |     target_pos.x = std::max(0, min(state.im_w, target_pos.x));
288 |     target_pos.y = std::max(0, min(state.im_h, target_pos.y));
289 |     target_sz.x = float(std::max(10, min(state.im_w, int(target_sz.x))));
290 |     target_sz.y = float(std::max(10, min(state.im_h, int(target_sz.y))));
291 | 
292 |     state.target_pos = target_pos;
293 |     state.target_sz = target_sz;
294 |     return cls_score_max;
295 | }
296 | 
297 | 
298 | void NanoTrack::load_model(std::string T_model_backbone, std::string X_model_backbone, std::string model_head)
299 | {
300 |     this->module_T127.loadRKNN(T_model_backbone, 1, "model_T");
301 |     this->module_X255.loadRKNN(X_model_backbone, 1, "model_X");
302 |     this->net_head.loadRKNN(model_head, 2, "model_head");
303 | 
304 | }
305 | 
306 | void NanoTrack::create_window()
307 | {
308 |     int score_size= cfg.score_size;
309 |     std::vector<float> hanning(score_size,0);
310 |     this->window.resize(score_size*score_size, 0);
311 | 
312 |     for (int i = 0; i < score_size; i++)
313 |     {
314 |         float w = 0.5f - 0.5f * std::cos(2 * 3.1415926535898f * i / (score_size - 1));
315 |         hanning[i] = w;
316 |     } 
317 |     for (int i = 0; i < score_size; i++)
318 |     {
319 |         for (int j = 0; j < score_size; j++)
320 |         {
321 |             this->window[i*score_size+j] = hanning[i] * hanning[j]; 
322 |         }
323 |     }  
324 | 
325 | }
326 | 
327 | // 生成每一个格点的坐标 
328 | void NanoTrack::create_grids()
329 | {
330 |     /*
331 |     each element of feature map on input search image
332 |     :return: H*W*2 (position for each element)
333 |     */
334 |     int sz = cfg.score_size;   //16x16
335 | 
336 |     this->grid_to_search_x.resize(sz * sz, 0);
337 |     this->grid_to_search_y.resize(sz * sz, 0);
338 | 
339 |     for (int i = 0; i < sz; i++)
340 |     {
341 |         for (int j = 0; j < sz; j++)
342 |         {
343 |             this->grid_to_search_x[i*sz+j] = j*cfg.total_stride-128;   
344 |             this->grid_to_search_y[i*sz+j] = i*cfg.total_stride-128;
345 |         }
346 |     }
347 | }
348 | 
349 | cv::Mat NanoTrack::get_subwindow_tracking(cv::Mat im, cv::Point2f pos, int model_sz, int original_sz,cv::Scalar channel_ave)
350 | {
351 |     float c = (float)(original_sz + 1) / 2;
352 |     int context_xmin = pos.x - c + 0.5;
353 |     int context_xmax = context_xmin + original_sz - 1;
354 |     int context_ymin = pos.y - c + 0.5;
355 |     int context_ymax = context_ymin + original_sz - 1;
356 | 
357 |     int left_pad = int(std::max(0, -context_xmin));
358 |     int top_pad = int(std::max(0, -context_ymin));
359 |     int right_pad = int(std::max(0, context_xmax - im.cols + 1));
360 |     int bottom_pad = int(std::max(0, context_ymax - im.rows + 1));
361 |     context_xmin += left_pad;
362 |     context_xmax += left_pad;
363 |     context_ymin += top_pad;
364 |     context_ymax += top_pad;
365 | 
366 |     cv::Mat im_path_original;
367 | 
368 |     if (top_pad > 0 || left_pad > 0 || right_pad > 0 || bottom_pad > 0)
369 |     {
370 |         cv::Mat te_im = cv::Mat::zeros(im.rows + top_pad + bottom_pad, im.cols + left_pad + right_pad, CV_8UC3);
371 |        
372 |         cv::copyMakeBorder(im, te_im, top_pad, bottom_pad, left_pad, right_pad, cv::BORDER_CONSTANT, channel_ave);
373 |         im_path_original = te_im(cv::Rect(context_xmin, context_ymin, context_xmax - context_xmin + 1, context_ymax - context_ymin + 1));
374 |     }
375 |     else
376 |         im_path_original = im(cv::Rect(context_xmin, context_ymin, context_xmax - context_xmin + 1, context_ymax - context_ymin + 1));
377 | 
378 |     cv::Mat im_path;
379 |     cv::resize(im_path_original, im_path, cv::Size(model_sz, model_sz));
380 | 
381 |     return im_path; 
382 | }


--------------------------------------------------------------------------------
/nanotrack.hpp:
--------------------------------------------------------------------------------
  1 | #ifndef NANOTRACK_H
  2 | #define NANOTRACK_H
  3 | 
  4 | #include <vector> 
  5 | #include <map>  
  6 |  
  7 | #include <opencv2/core/core.hpp> 
  8 | // #include <opencv2/highgui/highgui.hpp> 
  9 | #include <opencv2/imgproc/imgproc.hpp> 
 10 | 
 11 | #include <string.h>
 12 | #include <string>
 13 | #include <vector>
 14 | 
 15 | #include "rknn_api.h"
 16 | #include "RKNNModel.h"
 17 | 
 18 | #include <dlfcn.h>
 19 | #include <stdio.h>
 20 | #include <stdlib.h>
 21 | #include <string.h>
 22 | #include <sys/time.h>
 23 | 
 24 | #include <iostream>
 25 | #include <dirent.h>
 26 | #include <vector>
 27 | #include <string>
 28 | #include <algorithm>
 29 | 
 30 | #define _BASETSD_H
 31 | 
 32 | #include "RgaUtils.h"
 33 | 
 34 | #include "rknn_api.h"
 35 | 
 36 | using namespace std;
 37 | 
 38 | #define PI 3.1415926 
 39 | 
 40 | // using namespace cv;
 41 | 
 42 | struct Config{ 
 43 |     
 44 |     std::string windowing = "cosine";
 45 |     std::vector<float> window;
 46 | 
 47 |     int stride = 16;
 48 |     float penalty_k = 0.15;
 49 |     float window_influence = 0.455;
 50 |     // float lr = 0.38;
 51 |     float lr = 0.37;
 52 | 
 53 |     int exemplar_size=127;
 54 |     int instance_size=255;
 55 |     int total_stride=16;
 56 |     int score_size=16;
 57 |     float context_amount = 0.5;
 58 | };
 59 | 
 60 | struct State { 
 61 |     int im_h; 
 62 |     int im_w;  
 63 |     cv::Scalar channel_ave; 
 64 |     cv::Point target_pos; 
 65 |     cv::Point2f target_sz = {0.f, 0.f}; 
 66 |     float cls_score_max; 
 67 | };
 68 | 
 69 | class NanoTrack {
 70 | 
 71 | public: 
 72 |     
 73 |     NanoTrack();
 74 |     
 75 |     ~NanoTrack(); 
 76 | 
 77 |     void init(cv::Mat img, cv::Rect bbox);
 78 |     
 79 |     void update(const cv::Mat &x_crops, cv::Point &target_pos, cv::Point2f &target_sz, float scale_z, float &cls_score_max);
 80 |         
 81 |     float track(cv::Mat& im);
 82 |     
 83 |     void load_model(std::string T_backbone_model, std::string X_backbone_model, std::string model_head);
 84 | 
 85 |     vector<float> result_T, result_X;
 86 | 
 87 | 
 88 |     int stride=16;
 89 |     
 90 |     // state  dynamic
 91 |     State state;
 92 |     
 93 |     // config static
 94 |     Config cfg; 
 95 | 
 96 |     const float mean_vals[3] = { 0.485f*255.f, 0.456f*255.f, 0.406f*255.f };  
 97 |     const float norm_vals[3] = {1/0.229f/255.f, 1/0.224f/255.f, 1/0.225f/255.f};
 98 | 
 99 | 	RKNNModel module_T127;
100 | 	RKNNModel module_X255;
101 | 	RKNNModel net_head;
102 | 
103 |     rknn_context ctx_T;
104 |     rknn_context ctx_X;
105 |     rknn_context ctx_head;
106 | 
107 | 
108 | 
109 | private: 
110 |     void create_grids(); 
111 |     void create_window();  
112 |     cv::Mat get_subwindow_tracking(cv::Mat im, cv::Point2f pos, int model_sz, int original_sz,cv::Scalar channel_ave);
113 | 
114 |     std::vector<float> grid_to_search_x;
115 |     std::vector<float> grid_to_search_y;
116 |     std::vector<float> window;
117 | };
118 | 
119 | #endif 


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