Train_Identify/nvidia_ascend_engine/nvidia_engine/MoveEngine/MoveEngine.cpp

#include "MoveEngine.h"
#include <opencv2/opencv.hpp>
#include "myutils.h"
#include "myqueue.h"

using namespace ai_matrix;

extern bool g_bHaveTrainFlag;
extern bool g_bNoDealStepTwoFlag;

MoveEngine::MoveEngine() {}

MoveEngine::~MoveEngine() {}

APP_ERROR MoveEngine::Init()
{
    strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";
    modelConfig_ = MyYaml::GetIns()->GetModelConfig("MoveEngine");
    bNeedMoveDetectFlag_ = MyYaml::GetIns()->GetBoolValue("gc_need_move_detect_flag");
    strResultPath_ = MyYaml::GetIns()->GetPathValue("gc_result_path");
    strBestPath_ = MyYaml::GetIns()->GetPathValue("gc_best_path");

    if (bNeedMoveDetectFlag_)
    {
        // 读取模型信息
        APP_ERROR ret = ReadModelInfo();
        if (ret != APP_ERR_OK)
        {
            LogError << "Failed to read model info, ret = " << ret;
            return ret;
        }

        ret = InitModel();
        if (ret != APP_ERR_OK)
        {
            LogError << "Failed to read model info, ret = " << ret;
            return ret;
        }
    }

    std::map<int, ai_matrix::DataSourceConfig> mapUseDataSouceCfg = MyYaml::GetIns()->GetUseDataSourceConfig();
    for (auto iter = mapUseDataSouceCfg.begin(); iter != mapUseDataSouceCfg.end(); iter++)
    {
        //端口0是主摄像头push存图，需给其他使用的数据源push(来车/结束)通知
        if (iter->first == 0)
        {
            dataSourceCfg_ = iter->second;
            continue;
        }
        setPushPort_.insert(iter->first);
    }

    InitParam();

    LogInfo << "MoveEngine Init ok";
    return APP_ERR_OK;
}

APP_ERROR MoveEngine::InitModel()
{
//    modelinfo.yolov5ModelParam.uiClassNum = class_num;
//    modelinfo.yolov5ModelParam.uiDetSize = det_size;
//    modelinfo.yolov5ModelParam.fScoreThreshold = score_threshold;
//    modelinfo.yolov5ModelParam.fNmsThreshold = nms_threshold;

    modelinfo.modelCommonInfo.uiModelWidth = model_width;
    modelinfo.modelCommonInfo.uiModelHeight = model_height;
    modelinfo.modelCommonInfo.uiInputSize = input_size;
    modelinfo.modelCommonInfo.uiOutputSize = output_size;
    modelinfo.modelCommonInfo.uiChannel = INPUT_CHANNEL;
    modelinfo.modelCommonInfo.uiBatchSize = batch_size;
    modelinfo.modelCommonInfo.strInputBlobName = INPUT_BLOB_NAME;
    modelinfo.modelCommonInfo.strOutputBlobName = OUTPUT_BLOB_NAME;

    string strModelName = "";

    int nRet = yolov8model.YoloV8InferenceInit(&modelinfo, strModelName, modelConfig_.strOmPath);
    if (nRet != 0)
    {
        LogInfo << "YoloV5ClassifyInferenceInit nRet:" << nRet;
        return APP_ERR_COMM_READ_FAIL;
    }
    return APP_ERR_OK;
}

