训练自己的yolov5样本, 并部署到rv1126 <六>

    // 创建推理线程
    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的库等待连接了.

训练自己的yolov5样本, 并部署到rv1126 <六>_第1张图片
追踪一个棉签头.

代码已经开源在:
https://github.com/MontaukLaw/rknn_yolo_rtsp

你可能感兴趣的:(人工智能,深度学习,嵌入式硬件,ai)