Train_Identify/nvidia_ascend_engine/common_engine/TransEngine/TransContainerEngine.cpp

#include "TransContainerEngine.h"
#include "myutils.h"
#include <regex>

using namespace ai_matrix;

namespace
{
    //按照x坐标排列
    bool CompareX(const SingleData &v1, const SingleData &v2)
    {
        return (v1.fLTX < v2.fLTX);
    }
    //按照y坐标排列
    bool CompareY(const SingleData &v1, const SingleData &v2)
    {
        return (v1.fLTY < v2.fLTY);
    }
}

TransContainerEngine::TransContainerEngine() {}

TransContainerEngine::~TransContainerEngine() {}

APP_ERROR TransContainerEngine::Init()
{
    bUseEngine_ = MyUtils::getins()->ChkIsHaveTarget("CONTAINER");
    if (!bUseEngine_)
    {
        LogWarn << "engineId_:" << engineId_ << " not use engine";
        return APP_ERR_OK;
    }
    strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";

    ai_matrix::DataSourceConfig mainCfg = MyYaml::GetIns()->GetDataSourceConfigById(0);

    //获取几个摄像头识别集装箱
    std::map<int, ai_matrix::DataSourceConfig> mapUseDataSouceCfg = MyYaml::GetIns()->GetUseDataSourceConfig();
    for (auto iter = mapUseDataSouceCfg.begin(); iter != mapUseDataSouceCfg.end(); iter++)
    {
        //顶部集装箱模型
        if (iter->second.strTarget.find("CONTAINER_T") != std::string::npos)
        {
            LogDebug << "sourceid:" << iter->first << " deal CONTAINER_T";
            mapModelType_[iter->first] = 1;
        }
        //侧部集装箱模型
        else if (iter->second.strTarget.find("CONTAINER") != std::string::npos)
        {
            LogDebug << "sourceid:" << iter->first << " deal CONTAINER";
            mapModelType_[iter->first] = 0;

            //主摄像头对侧的摄像头
            if (mainCfg.iLeftFirst != iter->second.iLeftFirst)
            {
                mapModelType_[iter->first] = 2;
            }
        }
    }

    //侧部和顶部集装箱模型(除模型外，其他参数一样)使用侧部模型参数即可。
    modelConfig_ = MyYaml::GetIns()->GetModelConfig("StepTwoContainerEngine");

    //读取模型参数信息文件
    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;
    }
    for (int i = 0; i < jvModelInfo["class"].size(); i++)
    {
        vecClassNames_.push_back(jvModelInfo["class"][i].asString());
    }

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

APP_ERROR TransContainerEngine::DeInit()
{
    if (!bUseEngine_)
    {
        LogWarn << "engineId_:" << engineId_ << " not use engine";
        return APP_ERR_OK;
    }
    LogInfo << "TransContainerEngine DeInit ok";
    return APP_ERR_OK;
}

/**
* push数据到队列，队列满时则休眠一段时间再push
* inParam : const std::string strPort                           push的端口
          : const std::shared_ptr<ProcessData> &pProcessData    push的数据
* outParam: N/A
* return  : N/A
*/
void TransContainerEngine::PushData(const std::string &strPort, const std::shared_ptr<ProcessData> &pProcessData)
{
    while (true)
    {
        int iRet = outputQueMap_[strPort]->push(std::static_pointer_cast<void>(pProcessData));
        if (iRet != 0)
        {
            LogDebug << "sourceid:" << pProcessData->iDataSource << " frameid:" << pProcessData->iFrameId << " push fail iRet:" << iRet;
            if (iRet == 2)
            {
                usleep(10000); // 10ms
                continue;
            }
        }
        break;
    }
}

