yolov3-tiny.cpp

#include "main.h"
#include <sys/time.h>
#include "DetectorYOLOV3.h"

static void DecoderThreadFunc(Queue* queue)
{
    int ret, end = 0;
    int decode_num = 0;
    Queue* que = queue;
    Decoder decoder;
    ret = decoder.DecoderInit();
    if (ret == -1)
    {
        av_log(NULL, AV_LOG_ERROR, "decoder init failed.\n");
        que->finish();
        return;
    }
    while(true)
    {
        if (av_read_frame(decoder.fmt_ctx, decoder.pkt) < 0)
        {
            if(end == 2)
            {
                que->DecodeEnd = true;
                break;
            }
            end = 1;
        }
        if (decoder.pkt->stream_index == decoder.video_stream_idx) {
            if(!end) {
                ret = avcodec_send_packet(decoder.video_dec_ctx, decoder.pkt);
            } else {
                ret = avcodec_send_packet(decoder.video_dec_ctx, NULL);
            }
            if (ret < 0 && ret != AVERROR_EOF) {
                fprintf(stderr, "Error submitting a packet for decoding\n");
                que->DecodeEnd = true;
                break;
            }
            while (ret >= 0 || end == 1)
            {
                ret = avcodec_receive_frame(decoder.video_dec_ctx, decoder.frame);
                if (ret == AVERROR(EAGAIN)) {
                    break;
                } else if (ret == AVERROR_EOF ) {
                    end = 2;
                    break;
                } else if (ret < 0) {
                    av_log(NULL, AV_LOG_ERROR, "Error while receiving a frame from the decoder\n");
                    que->finish();
                    return;
                }
                decode_num++;

                if (decoder.frame->format == AV_PIX_FMT_YUV420P)
                {
                    cv::Mat srcImage;
                    srcImage = cv::Mat::zeros(decoder.frame->height*3/2, decoder.frame->width, CV_8UC1);
                    memcpy(srcImage.data,                                                              (unsigned char*)decoder.frame->data[0], decoder.frame->width * decoder.frame->height);
                    memcpy(srcImage.data + decoder.frame->width * decoder.frame->height,     (unsigned char*)decoder.frame->data[1], decoder.frame->width * decoder.frame->height/4);
                    memcpy(srcImage.data + decoder.frame->width * decoder.frame->height*5/4, (unsigned char*)decoder.frame->data[2], decoder.frame->width * decoder.frame->height/4);
                    que->enQueue(srcImage);
                    av_frame_unref(decoder.frame);
                }
                else
                {
                    av_log(NULL, AV_LOG_ERROR, "video format not yuv420p, fix the code first.\n");
                    que->finish();
                    return;
                }
            }
        }
        av_packet_unref(decoder.pkt);
    }
    fprintf(stdout, "Decoder: ####### decode finish, frame count: %d\n", decode_num);
    que->finish();
}

void Ort_DetectorYOLOV3()
{
    const std::string modelPath = "../Resource/Models/Yolov3/yolov3-tiny.onnx";
    const std::string keysPath = "../Resource/Models/Yolov3/coco.names";

    Queue* queue = new Queue(1);

    DetectorYOLOV3 detector;
    detector.setNumThread(1);
    detector.setGpuIndex(0);
    detector.initModel(modelPath, keysPath);

    std::thread ThreadDecoder(DecoderThreadFunc, queue);

    while (!queue->DecodeEnd) {
        cv::Mat InferImage;
        queue->deQueue(&InferImage);
        if (InferImage.empty()) {
            continue;
        }

        cvtColor(InferImage, InferImage, COLOR_YUV420p2RGB);

        double time1 = getTickCount();
        detector.Detect(InferImage);
        double time2 = getTickCount();
        double elapsedTime = (time2 - time1)*1000 / getTickFrequency();
        fprintf(stdout, "inference time: %.3f ms\n", elapsedTime);
    }
    ThreadDecoder.join();
    delete queue;
    queue = NULL;
    return;
}