modularization/src/detection/detection_yolov4_shigong_monocular/main_multibatch.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 "obstacles.pb.h"
#include "yolodetect.h"
#include "detect_obstacle.h"

using namespace std;
using namespace cv;

uint max_batch_size = 2;

string yaml_path="/home/nvidia/models/camera/external.yaml";

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

void * g_obstacle;
string obstacledata="obstacledata";

void * mpa_camera;
string cameradata="picleft";

typedef struct frame_info
{
    cv::Mat  frame;
    long long timestamp;
}frame_info;
typedef struct obstacle_info
{
    string category;
    uint64 class_id;
    Point3f bbox_3d;
    Rect_<float> bbox;
}obstacle_info;

typedef struct obstacles_info
{
    vector<obstacle_info> bboxes;
    uint64 time;
}obstacles_info;


int capture_width=1280,capture_height=720;
cv::Mat cameraMatrix = (cv::Mat_<double>(3,3)<< 1.0538648949999999e+03, 0., 9.7134830299999999e+02, 0.,
                        9.6887771299999997e+02, 5.5096488099999999e+02, 0., 0., 1.);
cv::Mat distCoeffs = (cv::Mat_<double>(1,5)<<-2.9078999999999999e-01, 6.0994000000000000e-02,
                      3.6489999999999999e-03, -2.0590000000000001e-03, 0.);
cv::Mat cameraMatrix_inv = (cv::Mat_<double>(3,3)<<9.4888823486240150e-04, 0., -9.2170097667025908e-01, 0.,
                            1.0321219970099572e-03, -5.6866297326007331e-01, 0., 0., 1.);
cv::Mat rvecM_inv = (cv::Mat_<double>(3,3)<<9.9361999532145029e-01, 1.5481286202195165e-02,
                     1.1171228524619221e-01, -1.1071480863577153e-01,
                     -5.4789655210038264e-02, 9.9234083098031878e-01,
                     2.1483390005894807e-02, -9.9837789611924532e-01,
                     -5.2726089314150459e-02 );
cv::Mat tvec = (cv::Mat_<double>(3,1)<<1.5367854556890524e+02, 1.5365463964586140e+03,
                -1.2444937322334890e+03);

ConsumerProducerQueue<frame_info> * imageBuffer =  new ConsumerProducerQueue<frame_info>(3,true);
ConsumerProducerQueue<frame_info> * imageBuffer_30 =  new ConsumerProducerQueue<frame_info>(3,true);
ConsumerProducerQueue<frame_info> * imageBuffer_90 =  new ConsumerProducerQueue<frame_info>(3,true);

