modularization/src/detection/yolov4_turnstile_tensorrt7/main.cpp

#include <opencv2/opencv.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <vector>
#include <fstream>
#include <thread>
#include <cmath>

#include "imageBuffer.h"
#include "modulecomm.h"
#include "rawpic.pb.h"
#include "turnstile.pb.h"
#include "yolodetect.h"
#include "detect_turnstile.h"

using namespace std;
using namespace cv;

bool start_up = true;
bool broken_flag = false;

bool test_video = false;
string video_path = "20201231144029.avi";

void * g_turnstile;
string turnstiledata="turnstiledata";

void * mpa_camera;
string cameradata="picfront";

typedef struct frame_info
{
    cv::Mat  frame;
    long long timestamp;
}frame_info;

ConsumerProducerQueue<frame_info> * imageBuffer =  new ConsumerProducerQueue<frame_info>(5,true);

//读取视频数据
void ReadFunc(int n)
{
    cv::VideoCapture cap(video_path);
    if(!cap.isOpened())
    {
        cout<<"camera failed to open"<<endl;
    }

    while(1)
    {
        frame_info frameInfo;
        cv::Mat frame;
        //读视频的时候加上,读摄像头去掉
        if(imageBuffer->isFull())
        {
            continue;
        }
        if(cap.read(frame))
        {
            frameInfo.frame = frame;
            imageBuffer->add(frameInfo);
        }
        else
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(1));
        }
    }
}

//从共享内存中接收摄像头数据
void Listencamera(const char * strdata,const unsigned int nSize,const unsigned int index,const QDateTime * dt,const char * strmemname)
{
    if(nSize<1000)return;
    iv::vision::rawpic pic;
    if(false == pic.ParseFromArray(strdata,nSize))
    {
        std::cout<<"picview Listenpic fail."<<std::endl;
        return;
    }
    cv::Mat mat(pic.height(),pic.width(),pic.mattype());
    if(pic.type() == 1)
        memcpy(mat.data,pic.picdata().data(),mat.rows*mat.cols*mat.elemSize());
    else
    {
//     mat.release();
       std::vector<unsigned char> buff(pic.picdata().data(),pic.picdata().data() + pic.picdata().size());
       mat = cv::imdecode(buff,cv::IMREAD_COLOR);
    }

    frame_info img_info;
    img_info.frame = mat;
    img_info.timestamp = pic.time();
    imageBuffer->add(img_info);
    mat.release();
}

//向共享内存中存入闸机检测结果
void SendDetectResult(frame_info &img_info, bool broken_flag, void* g_name)
{
    iv::vision::turnstile detectResult;
    detectResult.set_state(broken_flag);
    iv::vision::rawpic cameraPic;
    cameraPic.set_time(img_info.timestamp);
    cameraPic.set_elemsize(img_info.frame.elemSize());
    cameraPic.set_width(img_info.frame.cols);
    cameraPic.set_height(img_info.frame.rows);
    cameraPic.set_mattype(img_info.frame.type());
    std::vector<int> param = std::vector<int>(2);
    param[0] = cv::IMWRITE_JPEG_QUALITY;
    param[1] = 95; // default(95) 0-100
    std::vector<unsigned char> buff;
    cv::imencode(".jpg",img_info.frame,buff,param);
    cameraPic.set_picdata(buff.data(),buff.size());
    buff.clear();
    cameraPic.set_type(2);
    detectResult.mutable_pic()->CopyFrom(cameraPic);
    std::string out_result = detectResult.SerializeAsString();
    iv::modulecomm::ModuleSendMsg(g_name,out_result.data(),out_result.length());
}