/**
* 集装箱框内容转换
* inParam : std::map<int, std::vector<SingleData>> &mapLine
* outParam: TransSubData &transSubData
* return  : N/A
*/
void TransContainerEngine::TransContainer(TransSubData &transSubData, std::map<int, std::vector<SingleData>> &mapLine)
{
    std::string strContainerNo;
    //箱主代码+设备识别码（4位英文字母）
    if (mapLine.find(0) != mapLine.end())
    {
        TransInfo Info0;
        Info0.iLine = 0;
        std::string strTemp = "";
        for (auto j = 0; j < mapLine.at(0).size(); j++)
        {
            //过滤非字母
            std::string strOne = vecClassNames_.at(mapLine.at(0).at(j).iClassId);
            if (strOne[0] < 'A' || strOne[0] > 'Z')
            {
                LogWarn << "engineId:" << engineId_ << " " << strOne << " not A-Z in container_0";
                continue;
            }
            strTemp += strOne;
            Info0.vecValue.emplace_back(strOne);
            Info0.vecScore.emplace_back(mapLine.at(0).at(j).fScore);
            transSubData.fScoreSum += mapLine.at(0).at(j).fScore;
        }
        if (strTemp.length() == 4)
        {
            Info0.IsChkFlag = true;  
        }
        strContainerNo += strTemp;
        transSubData.vecTransInfo.emplace_back(Info0);
    }

    //箱号（6位数字）
    if (mapLine.find(1) != mapLine.end())
    {
        TransInfo Info1;
        Info1.iLine = 1;
        std::string strTemp = "";
        for (auto j = 0; j < mapLine.at(1).size(); j++)
        {
            //过滤非数字
            std::string strOne = vecClassNames_.at(mapLine.at(1).at(j).iClassId);
            if (strOne[0] < '0' || strOne[0] > '9')
            {
                LogWarn << "engineId:" << engineId_ << " " << strOne << " not digit in container_1";
                continue;
            }
            strTemp += strOne;
            Info1.vecValue.emplace_back(strOne);
            Info1.vecScore.emplace_back(mapLine.at(1).at(j).fScore);
            transSubData.fScoreSum += mapLine.at(1).at(j).fScore;
        }
        if (strTemp.length() == 6)
        {
            Info1.IsChkFlag = true;
        }
        strContainerNo += strTemp;
        transSubData.vecTransInfo.emplace_back(Info1);
    }

    //校验码（1位数字）
    if (mapLine.find(2) != mapLine.end())
    {
        TransInfo Info2;
        Info2.iLine = 2;
        std::string strTemp = "";
        for (auto j = 0; j < mapLine.at(2).size(); j++)
        {
            //过滤非数字
            std::string strOne = vecClassNames_.at(mapLine.at(2).at(j).iClassId);
            if (strOne[0] < '0' || strOne[0] > '9')
            {
                LogWarn << "engineId:" << engineId_ << " " << strOne << " not digit in container_2";
                continue;
            }
            strTemp += strOne;
            Info2.vecValue.emplace_back(strOne);
            Info2.vecScore.emplace_back(mapLine.at(2).at(j).fScore);
            transSubData.fScoreSum += mapLine.at(2).at(j).fScore;
        }
        if (strTemp.length() == 1)
        {
            Info2.IsChkFlag = true;
        }
        strContainerNo += strTemp;
        transSubData.vecTransInfo.emplace_back(Info2);
    }

    /*2022-04-02反馈不需要该识别内容，暂时屏蔽
    //集装箱型号（4位 字母或数字）
    if (mapLine.find(3) != mapLine.end())
    {

        TransInfo Info3;
        Info3.iLine = 3;
        std::string strTemp = "";
        for (auto j = 0; j < mapLine.at(3).size(); j++)
        {
            strTemp += vecClassNames_.at(mapLine.at(3).at(j).iClassId);
            Info3.vecValue.emplace_back(vecClassNames_.at(mapLine.at(3).at(j).iClassId));
            Info3.vecScore.emplace_back(mapLine.at(3).at(j).fScore);
            //transSubData.fScoreSum += mapLine.at(3).at(j).fScore;  //顶部集装箱号没有该内容，因此这块不因加入总得分。
        }
        Info3.IsChkFlag = true;
        transSubData.vecTransInfo.emplace_back(Info3);
    }
    */

    transSubData.strAllValue = strContainerNo;
}

