VTrain/engine/MoveEngine/MoveEngine.cpp

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

using namespace ai_matrix;

MoveEngine::MoveEngine() {}

MoveEngine::~MoveEngine() {}

APP_ERROR MoveEngine::Init()
{
    strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";
    strPort1_ = engineName_ + "_" + std::to_string(engineId_) + "_1";

    this->modelConfig_ = Config::getins()->getModelByMoveConfig();
    this->identifyConfig_ = Config::getins()->getIdentifyConfig();
    this->baseConfig_ = Config::getins()->getBaseConfig();
    this->dataSourceConfig_ = Config::getins()->getDataSourceConfig();

    if (identifyConfig_.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;
        }
    }

    InitParam();

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

APP_ERROR MoveEngine::InitModel()
{
    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_.strModelPath);
    if (nRet != 0)
    {
        LogInfo << "YoloV5ClassifyInferenceInit nRet:" << nRet;
        return APP_ERR_COMM_READ_FAIL;
    }
    return APP_ERR_OK;
}

APP_ERROR MoveEngine::ReadModelInfo()
{
    int iFolderExist = access(modelConfig_.strModelPath.c_str(), R_OK);
    if (iFolderExist == -1)
    {
        LogError << "ModelPath " << modelConfig_.strModelPath << " doesn't exist or read failed!";
        return APP_ERR_COMM_NO_EXIST;
    }
    return APP_ERR_OK;
}

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


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

void MoveEngine::sendComeTrain() {
    std::string message = "{\"cometime\":\"" + this->strTrainDate_ + " " + this->strTrainTime_ + "\",\"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->strTrainTime_ + "\",\"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 iType)
{
    if(iMoveDataNO_ == 1 || strTrainDate_.empty())
    {
        strTrainDate_ = TimeUtil::getins()->getDate();
        strTrainTime_ = TimeUtil::getins()->getTime();
        //创建该该列车的数据存储目录(防止后续多线程创建报错)
        std::string strTrainPath = baseConfig_.strDebugResultPath + "/"
                + strTrainDate_ + "/"
                + StringUtil::getins()->replace_all_distinct(strTrainTime_, ":", "-") + "/jpg/";
        FileUtil::getins()->createDirPath(strTrainPath);
        std::string strBestImgPath = baseConfig_.strResultPath + "/"
                + strTrainDate_ + "/"
                + StringUtil::getins()->replace_all_distinct(strTrainTime_, ":", "-") + "/";
        FileUtil::getins()->createDirPath(strBestImgPath);
    }
    if (g_come_direction !=0 && (g_come_direction != this->dataSourceConfig_.iDirection && this->dataSourceConfig_.iDirection))
    {
        pProcessData->bIsTrainEnd = true;
    }
    if (pProcessData->bIsTrainEnd && this->iMoveDataNO_ == 1)
    {
        return;
    }

    std::string strFilePath = baseConfig_.strDebugResultPath + "/"
            + strTrainDate_ + "/"
            + StringUtil::getins()->replace_all_distinct(strTrainTime_, ":", "-") + "/jpg";

//    LogInfo << "-- " << iMoveDataNO_ << " --  " << strFilePath;
    // 通知存储来车检测结果
    std::shared_ptr<VMoveInfo> pVMoveInfo = std::make_shared<VMoveInfo>();
    pVMoveInfo->strFilePath = strFilePath;
    pVMoveInfo->strFileName = std::to_string(iMoveDataNO_) + ".json";
    pVMoveInfo->iFrameId = iMoveDataNO_; //当前帧号
    pVMoveInfo->i64TimeStamp = pProcessData->sourceFrameData.i64TimeStamp;
    pVMoveInfo->bIsEnd = pProcessData->bIsTrainEnd;
    pVMoveInfo->strTrainDate = strTrainDate_;
    pVMoveInfo->strTrainTime = strTrainTime_;
    pVMoveInfo->iTrainStage = iType;
    outputQueMap_[strPort1_]->push(std::static_pointer_cast<void>(pVMoveInfo), true);

    // push端口,存图
    std::shared_ptr<SaveImgData> pSaveImgData = std::make_shared<SaveImgData>();
    pSaveImgData->strFilePath = strFilePath;
    pSaveImgData->strFileName = std::to_string(iMoveDataNO_) + ".jpg";
    pSaveImgData->strTrainStage = this->intTrainStage_2_str(iType);
    pSaveImgData->frameData.pData = pProcessData->sourceFrameData.pData;
    pSaveImgData->frameData.iSize = pProcessData->sourceFrameData.iSize;
    pSaveImgData->frameData.iWidth = pProcessData->dataSourceInfo.iWidth;
    pSaveImgData->frameData.iHeight = pProcessData->dataSourceInfo.iHeight;
    pSaveImgData->frameData.i64TimeStamp = pProcessData->sourceFrameData.i64TimeStamp;
    pSaveImgData->bIsEnd = pProcessData->bIsTrainEnd;
    outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pSaveImgData), true);

    if (pProcessData->bIsTrainEnd)
    {
        this->InitParam();
    }
}