// 获取系统当前时间,从格林威治标准时间(1970-01-01 00:00:00.000)开始到指定时间的毫秒数。
long long getCurrentSystemTime()
{
    auto now = std::chrono::system_clock::now();
    long long millseconds = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count();
    /*
    uint64_t dis_millseconds = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count()
            - std::chrono::duration_cast<std::chrono::seconds>(now.time_since_epoch()).count() * 1000;
    auto tt = std::chrono::system_clock::to_time_t(now);
    struct tm* ptm = localtime(&tt);
    char date[60] = { 0 };
    sprintf(date, "%d/%02d/%02d/%02d/%02d/%02d/%03d",
            (int)ptm->tm_year + 1900, (int)ptm->tm_mon + 1, (int)ptm->tm_mday,
            (int)ptm->tm_hour, (int)ptm->tm_min, (int)ptm->tm_sec, (int)dis_millseconds);
    return std::string(date);
    */
    return millseconds;
}
//读取视频数据
void ReadFunc(int n)
{
    cv::VideoCapture cap_30(0);
    cv::VideoCapture cap_90(1);
    if(!cap_30.isOpened() | !cap_90.isOpened())
    {
        cout<<"camera failed to open"<<endl;
    }
    while(1)
    {
        frame_info frameInfo_30,frameInfo_90;
        cv::Mat frame_30,frame_90;
        //读视频的时候加上,读摄像头去掉
        if(imageBuffer_30->isFull() | imageBuffer_90->isFull())
        {
            continue;
        }
        if(cap_30.read(frame_30) && cap_90.read(frame_90))
        {
            long long timestamp = getCurrentSystemTime();
            frameInfo_30.frame = frame_30;
            frameInfo_30.timestamp = timestamp;
            imageBuffer_30->add(frameInfo_30);
            frameInfo_90.frame = frame_90;
            frameInfo_90.timestamp = timestamp;
            imageBuffer_90->add(frameInfo_90);
        }
        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,IMREAD_COLOR);
    }

    frame_info img_info;
    img_info.frame = mat;
    img_info.timestamp = pic.time();
    imageBuffer->add(img_info);
    mat.release();
}
//向共享内存中存入障碍物检测结果
void SendDetectResult(obstacles_info &obstacles,void* g_name)
{
    iv::vision::ObstacleInfo detectResult;
    detectResult.set_time(obstacles.time);
    for(unsigned int i=0;i <obstacles.bboxes.size();i++)
    {
        iv::vision::Bbox3D obstacle;
        obstacle.set_category(obstacles.bboxes[i].category);
        obstacle.set_class_id(obstacles.bboxes[i].class_id);
        obstacle.set_x_3d(obstacles.bboxes[i].bbox_3d.x);
        obstacle.set_y_3d(obstacles.bboxes[i].bbox_3d.y);
        obstacle.set_z_3d(obstacles.bboxes[i].bbox_3d.z);
        obstacle.set_x(obstacles.bboxes[i].bbox.x);
        obstacle.set_y(obstacles.bboxes[i].bbox.y);
        obstacle.set_w(obstacles.bboxes[i].bbox.width);
        obstacle.set_h(obstacles.bboxes[i].bbox.height);
        iv::vision::Bbox3D * bb = detectResult.add_bbox_3d();
        bb->CopyFrom(obstacle);
    }
    std::string out_result = detectResult.SerializeAsString();
    iv::modulecomm::ModuleSendMsg(g_name,out_result.data(),out_result.length());
}
//像素坐标系转换到车体坐标系
cv::Point3f CameraToWorld(cv::Point2f CameraInfo)
{
    cv::Mat imagePoint = cv::Mat::ones(3, 1, cv::DataType<double>::type);//u,v,1
    cv::Mat tempMat, tempMat2;
    //输入一个2D坐标点，便可以求出相应的s
    imagePoint.at<double>(0,0)=double(CameraInfo.x);
    imagePoint.at<double>(1,0)=double(CameraInfo.y);
    double zConst = 0;//实际坐标系的距离
    //计算参数s
    double s;
    tempMat = rvecM_inv * cameraMatrix_inv * imagePoint;
    tempMat2 = rvecM_inv * tvec;
    s = zConst + tempMat2.at<double>(2, 0);
    s /= tempMat.at<double>(2, 0);
    //cout<<"s : "<<s<<endl;
    //////////////////////camera_coordinates////////////////
    cv::Mat camera_cordinates=-rvecM_inv*tvec;
    /////////////////////2D to 3D///////////////////////
    cv::Mat wcPoint = rvecM_inv * (cameraMatrix_inv *s*imagePoint - tvec);
    cv::Point3f worldPoint(wcPoint.at<double>(0, 0), wcPoint.at<double>(1, 0),
                           wcPoint.at<double>(2, 0));
    return worldPoint;
}
int main(int argc, char** argv )
{

    if(argc==3)
    {
        test_video = (strcmp(argv[1], "true") == 0)?true:false;
        if(test_video)
            video_path = argv[2];
        else
            cameradata = argv[2];
    }
    if(argc==2)
        test_video = (strcmp(argv[1], "true") == 0)?true:false;

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

    g_obstacle = iv::modulecomm::RegisterSend(&obstacledata[0],1000000,1);

    //导入相机内参和畸变系数矩阵
    FileStorage file_storage(yaml_path, FileStorage::READ);
    if( file_storage.isOpened())
    {
        file_storage["camera_matrix"] >> cameraMatrix;
        file_storage["camera_matrix_inv"] >> cameraMatrix_inv;
        file_storage["distortion_coefficients"] >> distCoeffs;
        file_storage["RotationR_inv"] >> rvecM_inv;
        file_storage["tvec"] >> tvec;
        file_storage.release();
    }
    else{
        cout << "Error: can not open the external parameters file"<<endl;
        return 0;
    }

    NetworkInfo networkInfo;
    networkInfo.networkType     = "yolov4";
    networkInfo.configFilePath  = "/home/nvidia/models/camera/yolov4-shigong.cfg";
    networkInfo.wtsFilePath     = "/home/nvidia/models/camera/yolov4-shigong_final.weights";
    networkInfo.deviceType      = "kGPU";
    networkInfo.inputBlobName   = "data";
    networkInfo.maxbatchSize    = max_batch_size;
    std::string modelname = "/home/nvidia/models/camera/yolov4_shigong.engine";
    IExecutionContext* yolo_context{nullptr};
    YoloDetect detector(networkInfo,modelname);

    /*
    string classesFile = "model/shigong.names";
    //加载类别名
    vector<string> classes;
    ifstream ifs(classesFile.c_str());
    string line;
    while (getline(ifs, line)) classes.push_back(line);
    */
    vector<string> classes = {"shigong"};
    //加载网络模型，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_30,trackers_90;
    KalmanTracker::kf_count = 0; // tracking id relies on this, so we have to reset it in each seq.
    int frame_count = 0;
    frame_info frameInfo_30,frameInfo_90;
    Mat frame_30,frame_90;

    long long millseconds; //时间戳
    double waittime = (double)getTickCount();

    while (1)
    {
        if(imageBuffer_30->isEmpty() | imageBuffer_90->isEmpty())
        {
            double waittotal = (double)getTickCount() - waittime;
            if(waittotal/cv::getTickFrequency()>30.0)
            {
                cout<<"Can't get frame and quit"<<endl;
                break;
            }
            continue;
        }

        imageBuffer_30->consume(frameInfo_30);
        frame_30 = frameInfo_30.frame;
        imageBuffer_90->consume(frameInfo_90);
        frame_90 = frameInfo_90.frame;
        vector<Mat> frame_vec;
        frame_vec.push_back(frame_30);
        frame_vec.push_back(frame_90);

        obstacles_info detect_result_info;
        detect_result_info.time = frameInfo_30.timestamp;
        frame_count++;

        double start = (double)getTickCount();
        //前向预测
        float ignore_thresh=0.4;
        float nms_thresh = 0.4;

        //两个摄像头30,90度
        std::vector<std::vector<Detection>> detect_results_yolo;
        std::vector<Detection> result_30,result_90;
        detect_results_yolo.push_back(result_30);
        detect_results_yolo.push_back(result_90);
        bool flag_30 = false;
        bool flag_90 = false;
        od::bbox_t bbox_t_30;
        od::bbox_t bbox_t_90;
        vector<od::bbox_t> outs_30;
        vector<od::bbox_t> outs_90;

        if(detector.process(*yolo_context,max_batch_size,frame_vec,
                            detect_results_yolo,ignore_thresh,nms_thresh))
        {
            for (size_t i = 0; i < detect_results_yolo[0].size(); i++) {
                cv::Rect r = detector.get_rect(frame_30, detect_results_yolo[0][i].bbox,detector.m_input_w,detector.m_input_h);
                bbox_t_30.x = r.x;
                bbox_t_30.y = r.y;
                bbox_t_30.w = r.width;
                bbox_t_30.h = r.height;
                bbox_t_30.prob = detect_results_yolo[0][i].det_confidence;
                bbox_t_30.obj_id = detect_results_yolo[0][i].class_id;
                outs_30.push_back(bbox_t_30);
            }
            for (size_t i = 0; i < detect_results_yolo[1].size(); i++) {
                cv::Rect r = detector.get_rect(frame_90, detect_results_yolo[1][i].bbox,
                                               detector.m_input_w,detector.m_input_h);
                bbox_t_90.x = r.x;
                bbox_t_90.y = r.y;
                bbox_t_90.w = r.width;
                bbox_t_90.h = r.height;
                bbox_t_90.prob = detect_results_yolo[1][i].det_confidence;
                bbox_t_90.obj_id = detect_results_yolo[1][i].class_id;
                outs_90.push_back(bbox_t_90);
                od::Drawer(frame_90, outs_90, classes);
            }
        }

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

        vector<od::TrackingBox>track_result_30,track_result_90;
        bool track_flag_30 = od::TrackObstacle(frame_count,trackers_30,outs_30,track_result_30);
        bool track_flag_90 = od::TrackObstacle(frame_count,trackers_90,outs_90,track_result_90);
        if(track_flag_30)
        {
            vector<obstacle_info>obstacle_vector;
            for(unsigned int i=0;i < track_result_30.size(); i++)
            {
                obstacle_info obstacle;
                obstacle.category = classes[track_result_30[i].class_id];
                obstacle.class_id = track_result_30[i].class_id;
                //求bbox的接地点
                Point2f bottom_center = Point2f(track_result_30[i].box.x+0.5*track_result_30[i].box.width,
                                                track_result_30[i].box.y+track_result_30[i].box.height);
                obstacle.bbox_3d = CameraToWorld(bottom_center);
                obstacle.bbox = track_result_30[i].box;
                obstacle_vector.push_back(obstacle);
            }
            detect_result_info.bboxes = obstacle_vector;
            SendDetectResult(detect_result_info,g_obstacle);
            od::Drawer(frame_30, track_result_30, classes);
        }
        else if(track_flag_90)
        {
            vector<obstacle_info>obstacle_vector;
            for(unsigned int i=0;i < track_result_90.size(); i++)
            {
                obstacle_info obstacle;
                obstacle.category = classes[track_result_90[i].class_id];
                obstacle.class_id = track_result_90[i].class_id;
                //求bbox的接地点
                Point2f bottom_center = Point2f(track_result_90[i].box.x+0.5*track_result_90[i].box.width,
                                                track_result_90[i].box.y+track_result_90[i].box.height);
                obstacle.bbox_3d = CameraToWorld(bottom_center);
                obstacle.bbox = track_result_90[i].box;
                obstacle_vector.push_back(obstacle);
            }
            detect_result_info.bboxes = obstacle_vector;
            SendDetectResult(detect_result_info,g_obstacle);
            od::Drawer(frame_90, track_result_90, classes);
        }
        double total = (double)getTickCount() - start;
        std::cout<< "Total Cost of Detection: "  <<total*1000.0/cv::getTickFrequency()<<" ms"<<endl;

        namedWindow("Result_30",WINDOW_NORMAL);
        imshow("Result_30",frame_30);

        namedWindow("Result_90",WINDOW_NORMAL);
        imshow("Result_90",frame_90);

        if(waitKey(1) == 'q')
            break;
        if(waitKey(1) == 's')
            waitKey(0);
        waittime = (double)getTickCount();
    }
    destroyAllWindows();
    trackers_30.clear();
    trackers_90.clear();
    yolo_context->destroy();
    return 0;
}