VTrain/engine/Step1InferenceEngine/TrainStep1InferenceEngine.cpp

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

using namespace ai_matrix;

TrainStep1InferenceEngine::TrainStep1InferenceEngine() {}

TrainStep1InferenceEngine::~TrainStep1InferenceEngine() {}

APP_ERROR TrainStep1InferenceEngine::Init()
{
    strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";
    this->modelConfig_ = Config::getins()->getModelByTrainStep1Config();
    this->dataSourceConfig_ = Config::getins()->getDataSourceConfig();
    this->identifyConfig_ = Config::getins()->getIdentifyConfig();

    int iFolderExist = access(modelConfig_.strModelPath.c_str(), R_OK);
    if (iFolderExist == -1)
    {
        LogError << "模型：" << modelConfig_.strModelPath << " 不存在!";
        return false;
    }
    class_num = this->modelConfig_.vecClass.size();
    score_threshold = this->modelConfig_.fScoreThreshold;

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

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

APP_ERROR TrainStep1InferenceEngine::initModel()
{
    modelinfo.yolov5ClearityModelParam.uiClassNum = class_num;
    modelinfo.yolov5ClearityModelParam.uiClearNum = clear_num;
    modelinfo.yolov5ClearityModelParam.uiDetSize = det_size;
    modelinfo.yolov5ClearityModelParam.fScoreThreshold = score_threshold;
    modelinfo.yolov5ClearityModelParam.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 = yolov5model.YoloV5ClearityInferenceInit(&modelinfo,
                                                       strModelName,
                                                       this->modelConfig_.strModelPath);
    if (nRet != 0)
    {
        LogError << "YoloV5ClassifyInferenceInit nRet:" << nRet;
        return APP_ERR_COMM_READ_FAIL;
    }
    return APP_ERR_OK;
}

APP_ERROR TrainStep1InferenceEngine::DeInit()
{
    yolov5model.YoloV5ClearityInferenceDeinit();
    LogInfo << "Step1InferenceEngine DeInit ok";
    return APP_ERR_OK;
}

/**
 * 获取第1步得分最高框
 * inParam : std::vector<stDetection> &vecResult     推理符合结果
 * outParam: std::vector<stDetection> &vecResult     每个类别得分最高结果
 * return  : N/A
 */
void TrainStep1InferenceEngine::getMaxScoreResult(std::vector<stDetection> &vecResult)
{
    if (vecResult.size() < 2)
    {
        return;
    }
    std::map<Target, std::vector<stDetection>> mapResult;
    for (size_t i = 0; i < vecResult.size(); i++)
    {
        stDetection stDTemp = vecResult.at(i);
        if (stDTemp.class_id == 0)
        {
            mapResult[HEAD].emplace_back(stDTemp);
        }
        else if (stDTemp.class_id == 1)
        {
            mapResult[PRO].emplace_back(stDTemp);
        }
        else if ((stDTemp.class_id >= 2 && stDTemp.class_id <= 6) || stDTemp.class_id == 8 || stDTemp.class_id == 15)
        {
            mapResult[NUM].emplace_back(stDTemp);
        }
        else if (stDTemp.class_id >= 9 && stDTemp.class_id <= 17 && stDTemp.class_id != 15)
        {
            mapResult[TRAINSPACE].emplace_back(stDTemp);
        }
        else if (stDTemp.class_id == 18)
        {
            mapResult[SPACE].emplace_back(stDTemp);
        }
        else if (stDTemp.class_id == 7)
        {
            mapResult[CONTAINER].emplace_back(stDTemp);
        }
    }

    //清空之前的结果
    vecResult.clear();
    // 每个类别中，获取得分最高的框
    for (auto iter = mapResult.begin(); iter != mapResult.end(); iter++)
    {
        int iMaxPos = -1;
        for (size_t i = 0; i < iter->second.size(); i++)
        {
            if (iMaxPos == -1)
            {
                iMaxPos = i;
            }
            else if (iter->second.at(i).class_conf > iter->second.at(iMaxPos).class_conf)
            {
                iMaxPos = i;
            }
        }

        if (iMaxPos >= 0)
        {
            vecResult.emplace_back(iter->second.at(iMaxPos));
        }
    }
}

/**
* 设置大框类型
* inParam : PostSubData &postSubData   :推理结果
* outParam: PostSubData &postSubData   :推理结果
* return  : N/A
*/
void TrainStep1InferenceEngine::getTargetType(SingleData &singleData)
{
    if (singleData.iClassId == TRAIN_HEAD)
    {
        singleData.iTargetType = HEAD;
    }
    else if (singleData.iClassId == TRAIN_PRO)
    {
        singleData.iTargetType = PRO;
    }
    else if ((singleData.iClassId >= 2 && singleData.iClassId <= 6) ||
             singleData.iClassId == J_TRAIN_NUM ||
             singleData.iClassId == W_TRAIN_NUM)
    {
        singleData.iTargetType = NUM;
    }
    else if (singleData.iClassId >= 9 && singleData.iClassId <= 17 && singleData.iClassId != 15)
    {
        singleData.iTargetType = TRAINSPACE;
    }
    else if (singleData.iClassId == U_TRAIN_SPACE)
    {
        singleData.iTargetType = SPACE;
    }
    else if (singleData.iClassId == CONTAINERNUM)
    {
        singleData.iTargetType = CONTAINER;
    }
}

/**
* 过滤无效信息
* inParam : std::vector<stDetection> &vecRet            :识别结果数据
          : std::shared_ptr<ProcessData> pProcessData   :帧信息数据
* outParam: N/A
* return  : N/A
*/
void TrainStep1InferenceEngine::filterInvalidInfo(std::vector<stDetection> &vecInferenceResult,
                                                  std::shared_ptr<VTrainStep1Data> &pVTrainStep1Data)
{
    std::vector<stDetection> vecSpaceInfo;
    for (auto it = vecInferenceResult.begin(); it != vecInferenceResult.end();)
    {
//        LogDebug << " 帧:" << pVTrainStep1Data->iFrameId
//                 << " --iClassId:" << it->class_id
////                 << " iLine:" << it->clear_conf
//                 << " confidence=" << it->class_conf
//                 << " lx=" << it->bbox[0]
//                 << " ly=" << it->bbox[1]
//                 << " rx=" << it->bbox[2]
//                 << " ry=" << it->bbox[3]
//                 << " clear:" << it->clear_conf;
        // 根据配置文件中  设置的识别范围，过滤掉无效数据
        if (!(it->bbox[0] >= this->dataSourceConfig_.vecIdentifyAreas[0] &&
              it->bbox[1] >= this->dataSourceConfig_.vecIdentifyAreas[1] &&
              it->bbox[2] <= this->dataSourceConfig_.vecIdentifyAreas[2] &&
              it->bbox[3] <= this->dataSourceConfig_.vecIdentifyAreas[3]))
        {
//            LogDebug << "frameId:" << pVTrainStep1Data->iFrameId
//                     << " 类别:" << it->class_id << " 超出识别区域-识别区域：("
//                     << this->dataSourceConfig_.vecIdentifyAreas[0] << ","
//                     << this->dataSourceConfig_.vecIdentifyAreas[1] << "),("
//                     << this->dataSourceConfig_.vecIdentifyAreas[2] << ","
//                     << this->dataSourceConfig_.vecIdentifyAreas[2] << ")";
            it = vecInferenceResult.erase(it);
            continue;
        }

        // 如果设置了不识别车头，则去掉车头标记的大框 !this->identifyConfig_.bTrainHeardDetect &&
        if (it->class_id == TRAIN_HEAD)
        {
            LogDebug << "frameId:" << pVTrainStep1Data->iFrameId << " 过滤掉车头编号";
            it = vecInferenceResult.erase(it);
            continue;
        }

        // 去除车头时的非车头编号信息
        if(pVTrainStep1Data->iTrainStage == MONITOR_MODEL_TRAIN_HEAD )
        {
            if(it->class_id  != TRAIN_HEAD)
            {
                LogDebug << " 帧号:" << pVTrainStep1Data->iFrameId
                         << " 大类:" << it->class_id << " 识别于车头位置，无效！";
                it = vecInferenceResult.erase(it);
                continue;
            }
        }

        // 去除车尾的车头编号信息
        if (pVTrainStep1Data->iTrainStage != MONITOR_MODEL_TRAIN_HEAD)
        {
            if (it->class_id == TRAIN_HEAD)
            {
                LogDebug << " 帧号:" << pVTrainStep1Data->iFrameId
                         << " 大类:" << it->class_id << " 识别于非车头位置，无效！";
                it = vecInferenceResult.erase(it);
                continue;
            }
        }

        // 去除车尾 和 车头车体之间 的间隔信息
        if ((pVTrainStep1Data->iTrainStage == MONITOR_MODEL_TRAIN_TAIL || pVTrainStep1Data->iTrainStage == MONITOR_MODEL_HEAD_FIRST)
            && (it->class_id >= C_TRAIN_SPACE && it->class_id <= U_TRAIN_SPACE && it->class_id != W_TRAIN_NUM))
        {
            LogDebug << " frameId:" << pVTrainStep1Data->iFrameId
                     << " bigclassid:" << it->class_id
                     <<" 识别于车尾或者车头与车身交接部分,无效！";
            it = vecInferenceResult.erase(it);
            continue;
        }

        // 过滤掉识别于模型反馈无车状态下的所有大框信息
        if (pVTrainStep1Data->iTrainStage == MONITOR_MODEL_NO_TRAIN)
        {
            LogDebug << " frameId:" << pVTrainStep1Data->iFrameId
                     << " bigclassid:" << it->class_id
                     <<" 识别于模型反馈的无车状态下,无效！";
            it = vecInferenceResult.erase(it);
            continue;
        }

        //剔除高度大于宽的车号大框
        if (((it->class_id >= K_TRAIN_NUM && it->class_id <= NX_TRAIN_NUM)
             || it->class_id == J_TRAIN_NUM
             || it->class_id == W_TRAIN_NUM) &&
            (it->bbox[3] - it->bbox[1]) > (it->bbox[2] - it->bbox[0]))
        {
            LogWarn << " frameId:" << pVTrainStep1Data->iFrameId
                    << " bigclassid:" << it->class_id << " 过滤 高度大于宽度的车号";
            it = vecInferenceResult.erase(it);
            continue;
        }

        if (it->class_id == K_TRAIN_NUM)
        {
            int iCenterY = IMAGE_HEIGHT / 2;
            int iHeight0 = it->bbox[1] / 2 + it->bbox[3] / 2;
            if (iHeight0 > iCenterY) {
                LogWarn << "矿车编号大框在画面Y轴中线以下，帧号:"
                        << pVTrainStep1Data->iFrameId
                        << " 画面Y轴中心：" << iCenterY
                        << " 大框Y轴中心：" << iHeight0 ;
                //  << "[" << it->bbox[0] << "," << it->bbox[1] << "]"
                //  << "[" << it->bbox[2] << "," << it->bbox[3] << "]";
                it = vecInferenceResult.erase(it);
                continue;
            }
        }

        if (it->class_id >= C_TRAIN_SPACE && it->class_id <= U_TRAIN_SPACE && it->class_id != W_TRAIN_NUM)
        {
            vecSpaceInfo.emplace_back(*it);
        }
        ++it;
    }

    if (vecInferenceResult.size() <= 0) return;

    // 过滤与间隔X轴重合的其他大框
//    for (auto it = vecInferenceResult.begin(); it != vecInferenceResult.end();)
//    {
//        if (!((it->class_id >= 9 && it->class_id <= 17 && it->class_id != 15) || it->class_id == 18))
//        {
//            for (int i = 0; i < vecSpaceInfo.size(); i++)
//            {
//                if (
//                        (it->bbox[0] > vecSpaceInfo[i].bbox[0]
//                && it->bbox[0] < vecSpaceInfo[i].bbox[2])
//                ||
//                (it->bbox[2] > vecSpaceInfo[i].bbox[0]
//                    && it->bbox[2] < vecSpaceInfo[i].bbox[2])
//                    )
//                {
//                    LogWarn << "-- " << it->class_id << " _ " << vecSpaceInfo[i].class_id
//                    << "   " << it->bbox[0] << "," << it->bbox[2] << " || " << vecSpaceInfo[i].bbox[0] << "," << vecSpaceInfo[i].bbox[2];
//                    it = vecInferenceResult.erase(it);
//                    break;
//                }
//            }
//        }
//    }

    //主摄像头1帧如果只识别2个大框，如果非平车的车号和属性场景，则必有间隔框
    if (vecInferenceResult.size() > 2)
    {
        int iHeight0 = vecInferenceResult[0].bbox[1] / 2 + vecInferenceResult[0].bbox[3] / 2;
        int iHeight1 = vecInferenceResult[1].bbox[1] / 2 + vecInferenceResult[1].bbox[3] / 2;
        int iCenterY = IMAGE_HEIGHT / 2;
        if (iHeight0 < iCenterY && iHeight1 < iCenterY) //非平车
        {
            bool bHaveSpace = false;
            for (auto &it : vecInferenceResult)
            {
                if (it.class_id >= C_TRAIN_SPACE
                    && it.class_id <= U_TRAIN_SPACE
                    && it.class_id != W_TRAIN_NUM)
                {
                    bHaveSpace = true;
                }
            }

            if (!bHaveSpace)
            {
                LogDebug << " frameId:" << pVTrainStep1Data->iFrameId << " no space";
                vecInferenceResult.clear();
            }
//            if (!(vecInferenceResult[0].class_id >= C_TRAIN_SPACE
//                   && vecInferenceResult[0].class_id <= U_TRAIN_SPACE
//                   && vecInferenceResult[0].class_id != W_TRAIN_NUM)
//                && !(vecInferenceResult[1].class_id >= C_TRAIN_SPACE
//                      && vecInferenceResult[1].class_id <= U_TRAIN_SPACE
//                      && vecInferenceResult[1].class_id != W_TRAIN_NUM))
//            {
//                LogDebug << " frameId:" << pVTrainStep1Data->iFrameId << " no space";
//                vecInferenceResult.clear();
//            }
        }
    }
}

APP_ERROR TrainStep1InferenceEngine::Process()
{
    int iRet = APP_ERR_OK;

    while (!isStop_)
    {
        std::shared_ptr<void> pVoidData0 = nullptr;
        inputQueMap_[strPort0_]->pop(pVoidData0);
        if (nullptr == pVoidData0)
        {
            usleep(1000); //1ms
            continue;
        }

        std::shared_ptr<VTrainStep1Data> pVTrainStep1Data = std::static_pointer_cast<VTrainStep1Data>(pVoidData0);

        if (pVTrainStep1Data->cvImage.empty())
        {
            usleep(1000); //1ms
            continue;
        }
//        else
//        {
//            vector<int> compression_params;
//            compression_params.push_back(cv::IMWRITE_JPEG_QUALITY);  //选择jpeg
//            compression_params.push_back(100); //图片质量
//            cv::imwrite("./jpg/" + std::to_string(pVTrainStep1Data->iFrameId) + ".jpg", pVTrainStep1Data->cvImage, compression_params);
//
//        }

        //进行推理
        std::vector<stDetection> vecInferenceResult;
        yolov5model.YoloV5ClearityInferenceModel(pVTrainStep1Data->cvImage, vecInferenceResult);

        //过滤无效信息
        this->filterInvalidInfo(vecInferenceResult, pVTrainStep1Data);

        this->getMaxScoreResult(vecInferenceResult);

        std::shared_ptr<InferenceResultData> pInferenceResultData = std::make_shared<InferenceResultData>();

        pInferenceResultData->iFrameId = pVTrainStep1Data->iFrameId;
        pInferenceResultData->bIsEnd = pVTrainStep1Data->bIsEnd;
        pInferenceResultData->strTrainDate = pVTrainStep1Data->strTrainDate;
        pInferenceResultData->strTrainTime = pVTrainStep1Data->strTrainTime;

        for (size_t j = 0; j < vecInferenceResult.size(); j++)
        {
            /*
            [0:车头; 1:属性; 2:煤炭漏斗车(兖矿自备,枣矿自备); 3:敞车; 4:棚车; 5:罐车; 6:平车
            7:集装箱; 8:牲畜车; 9:敞车间隔; 10:自备车间隔; 11:平车间隔; 12:罐车间隔; 13:棚车车间隔;
            14:牲畜车间隔; 15:毒品车; 16: 毒品车间隔; 17:混合车厢间隔; 18:连接轴通用间隔; 19:集装箱号; 20:倒集装箱号]
            */
            if (vecInferenceResult[j].class_id < 0 || vecInferenceResult[j].class_id > 20)
            {
                continue;
            }
            SingleData singledata;
            singledata.iClassId = vecInferenceResult[j].class_id;
            singledata.fScore = vecInferenceResult[j].class_conf;
            singledata.fLTX = vecInferenceResult[j].bbox[0];
            singledata.fLTY = vecInferenceResult[j].bbox[1];
            singledata.fRBX = vecInferenceResult[j].bbox[2];
            singledata.fRBY = vecInferenceResult[j].bbox[3];
            singledata.fClear = vecInferenceResult[j].clear_id;

            this->getTargetType(singledata);
            pInferenceResultData->vecSingleData.emplace_back(singledata);

            LogDebug << " 帧:" << pInferenceResultData->iFrameId
                     << " --iClassId:" << singledata.iClassId
                     << " iLine:" << singledata.iLine
                     << " confidence=" << singledata.fScore
                     << " lx=" << singledata.fLTX
                     << " ly=" << singledata.fLTY
                     << " rx=" << singledata.fRBX
                     << " ry=" << singledata.fRBY
                     << " clear:" << singledata.fClear;
        }

        outputQueMap_[strPort0_]->push(std::static_pointer_cast<void>(pInferenceResultData), true);
    }
    return APP_ERR_OK;
}