/**
* 根据大框位置，重置总得分transSubData.fScoreSum
* inParam : TransSubData &transSubData
          : std::shared_ptr<ProcessData> pProcessData
* outParam: 
* return  : N/A
*/
void TransContainerEngine::ReSetScoreSumByPosition(TransSubData &transSubData, std::shared_ptr<ProcessData> pProcessData)
{
    /*
    侧面数据来源：第一步大框位置X轴中心坐标离图片中心越近得分越高 
    顶部数据来源：第一步大框位置Y轴中心坐标离图片中心越近得分越高
    因图片宽大于高，因此侧部图片的得分会大于顶部，都识别的情况下会优先使用侧部图片。
    */
    int iImgCenterX = pProcessData->iWidth / 2;
    int iImgCenterY = pProcessData->iHeight / 2;

    //大框X轴中心点坐标
    int iCenterX = transSubData.step1Location.fLTX + (transSubData.step1Location.fRBX - transSubData.step1Location.fLTX) / 2;
    int iCenterY = transSubData.step1Location.fLTY + (transSubData.step1Location.fRBY - transSubData.step1Location.fLTY) / 2;

    //LogDebug << "fScoreSum:" << transSubData.fScoreSum;
    if (mapModelType_[pProcessData->iDataSource] == 0)
    {
        transSubData.fScoreSum = pProcessData->iWidth - abs(iImgCenterX - iCenterX);
    }
    else if (mapModelType_[pProcessData->iDataSource] == 1)
    {
        transSubData.fScoreSum = pProcessData->iHeight - abs(iImgCenterY - iCenterY);
    }
    else if (mapModelType_[pProcessData->iDataSource] == 2)
    {
        //主摄像头另一侧得分稍低些，即优先使用主摄像头识别数据。
        transSubData.fScoreSum = pProcessData->iWidth - abs(iImgCenterX - iCenterX) - 300;
    }

    LogDebug << "frameid:" << pProcessData->iFrameId << " iImgCenterX:" << iImgCenterX << " iImgCenterY:" << iImgCenterY
             << " iCenterX:" << iCenterX << " iCenterY:" << iCenterY << " fScoreSum:" << transSubData.fScoreSum;
}

APP_ERROR TransContainerEngine::Process()
{
    if (!bUseEngine_)
    {
        LogWarn << "engineId_:" << engineId_ << " not use engine";
        return APP_ERR_OK;
    }

    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<ProcessData> pProcessData = std::static_pointer_cast<ProcessData>(pVoidData0);
        //第2步后处理结果
        std::shared_ptr<PostData> pPostData = std::static_pointer_cast<PostData>(pProcessData->pVoidData);
        
        //组织输出数据
        std::shared_ptr<TransData> pTransData = std::make_shared<TransData>();
        for (size_t i = 0; i < pPostData->vecPostSubData.size(); i++)
        {
            PostSubData postSubData = pPostData->vecPostSubData[i];

            //按字段代号分类
            std::map<int, std::vector<SingleData>> mapLine;
            for (size_t j = 0; j < postSubData.vecSingleData.size(); j++)
            {
                mapLine[postSubData.vecSingleData.at(j).iLine].push_back(postSubData.vecSingleData.at(j));
            }

            //每一行按x或y坐标排序
            uint32_t iWidth = postSubData.step1Location.fRBX - postSubData.step1Location.fLTX;
            uint32_t iHeight = postSubData.step1Location.fRBY - postSubData.step1Location.fLTY;
            bool bSortByX = iHeight >= (iWidth * 1.5) ? false : true;   //竖排集装箱按y坐标排序
            for (auto it = mapLine.begin(); it != mapLine.end(); it++)
            {
                if (bSortByX)
                {
                    std::sort(it->second.begin(), it->second.end(), CompareX);
                }
                else
                {
                    std::sort(it->second.begin(), it->second.end(), CompareY);
                }

                std::string strTemp = "";
                for (auto j = 0; j < it->second.size(); j++)
                {
                    strTemp += vecClassNames_.at(it->second.at(j).iClassId);
                }
                LogDebug << "dataSource:" << pProcessData->iDataSource << " frameid:" << pProcessData->iFrameId
                         << " bigclassid:" << postSubData.iBigClassId << " line:" << it->first << "," << strTemp;
            }

            //非集装箱大类不处理
            if (postSubData.iBigClassId != 0 && postSubData.iBigClassId != 1)
            {
                continue;
            }
            
            TransSubData transSubData;
            transSubData.iBigClassId = postSubData.iBigClassId;
            transSubData.iCarXH = postSubData.iCarXH;
            transSubData.step1Location = postSubData.step1Location;
            TransContainer(transSubData, mapLine);

            //根据大框位置，重置总得分transSubData.fScoreSum
            ReSetScoreSumByPosition(transSubData, pProcessData);

            pTransData->vecTransSubData.emplace_back(transSubData);
        }

        pProcessData->pVoidData = std::static_pointer_cast<void>(pTransData);
        //push端口1，集装箱选优处理
        PushData(strPort0_, pProcessData);
    }

    return APP_ERR_OK;
}