APP_ERROR MoveEngine::ReadModelInfo()
{
    char szAbsPath[PATH_MAX];
    // Get the absolute path of model file
    if (realpath(modelConfig_.strOmPath.c_str(), szAbsPath) == nullptr)
    {
        LogError << "Failed to get the real path of " << modelConfig_.strOmPath.c_str();
        return APP_ERR_COMM_NO_EXIST;
    }

    // Check the validity of model path
    int iFolderExist = access(szAbsPath, R_OK);
    if (iFolderExist == -1)
    {
        LogError << "ModelPath " << szAbsPath << " doesn't exist or read failed!";
        return APP_ERR_COMM_NO_EXIST;
    }

    //读取模型参数信息文件
    Json::Value jvModelInfo;
    if (!MyUtils::getins()->ReadJsonInfo(jvModelInfo, modelConfig_.strModelInfoPath))
    {
        LogError << "ModelInfoPath:" << modelConfig_.strModelInfoPath << " doesn't exist or read failed!";
        return APP_ERR_COMM_NO_EXIST;
    }

    model_width = jvModelInfo["model_width"].asInt();
    model_height = jvModelInfo["model_height"].asInt();
    input_size = GET_INPUT_SIZE(model_width , model_height);
    output_size = GET_OUTPUT_SIZE(model_width , model_height, clear_num , class_num);

    return APP_ERR_OK;
}

APP_ERROR MoveEngine::DeInit()
{
    if (bNeedMoveDetectFlag_)
    {
        yolov8model.YoloV8InferenceDeinit();
    }
    LogInfo << "MoveEngine DeInit ok";
    return APP_ERR_OK;
}


/**
* 参数初始化(列车结束时需调用)
* inParam : N/A
* outParam: N/A
* return  : N/A
*/
void MoveEngine::InitParam()
{
    iStepInter_ = 0;
    iMoveDataNO_ = 1;
}

void MoveEngine::sendComeTrain() {
    std::string message = "{\"cometime\":\"" + this->strTrainDate_ + " " + this->strTrainName_ + "\",\"type\":\"1\"}";
    outputQueMap_[engineName_ + "_" + std::to_string(engineId_) + "_1"]->push(std::static_pointer_cast<void>(std::make_shared<std::string>(message)));
}

void MoveEngine::sendEndTrain() {
    std::string message = "{\"cometime\":\"" + this->strTrainDate_ + " " + this->strTrainName_ + "\",\"type\":\"0\"}";
    outputQueMap_[engineName_ + "_" + std::to_string(engineId_) + "_1"]->push(std::static_pointer_cast<void>(std::make_shared<std::string>(message)));
}


/**
* 使用单device处理
* inParam : std::shared_ptr<ProcessData> pProcessData
* outParam: N/A
* return  : N/A
*/
void MoveEngine::SingleDeviceProcess(std::shared_ptr<ProcessData> pProcessData, int nType)
{
    if(iMoveDataNO_ == 1)
    {
        strTrainDate_ = MyUtils::getins()->GetDate();
        strTrainName_ = MyUtils::getins()->GetTime();
        //创建该该列车的数据存储目录(防止后续多线程创建报错)
        std::string strTrainPath = strResultPath_ + strTrainDate_ + "/" + strTrainName_ + "/";
        MyUtils::getins()->CreateDirPath(strTrainPath);
        std::string strBestImgPath = strBestPath_ + strTrainDate_ + "/" + strTrainName_ + "/";
        MyUtils::getins()->CreateDirPath(strBestImgPath);
    }
    pProcessData->strTrainDate = strTrainDate_;
    pProcessData->strTrainName = strTrainName_;
    pProcessData->iFrameId = iMoveDataNO_ * dataSourceCfg_.iSkipInterval;
    //组织数据, push其他端口 (只通知2次，车来一次，车结束一次)
    if (iMoveDataNO_ == 1 || pProcessData->bIsEnd || (nPreMonitorState != nType))
    {
        if(nPreMonitorState != nType){
            nPreMonitorState = nType;
        }
        LogDebug << "来车检测--> train_date:" << strTrainDate_ << " train_name:" << strTrainName_
                 << " frameid:" << pProcessData->iFrameId << " isEnd:" << pProcessData->bIsEnd;
        std::shared_ptr<MoveData> pMoveData = std::make_shared<MoveData>();
        pMoveData->iFrameId = iMoveDataNO_ * dataSourceCfg_.iSkipInterval; //当前帧号
        pMoveData->i64TimeStamp = pProcessData->i64TimeStamp;
        pMoveData->bHasTrain = true;
        pMoveData->bIsEnd = pProcessData->bIsEnd;
        pMoveData->strTrainDate = strTrainDate_;
        pMoveData->strTrainName = strTrainName_;
        pMoveData->nMonitorState = nType;
        for (auto iter = setPushPort_.begin(); iter != setPushPort_.end(); iter++)
        {
            outputQueMap_[engineName_ + "_" + std::to_string(engineId_) + "_" + std::to_string(*iter)]->push(std::static_pointer_cast<void>(pMoveData));
        }

        // LogInfo << "【来车检测】" << pMoveData->iFrameId << "  isEnd:" << pProcessData->bIsEnd;
        //通知第一步开始识别        
        outputQueMap_[engineName_ + "_" + std::to_string(engineId_) + "_5"]->push(std::static_pointer_cast<void>(pMoveData));
    }

    pProcessData->iDataNO = iMoveDataNO_++;
    //push端口,存图
    std::shared_ptr<SaveImgData> pSaveImgData = std::make_shared<SaveImgData>();
    pSaveImgData->pData = pProcessData->pData;
    pSaveImgData->iSize = pProcessData->iSize;
    pSaveImgData->iFrameId = pProcessData->iFrameId; // 帧号
    pSaveImgData->iWidth = pProcessData->iWidth;
    pSaveImgData->iHeight = pProcessData->iHeight;
    char szCameraNo[4] = {0};
    sprintf(szCameraNo, "%03d", pProcessData->iDataSource + 1);
    pSaveImgData->strImgPath = strResultPath_ + pProcessData->strTrainDate + "/" + pProcessData->strTrainName + "/" + szCameraNo;
    pSaveImgData->strImgName = std::to_string(pSaveImgData->iFrameId) + ".jpg";
    pSaveImgData->bIsEnd = pProcessData->bIsEnd;
    pSaveImgData->bSaveToFtp = true;
    pSaveImgData->i64TimeStamp = pProcessData->i64TimeStamp;
    pSaveImgData->nMonitorState = nType;
    outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pSaveImgData));
}

