// 创建推理线程
pthread_create(&rknn_yolo_thread_id, NULL, rknn_yolo_thread, NULL);
// 从vi中获取数据,然后那推理结果,画框子,然后送进编码器
pthread_create(&observer_thread_id, NULL, observer_thread, NULL);
新建两个线程,先看rknn_yolo_thread
while (g_flag_run) {
ifDetecting = RK_TRUE;
buffer = RK_MPI_SYS_GetMediaBuffer(RK_ID_RGA, RK_NN_RGA_CHN_INDEX, -1);
if (!buffer) {
usleep(1000);
continue;
}
第一步获取数据,前面说了,vi绑定到rga,直接从rga获取剪裁好的640*640的数据,先从nv12转成rgb:
void *pRknnInputData = malloc(YOLO_INPUT_SIZE);
ret = nv12_to_rgb24_640x640(RK_MPI_MB_GetPtr(buffer), pRknnInputData);
if (ret < 0) {
printf("nv12_to_rgb24_640x640 failed\n");
}
转成RGB之后,数据指针为pRknnInputData
int pResult = predict(pRknnInputData, &detect_result_group);
// put detect result to list
if (detect_result_group.detect_count > 0) {
rknn_list_push(rknn_list_, get_current_time_ms(), detect_result_group);
int size = rknn_list_size(rknn_list_);
if (size >= MAX_RKNN_LIST_NUM) {
rknn_list_drop(rknn_list_);
}
printf("size is %d\n", size);
}
predict即推理部分
rknn_inputs_set(m_ctx, 1, inputs);
ret = rknn_run(m_ctx, NULL);
ret = rknn_outputs_get(m_ctx, 3, outputs, NULL);
/* 后处理 */
post_process_640_v5(outputs, &m_info, detect_result_group);
根据rknn的流程,直接前向传播,经过骨干网络,然后把3个粒度的输出,做后处理, IOU阈值,nms非极大值抑制之后,得到最终的predict的结果,是个detect_result_group_t结构体实例,里面有个结果的数组, detect_result_t results[OBJ_NUMB_MAX_SIZE],detect_result_t就是前面我说的,检测出来的物理类别名称,框子的xywh等
回到推理的线程继续:
if (detect_result_group.detect_count > 0) {
rknn_list_push(rknn_list_, get_current_time_ms(), detect_result_group);
int size = rknn_list_size(rknn_list_);
if (size >= MAX_RKNN_LIST_NUM) {
rknn_list_drop(rknn_list_);
}
printf("size is %d\n", size);
}
将推理的结果放入rknn_list, 其实这里涉及一个生产者/消费者的模型,搞个锁管理一下.
上面就是推理线程的工作,简单来说就是从rga获取640x640的nv12数据,转成rgb,放进rknn做推理,推理结果推入rknn的队列栈.
接下来就是monitor线程
while (g_flag_run) {
buffer = RK_MPI_SYS_GetMediaBuffer(RK_ID_VI, DRAW_RESULT_BOX_CHN_INDEX, -1);
直接从vi通道0拿到数据.
if (rknn_list_size(rknn_list_)) {
long time_before;
detect_result_group_t detect_result_group;
memset(&detect_result_group, 0, sizeof(detect_result_group));
// pick up the first one
rknn_list_pop(rknn_list_, &time_before, &detect_result_group);
// printf("result count:%d \n", detect_result_group.count);
如果推理结果队列不是空的即发现了物体.
用队列的pop功能拿到最先进入队列的数据
for (int j = 0; j < detect_result_group.detect_count; j++) {
int x = detect_result_group.results[j].box.left + X_START;
int y = detect_result_group.results[j].box.top + Y_START;
int w = (detect_result_group.results[j].box.right -
detect_result_group.results[j].box.left);
int h = (detect_result_group.results[j].box.bottom -
detect_result_group.results[j].box.top);
while ((uint32_t) (x + w) >= RTSP_INPUT_VI_WIDTH) {
w -= 16;
}
while ((uint32_t) (y + h) >= RTSP_INPUT_VI_HEIGHT) {
h -= 16;
}
printf("border=(%d %d %d %d)\n", x, y, w, h);
boxInfoList[j] = {x, y, w, h};
boxInfoListNumber++;
}
boxDisplayCounterDown = 1;
}
因为推测结果的框子xy是基于640x640分辨率的,而显示用的视频流是19201080,所以需要做一个坐标的转换,打个比方,有个框的xy是0,0,实际放在19201080的画面中,x在大画面中就是(1920-640)/2 = 640, y在大画面中就是(1080-640)2=220,
if (boxInfoListNumber > 0) {
for (int i = 0; i < boxInfoListNumber; i++) {
nv12_border((char *) RK_MPI_MB_GetPtr(buffer), RTSP_INPUT_VI_WIDTH, RTSP_INPUT_VI_HEIGHT,
boxInfoList[i].x, boxInfoList[i].y, boxInfoList[i].w, boxInfoList[i].h, 0, 0, 255);
}
boxDisplayCounterDown--;
}
if (boxDisplayCounterDown == 0) {
boxInfoListNumber = 0;
memset(boxInfoList, 0, sizeof(boxInfoList));
}
nv12_border((char *) RK_MPI_MB_GetPtr(buffer), RTSP_INPUT_VI_WIDTH, RTSP_INPUT_VI_HEIGHT,
X_START, Y_START, 640, 640, 0, 200, 0);
使用nv12_border在nv12格式的画面中画框子,前面是画检测到的目标,后面是画一个640x640的框,表示算法检测的画面范围.这个640x640以外的画面不被检测.
RK_MPI_SYS_SendMediaBuffer(RK_ID_VENC, DRAW_RESULT_BOX_CHN_INDEX, buffer);
RK_MPI_MB_ReleaseBuffer(buffer);
最后把画完框的数据放进编码器,那边编好码的数据帧就被传入了rtsp的库等待连接了.
代码已经开源在:
https://github.com/MontaukLaw/rknn_yolo_rtsp