SSD.cpp

#include <Sample.h>
#include <SimpleLog.h>
#include <Filesystem.h>
#include <DetectorSSD.h>
#include <sys/time.h>
#include <Decoder.h>
#include <Queuethread.h>

using namespace cv;
using namespace std;
using namespace cv::dnn;
using namespace migraphx;
using namespace migraphxSamples;

static void DecoderThreadFunc(Queue* queue)
{
    int ret, end = 0;
    int frame_cnt = 0;
    Queue* que = queue;
    Decoder decoder(que->device);
    InitializationParameterOfDecoder initParamOfDecoderSSD;
    #ifndef DMA
    initParamOfDecoderSSD.src_filename = "../Resource/Images/Mean.mp4";
    if( que->device == _HW) {
        initParamOfDecoderSSD.str_devid[0] = {0};
        initParamOfDecoderSSD.xcoder_params = "out=hw";
        initParamOfDecoderSSD.dec_name = "h264_ni_quadra_dec";
        initParamOfDecoderSSD.filters_descr = "ni_quadra_scale=640:480:format=bgrp,hwdownload,format=bgrp";
    } else if (que->device == _HW_DMA) {
        LOG_ERROR(stdout, "Error program param or cmake param, not USE_P2P can`t set '--dma'!\n");
        que->finish();
        return;
    }
    #else
    if( que->device == _HW_DMA) {
        initParamOfDecoderSSD.str_devid[0] = {0};
        initParamOfDecoderSSD.xcoder_params = "out=hw";
        initParamOfDecoderSSD.dec_name = "h264_ni_quadra_dec";
        initParamOfDecoderSSD.filters_descr = "ni_quadra_scale=640:480:format=rgba:is_p2p=1";
        initParamOfDecoderSSD.src_filename = "../Resource/Images/cr7_1920x1080.h264";
    } else {
        LOG_ERROR(stdout, "Error program param or cmake param, USE_P2P need set '--dma'!\n");
        que->finish();
        return;
    }
    #endif
    ret = decoder.DecoderInit(initParamOfDecoderSSD);
    if (ret == -1)
    {
        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;
                }
                decoder.frame->pts = decoder.frame->best_effort_timestamp;
                frame_cnt++;

                if (que->device == CPU)
                {
                    cv::Mat 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);
                    cvtColor(srcImage, srcImage, COLOR_YUV420p2RGB);
                    que->enQueue(srcImage);
                }
                if (que->device == _HW || que->device == _HW_DMA)
                {
                    if (av_buffersrc_add_frame_flags(decoder.buffersrc_ctx, decoder.frame, AV_BUFFERSRC_FLAG_KEEP_REF) < 0) {
                        av_log(NULL, AV_LOG_ERROR, "Error while feeding the filtergraph\n");
                        break;
                    }
                    while (1)
                    {
                        ret = av_buffersink_get_frame(decoder.buffersink_ctx, decoder.filt_frame);
                       if (ret == AVERROR(EAGAIN))
                       {
                           break;
                       }
                       else if(ret == AVERROR_EOF)
                       {
                           end = 2;
                           break;
                       }
                       if (ret < 0)
                       {
                           que->finish();
                           return;
                       }

                       #ifndef DMA
                       if (que->device == _HW)
                       {
                           cv::Mat srcImage;
                           switch (decoder.filt_frame->format)
                           {
                               case AV_PIX_FMT_BGRP:
                               {
                                   srcImage = cv::Mat::zeros(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC3);
                                   cv::Mat mat_r = cv::Mat(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC1, (unsigned char*)decoder.filt_frame->data[0]);
                                   cv::Mat mat_g = cv::Mat(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC1, (unsigned char*)decoder.filt_frame->data[1]);
                                   cv::Mat mat_b = cv::Mat(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC1, (unsigned char*)decoder.filt_frame->data[2]);
                                   cv::Mat Channels[3]{mat_r, mat_g, mat_b};
                                   cv::merge(Channels, 3, srcImage);
                                   break;
                               }
                               case AV_PIX_FMT_YUV420P:
                               {
                                   srcImage = cv::Mat::zeros(decoder.filt_frame->height*3/2, decoder.filt_frame->width, CV_8UC1);
                                   memcpy(srcImage.data,							       (unsigned char*)decoder.filt_frame->data[0], decoder.filt_frame->width * decoder.filt_frame->height);
                                   memcpy(srcImage.data + decoder.filt_frame->width * decoder.filt_frame->height,     (unsigned char*)decoder.filt_frame->data[1], decoder.filt_frame->width * decoder.filt_frame->height/4);
                                   memcpy(srcImage.data + decoder.filt_frame->width * decoder.filt_frame->height*5/4, (unsigned char*)decoder.filt_frame->data[2], decoder.filt_frame->width * decoder.filt_frame->height/4);
                                   cvtColor(srcImage, srcImage, COLOR_YUV420p2RGB);
                                   break;
                               }
                               case AV_PIX_FMT_RGBA:
                               {
                                   srcImage = cv::Mat::zeros(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC4);
                                   memcpy(srcImage.data, (unsigned char*)decoder.filt_frame->data[0], decoder.filt_frame->width * decoder.filt_frame->height * 4);
                                   cvtColor(srcImage, srcImage, COLOR_BGRA2RGB);
                                   break;
                               }
                               default:
                                   break;
                           }
                           que->enQueue(srcImage);
                           av_frame_unref(decoder.filt_frame);
                       }
                       #else
                       if (que->device == _HW_DMA)
                       {
                           DCU_Frame dcu_frame;
                           AVHWFramesContext *hwfc  = (AVHWFramesContext *)decoder.filt_frame->hw_frames_ctx->data;
                           switch (hwfc->sw_format)
                           {
                               case AV_PIX_FMT_BGRP:
                               {
                                   dcu_frame.format = AV_PIX_FMT_BGRP;
                                   dcu_frame.data_len = decoder.filt_frame->width * decoder.filt_frame->height * 3;
                                   dcu_frame.srcImage = cv::Mat::zeros(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC3);
                                   break;
                               }
                               case AV_PIX_FMT_YUV420P:
                               {
                                   dcu_frame.format = AV_PIX_FMT_YUV420P;
                                   dcu_frame.data_len = decoder.filt_frame->width * decoder.filt_frame->height * 3 / 2;
                                   dcu_frame.srcImage = cv::Mat::zeros(decoder.filt_frame->height*3/2, decoder.filt_frame->width, CV_8UC1);
                                   break;
                               }
                               case AV_PIX_FMT_RGBA:
                               {
                                   dcu_frame.format = AV_PIX_FMT_RGBA;
                                   dcu_frame.data_len = decoder.filt_frame->width * decoder.filt_frame->height * 4;
                                   dcu_frame.srcImage = cv::Mat::zeros(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC4);
                                   break;
                               }
                               default:
                                   break;
                           }

                           dcu_frame.width = decoder.filt_frame->width;
                           dcu_frame.height = decoder.filt_frame->height;
                           hipMalloc((void**)&(dcu_frame.dcu_data), dcu_frame.data_len * sizeof(unsigned char));
                           ret = decoder.retrieve_filter_frame(dcu_frame, decoder.filt_frame);
                           if (ret)
                               av_log(NULL, AV_LOG_ERROR, "Error while retrieve_filter_frame with p2p.\n");
                           if(dcu_frame.format == AV_PIX_FMT_BGRP)
                           {
                               cv::Mat mat_b = cv::Mat(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC1, (unsigned char*)dcu_frame.srcImage.data);
                               cv::Mat mat_g = cv::Mat(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC1, (unsigned char*)(dcu_frame.srcImage.data + decoder.filt_frame->height * decoder.filt_frame->width));
                               cv::Mat mat_r = cv::Mat(decoder.filt_frame->height, decoder.filt_frame->width, CV_8UC1, (unsigned char*)(dcu_frame.srcImage.data + decoder.filt_frame->height * decoder.filt_frame->width * 2));
                               cv::Mat Channels[3]{mat_r, mat_g, mat_b};
                               cv::merge(Channels, 3, dcu_frame.srcImage);
                           }
                           if(dcu_frame.format == AV_PIX_FMT_YUV420P)
                               cvtColor(dcu_frame.srcImage, dcu_frame.srcImage, COLOR_YUV420p2RGB);
                           if(dcu_frame.format == AV_PIX_FMT_RGBA)
                               cvtColor(dcu_frame.srcImage, dcu_frame.srcImage, COLOR_BGRA2RGB);

                           queue->enQueue(dcu_frame);
                           av_frame_unref(decoder.filt_frame);
                       }
                       #endif
                   }
               }
               av_frame_unref(decoder.frame);
           }
       }
       av_packet_unref(decoder.pkt);
    }
    LOG_INFO(stdout, "Decoder: ####### frame count: %d\n", frame_cnt);
    que->finish();
}