int main(int argc, char** argv )
{

    if(argc==3)
    {
        test_video = argv[1];
        video_path = argv[2];
    }
    if(argc==2)
        test_video = argv[1];

    if(test_video)
        std::thread * readthread = new std::thread(ReadFunc,1);
    else
        mpa_camera= iv::modulecomm::RegisterRecv(&cameradata[0],Listencamera);

    g_turnstile = iv::modulecomm::RegisterSend(&turnstiledata[0],1000000,1);

    NetworkInfo networkInfo;
    networkInfo.networkType     = "yolov4-turnstile";
    networkInfo.configFilePath  = "model/yolov4-turnstile.cfg";
    networkInfo.wtsFilePath     = "model/yolov4-turnstile.weights";
    networkInfo.deviceType      = "kGPU";
    networkInfo.inputBlobName   = "data";
    std::string modelname = "model/yolov4-turnstile.engine";
    IExecutionContext* yolo_context{nullptr};
    YoloDetect detector(networkInfo,modelname);

    vector<string> classes = {"open", "close"};
    if(start_up)
    {
        //加载网络模型，0是指定第一块GPU
        cudaSetDevice(0);
        if(!detector.loadModel(yolo_context))
        {
            cout<<"load yolo model failed"<<endl;
            return -1;
        }

        //Size size(cap.get(CV_CAP_PROP_FRAME_WIDTH), cap.get(CV_CAP_PROP_FRAME_HEIGHT));
        //VideoWriter writer("./data/result.avi", cv::VideoWriter::fourcc('M', 'J', 'P', 'G'), 25, size);

        vector<KalmanTracker> trackers;
        KalmanTracker::kf_count = 0; // tracking id relies on this, so we have to reset it in each seq.
        int frame_count = 0;
        bool calculation_flag = false;
        frame_info frameInfo;
        Mat frame;
        long long millseconds; //时间戳
        double waittime = (double)getTickCount();
        while (1)
        {

          // if(frame_count>45*30)
           //{
               //calculation_flag = true;
               //break;
           //}

           if(imageBuffer->isEmpty())
           {
               double waittotal = (double)getTickCount() - waittime;
               if(waittotal/cv::getTickFrequency()>30.0)
               {
                   cout<<"Cant't get frame and quit"<<endl;
                   break;
               }
               continue;
           }
           imageBuffer->consume(frameInfo);
           frame = frameInfo.frame;
           frame_count++;

           //cout<<"Frame_count: "<<frame_count<<" "<<frame.size()<<endl;

           double start = (double)getTickCount();

           //前向预测
           float ignore_thresh=0.4;
           float nms_thresh = 0.4;
           std::vector<Detection> detect_results;
           vector<td::bbox_t> outs;
           td::bbox_t result;
           if(detector.process(*yolo_context,frame,detect_results,0.4,0.4))
           {
               for (size_t i = 0; i < detect_results.size(); i++) {
                   cv::Rect r = detector.get_rect(frame, detect_results[i].bbox,detector.m_input_w,detector.m_input_h);
                   result.x = r.x;
                   result.y = r.y;
                   result.w = r.width;
                   result.h = r.height;
                   result.prob = detect_results[i].det_confidence;
                   result.obj_id = detect_results[i].class_id;
                   outs.push_back(result);
               }
           }

           double infertime = (double)getTickCount() - start;

           //std::cout<< "Total Cost of infertime: "  <<infertime*1000.0/cv::getTickFrequency()<<" ms"<<endl;

           vector<td::TrackingBox>track_result;
           //tracking turnstile
           bool track_flag = td::TrackTurnstile(frame_count,trackers,outs,track_result);
           td::Drawer(frame, track_result, classes);

           double total = (double)getTickCount() - start;
           std::cout<< "Total Cost of Detection: "  <<total*1000.0/cv::getTickFrequency()<<" ms"<<endl;

           SendDetectResult(frameInfo, broken_flag, g_turnstile);

           namedWindow("Result",WINDOW_NORMAL);
           imshow("Result",frame);
           if(waitKey(1) == 'q')
               break;
           if(waitKey(1) == 's')
               waitKey(0);
           //writer << frame;

           waittime = (double)getTickCount();
        }
        destroyAllWindows();
        if(trackers.size()>0 && calculation_flag)
        {
            cout<<">>>>>>>>Start Calculation<<<<<<<<"<<endl;
            for (auto it = trackers.begin(); it != trackers.end();)
            {
                int open_count = 0;
                int close_count = 0;
                int open_hit_streak = 0;
                int since_close = 0;
                unsigned int history_size = (*it).m_class_history.size();
                if (history_size >= 15*30)
                {
                    for(unsigned int j=0;j<history_size;j++)
                    {
                        if((*it).m_class_history[j] == 0)
                        {
                            open_count ++;
                            open_hit_streak ++;
                        }
                        else
                        {
                            close_count ++;
                            open_hit_streak = 0;

                        }
                        if(open_hit_streak>12*30)
                        {
                            broken_flag = true;
                            cout<<"The turnstile is broken"<<endl;
                            break;
                        }
                    }
                    if(close_count<history_size*0.2 | broken_flag)
                    {
                        broken_flag = true;
                        cout<<"The turnstile is broken"<<endl;
                        break;
                    }
                    it ++;
                }
                else if(history_size <15*30)
                    it = trackers.erase(it);
                else
                    it ++;
             }
            if(!broken_flag)
                cout<<"The turnstile is normal"<<endl;
        }
        else{
            cout<<"The time is not enough"<<endl;
        }
        //将结果传到共享内存
        if(!frameInfo.frame.empty())
        {
            SendDetectResult(frameInfo, broken_flag, g_turnstile);
            cout<<"Send result to share memory done !"<<endl;
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(1000));
        trackers.clear();
        //cap.release();
        //writer.release();
    }
    yolo_context->destroy();
    return 0;
}