std::string MoveEngine::intTrainStage_2_str(int iTrainStage)
{
    switch (iTrainStage) {
        case MONITOR_MODEL_TRAIN_HEAD:
            return "train_head";
        case MONITOR_MODEL_HEAD_FIRST:
            return "head_and_carriage";
        case MONITOR_MODEL_NO_TRAIN:
            return "no_train";
        case MONITOR_MODEL_TRAIN_TAIL:
            return "train_tail";
        case MONITOR_MODEL_TRAIN_BODY:
            return "carriage";
        default:
            return "";
    }
}

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 iType = MONITOR_MODEL_INIT_STATE;

        //1. 无需动态检测
        if (!this->identifyConfig_.bNeedMoveDetectFlag)
        {
            SingleDeviceProcess(pProcessData, MONITOR_MODEL_INIT_STATE);

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

        // 进行推理

        if (!pProcessData->sourceFrameData.pData || !pProcessData->sourceFrameData.iSize)
        {
            continue;
        }

        cv::Mat img; // BGR
        cv::Mat img_gray(pProcessData->dataSourceInfo.iHeight, pProcessData->dataSourceInfo.iWidth, CV_8UC1);
        cvtColor(cv::Mat(pProcessData->dataSourceInfo.iHeight, pProcessData->dataSourceInfo.iWidth, CV_8UC3, static_cast<uint8_t*>(pProcessData->sourceFrameData.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 < STEP0_OUTPUT_ARRAY; n++){
            if(fReturnVal[n] > fScore){
                fScore = fReturnVal[n];
                iType = n;
            }
        }
//            LogDebug <<"模型得分  车头:"<< fReturnVal[0]<<"  无车:"<< fReturnVal[1]<<"  车尾:"<< fReturnVal[2]<<"  有车:"<< fReturnVal[3];
//            LogInfo<<"来车当前状态:"<< (nType == 0 ? "有车头" : (nType == 1 ? "无车"));
        switch (iType) {
            case MONITOR_MODEL_INIT_STATE:
                LogDebug << "来车检测初始状态";
                break;
            case MONITOR_MODEL_TRAIN_HEAD:
                LogDebug << "来车状态：有车头";
                break;
            case MONITOR_MODEL_HEAD_FIRST:
                LogDebug << "来车状态：车头与车身之间";
            case MONITOR_MODEL_NO_TRAIN:
//                    LogDebug << "来车状态：无车";
                break;
            case MONITOR_MODEL_TRAIN_TAIL:
                LogDebug << "来车状态：车尾";
                break;
            case MONITOR_MODEL_TRAIN_BODY:
                LogDebug << "来车状态：车身";
                break;
            default:
                LogWarn << "来车状态：未定义状态";
                break;
        }

        bool bGetTrainExist = false;

        if (this->identifyConfig_.bTrainHeardDetect)
        {
            bGetTrainExist = (iType == MONITOR_MODEL_TRAIN_BODY
                              || iType == MONITOR_MODEL_TRAIN_HEAD
                              || iType == MONITOR_MODEL_HEAD_FIRST
                              || iType == MONITOR_MODEL_TRAIN_TAIL);
        }
        else
        {
            bGetTrainExist = (iType == MONITOR_MODEL_TRAIN_BODY
                              || iType == MONITOR_MODEL_TRAIN_TAIL
                              || iType == MONITOR_MODEL_HEAD_FIRST);
        }

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

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

        if (iHasTrainNum_ > 0)
            //有车开始识别
        {
//            if (iStepInter_ != 1) this->sendComeTrain();
            iStepInter_ = 1;

        }
        else
            //无车停止识别
        {
            if (iStepInter_ == 1)
            {
                iStepInter_ = 2;
//                this->sendEndTrain();
            }
        }
        //有车识别处理
        if (iStepInter_ != 0)
        {
            pProcessData->bIsTrainEnd = (iStepInter_ == 2 ? true : false); //动态检测无车，设置列车结束标识

            SingleDeviceProcess(pProcessData, iType);

            iMoveDataNO_++;

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