static void DetectorThreadFunc(Queue* que)
{
    Queue* queue = que;
    // DetectorSSD Init
    DetectorSSD detector;
    InitializationParameterOfDetector initParamOfDetectorSSD;
    initParamOfDetectorSSD.parentPath = "";
    initParamOfDetectorSSD.configFilePath = CONFIG_FILE;
    initParamOfDetectorSSD.logName = "";
    ErrorCode errorCode=detector.Initialize(initParamOfDetectorSSD);
    if(errorCode!=SUCCESS)
    {
	LOG_ERROR(stdout, "fail to initialize detector!\n");
	exit(-1);
    }
    LOG_INFO(stdout, "succeed to initialize detector\n");

    int frame_cnt = 0;
    double start_time = getTickCount();
    while (!queue->DecodeEnd) {
        #ifdef DMA
        DCU_Frame dcu_frame;
        queue->deQueue(&dcu_frame);
        if(dcu_frame.srcImage.empty()) {
            continue;
        }
        #else
        cv::Mat InferImage;
        queue->deQueue(&InferImage);
        if (InferImage.empty()) {
            continue;
        }
        #endif

        // detect
        std::vector<ResultOfDetection> predictions;
        double time1 = getTickCount();
        #ifdef DMA
        detector.Detect(dcu_frame, predictions);
        #else
        detector.Detect(InferImage, predictions);
        #endif
        double time2 = getTickCount();
        double elapsedTime = (time2 - time1)*1000 / getTickFrequency();
        LOG_INFO(stdout, "inference time:%f ms\n", elapsedTime);
        frame_cnt++;
        #ifdef DMA
        hipFree(dcu_frame.dcu_data);
        #endif

        // process result
        LOG_INFO(stdout,"////////////////Detection Results////////////////\n");
        for( int i = 0; i < predictions.size(); ++i)
        {
            ResultOfDetection result = predictions[i];
            #ifdef DMA
            cv::rectangle(dcu_frame.srcImage, result.boundingBox, Scalar(0,255,255),2);
            cv::putText(dcu_frame.srcImage, result.className, cv::Point(result.boundingBox.x, result.boundingBox.y-20), cv::FONT_HERSHEY_PLAIN, 2.0, Scalar(0, 0, 255), 2);
            #else
            cv::rectangle(InferImage, result.boundingBox, Scalar(0,255,255),2);
            cv::putText(InferImage, result.className, cv::Point(result.boundingBox.x, result.boundingBox.y-20), cv::FONT_HERSHEY_PLAIN, 2.0, Scalar(0, 0, 255), 2);
            #endif

            LOG_INFO(stdout,"box:%d %d %d %d,label:%d,confidence:%f\n",result.boundingBox.x,
            result.boundingBox.y,result.boundingBox.width,result.boundingBox.height,result.classID,result.confidence);
        }
        // X11 display can`t support in docker.
        /*namedWindow("video", WINDOW_NORMAL | WINDOW_KEEPRATIO);
        #ifdef DMA
        imshow("video", dcu_frame.srcImage);
        #else
        imshow("video", InferImage);
        #endif
        if (waitKey(10) == 'q') {
            break;
        }*/
    }
    #ifdef DMA
    hipFree(detector.preprocess_Image);
    #endif
    double end_time = getTickCount();
    fprintf(stdout, "Finish ####### frame_cnt: %d, Inference fps: %.2f, all time: %.2f ms\n", frame_cnt, float(frame_cnt/((end_time - start_time)/getTickFrequency())), (end_time - start_time)/getTickFrequency()*1000);
    queue->finish();
}


void Sample_DetectorSSD(int device)
{
    Queue* queue = new Queue(1);
    queue->device = device;

    std::thread ThreadDecoder(DecoderThreadFunc, queue);
    std::thread ThreadDetector(DetectorThreadFunc, queue);

    ThreadDecoder.join();
    ThreadDetector.join();

    delete queue;
    queue = NULL;
    return;
}