APP_ERROR MoveEngine::Process()
{
    int iRet = APP_ERR_OK;
    while (!isStop_)
    {
        //pop端口0
        std::shared_ptr<void> pVoidData0 = nullptr;
        inputQueMap_[strPort0_]->pop(pVoidData0);
        if (nullptr == pVoidData0)
        {
            usleep(1000); //1ms
            continue;
        }

        std::shared_ptr<ProcessData> pProcessData = std::static_pointer_cast<ProcessData>(pVoidData0);

        int iQueueSize = inputQueMap_[strPort0_]->getSize();
        int nType = MONITOR_MODEL_INIT_STATE;
        if (iQueueSize > 5)
        {
            LogDebug << "iQueueSize: " << iQueueSize;
            g_bNoDealStepTwoFlag = true;
        }
        else if (g_bNoDealStepTwoFlag)
        {
            g_bNoDealStepTwoFlag = false;
        }

        //1. 无需动态检测
        if (!bNeedMoveDetectFlag_)
        {
            SingleDeviceProcess(pProcessData, MONITOR_MODEL_INIT_STATE);

            if (pProcessData->bIsEnd)
            {
                InitParam();
            }
            continue;
        }

        // 进行推理
        bool bGetTrainExist = false;
        uint64_t i64Time = 0;
        if (pProcessData->pData != nullptr && pProcessData->iSize != 0)
        {
            cv::Mat img; // BGR
            cv::Mat img_gray(pProcessData->iHeight, pProcessData->iWidth, CV_8UC1);
            cvtColor(cv::Mat(pProcessData->iHeight, pProcessData->iWidth, CV_8UC3, static_cast<uint8_t*>(pProcessData->pData.get())), img_gray, COLOR_BGR2GRAY);
            cvtColor(img_gray, img, COLOR_GRAY2BGR);

            auto start = std::chrono::system_clock::now(); // 计时开始
            float fReturnVal[STEP0_OUTPUT_ARRAY];
            memset(fReturnVal, 0x00, sizeof(fReturnVal));

            yolov8model.YoloV8InferenceModelGetType(img, fReturnVal, STEP0_OUTPUT_ARRAY * sizeof(float));
            // exit(0);

            float fScore = 0.0f;
            for(int n = 0; n < 4; n++){
                if(fReturnVal[n] > fScore){
                    fScore = fReturnVal[n];
                    nType = n;
                }
            }
//            LogDebug <<"模型得分  车头:"<< fReturnVal[0]<<"  无车:"<< fReturnVal[1]<<"  车尾:"<< fReturnVal[2]<<"  有车:"<< fReturnVal[3];
//            LogInfo<<"来车当前状态:"<< (nType == 0 ? "有车头" : (nType == 1 ? "无车"));
            switch (nType) {
                case 0:
                    LogDebug << iMoveDataNO_ * dataSourceCfg_.iSkipInterval << "来车状态：有车头";
                    break;
                case 1:
                    LogDebug << iMoveDataNO_ * dataSourceCfg_.iSkipInterval << "来车状态：无车";
                    break;
                case 2:
                    LogDebug << iMoveDataNO_ * dataSourceCfg_.iSkipInterval << "来车状态：车尾";
                    break;
                case 3:
                    LogDebug << iMoveDataNO_ * dataSourceCfg_.iSkipInterval << "来车状态：有车";
                    break;
            }

            if (MyYaml::GetIns()->GetBoolValue("gc_train_heard_detect"))
            {
                bGetTrainExist = ((nType == MONITOR_MODEL_TRAIN_BODY) || (nType == MONITOR_MODEL_TRAIN_HEAD) || (nType == MONITOR_MODEL_TRAIN_TAIL));
            }
            else
            {
                bGetTrainExist = ((nType == MONITOR_MODEL_TRAIN_BODY) || (nType == MONITOR_MODEL_TRAIN_TAIL));
            }

            auto end = std::chrono::system_clock::now();
            i64Time = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
        }

        if (bGetTrainExist == true)
        {
            iHasTrainNum_ = iHasTrainNum_ > 20 ? iHasTrainNum_ : iHasTrainNum_ + 1;
//            if (iHasTrainNum_ > 0) LogDebug << "当前有车, 计数：" << iHasTrainNum_;
        }
        else
        {
            iHasTrainNum_ = iHasTrainNum_ == 0 ? iHasTrainNum_ : iHasTrainNum_ - 1;
            if (iHasTrainNum_ == 0) LogInfo << "----- 当前无车 -----";
        }



        g_bHaveTrainFlag = bGetTrainExist;

        //有车开始识别
        if (iHasTrainNum_ > 0)
        {
            if (iStepInter_ != 1 && queProcessData_.size() < 3)
            {

                queProcessData_.push(pProcessData);
                LogDebug << "iStepInter_: " << iStepInter_ << " queSize:" << queProcessData_.size() << " continue";
                continue;
            }
            if (iStepInter_ != 1) this->sendComeTrain();
            iStepInter_ = 1;

        }
            //无车停止识别
        else
        {
            if (iStepInter_ == 1)
            {
                iStepInter_ = 2;
                this->sendEndTrain();
            }
            while (!queProcessData_.empty())
            {
                LogDebug << "while iStepInter_: " << iStepInter_ << " queSize:" << queProcessData_.size();
                queProcessData_.pop();
            }
        }
        //有车识别处理
        if (iStepInter_ != 0)
        {
            while (!queProcessData_.empty())
            {
                LogDebug << "while2 iStepInter_: " << iStepInter_ << " queSize:" << queProcessData_.size();
                std::shared_ptr<ProcessData> pProcessDataTemp = queProcessData_.front();
                queProcessData_.pop();
                pProcessDataTemp->iStatus = TRAINSTATUS_RUN;
                pProcessDataTemp->bIsEnd = false;
                SingleDeviceProcess(pProcessDataTemp, nType);
            }

            pProcessData->iStatus = TRAINSTATUS_RUN;
            pProcessData->bIsEnd = ((iStepInter_ == 2) ? true : false); //动态检测无车，设置列车结束标识

            SingleDeviceProcess(pProcessData, nType);

            if (iStepInter_ == 2)
            {
                // this->sendEndTrain();
                InitParam();
            }
        }
        nPreMonitorState = nType;
    }
}