VTrain/engine/Step2DataReadEngine/TrainStep2DataReadEngine.cpp

#include "TrainStep2DataReadEngine.h"

using namespace ai_matrix;

TrainStep2DataReadEngine::TrainStep2DataReadEngine() {}

TrainStep2DataReadEngine::~TrainStep2DataReadEngine() {}

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

    this->baseConfig_ = Config::getins()->getBaseConfig();

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

APP_ERROR TrainStep2DataReadEngine::DeInit()
{

    LogInfo << "TrainStep2DataReadEngine DeInit ok";
    return APP_ERR_OK;
}

bool TrainStep2DataReadEngine::readJson(std::string &strFilePath, Json::Value &jvInfo)
{
    if (access(strFilePath.c_str(), F_OK) != 0)
    {
        LogWarn << "文件:" << strFilePath << " 不存在";
        return false;
    }

    if (!FileUtil::getins()->readJsonInfo(jvInfo, strFilePath))
    {
        LogError << "读取json文件失败:" << strFilePath;
        return false;
    }
    return true;
}

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

    while (!isStop_)
    {
        //获取主摄像头检测的状态
        std::shared_ptr<void> pVoidData0 = nullptr;
        iRet = inputQueMap_[strPort0_]->pop(pVoidData0);

        if (nullptr == pVoidData0)
        {
            usleep(1000);
            continue;
        }
        std::shared_ptr<PartionInfo> pPartionInfo = std::static_pointer_cast<PartionInfo>(pVoidData0);

//        pPartionInfo->iStartframe++;
//        pPartionInfo->iEndframe--;

        std::string strDataDir = this->baseConfig_.strDebugResultPath + "/"
                + pPartionInfo->strTrainDate + "/"
                + StringUtil::getins()->replace_all_distinct(pPartionInfo->strTrainTime, ":", "-") + "/"
                + "jpg/";

        for (int iFrameId = pPartionInfo->iStartframe; iFrameId <= pPartionInfo->iEndframe; ++iFrameId)
        {
            std::string strJsonPath = strDataDir + std::to_string(iFrameId) + ".json";
            std::string strImagePath = strDataDir + std::to_string(iFrameId) + ".jpg";

            Json::Value jvFrameInfo;
            if (!FileUtil::getins()->readJsonInfo(jvFrameInfo, strJsonPath))
            {
                continue;
            }

            Json::Value jvSpaceInfo;
            if (jvFrameInfo.isMember("divide_space"))
            {
                jvSpaceInfo = jvFrameInfo["divide_space"];
            }

            bool bSpace_front = false;
            if ((iFrameId - pPartionInfo->iStartframe) < (pPartionInfo->iEndframe - iFrameId))
            {
                bSpace_front = true;
            }

            std::vector<SingleData> vecSingleData;

            for (int i = 0; i < jvFrameInfo["step1"].size(); i++)
            {
                Json::Value jvStep1 = jvFrameInfo["step1"][i];

                if (!jvSpaceInfo.isNull())
                {
                    if (jvFrameInfo["direction"].asInt() == DIRECTION_LEFT)
                    {
                        if (
                                ((bSpace_front && jvStep1["ltx"].asFloat() > jvSpaceInfo["ltx"].asFloat())
                                 ||
                                 (!bSpace_front && jvStep1["rbx"].asFloat() < jvSpaceInfo["rbx"].asFloat()))
                                )
                        {
                            jvStep1["train_index"] = pPartionInfo->iTrainIndex;
                        }
                        if (bSpace_front && jvStep1["ltx"].asFloat() < jvSpaceInfo["ltx"].asFloat())
                        {
                            jvStep1["train_index"] = pPartionInfo->iTrainIndex - 1;
                        }
                        if (!bSpace_front && jvStep1["rbx"].asFloat() > jvSpaceInfo["rbx"].asFloat())
                        {
                            jvStep1["train_index"] = pPartionInfo->iTrainIndex + 1;
                        }
                    }

                    if (jvFrameInfo["direction"].asInt() == DIRECTION_RIGHT)
                    {
                        if (
                                ((bSpace_front && jvStep1["ltx"].asFloat() < jvSpaceInfo["ltx"].asFloat())
                                 ||
                                 (!bSpace_front && jvStep1["rbx"].asFloat() > jvSpaceInfo["rbx"].asFloat()))
                                )
                        {
                            jvStep1["train_index"] = pPartionInfo->iTrainIndex;
                        }
                        if (bSpace_front && jvStep1["rbx"].asFloat() > jvSpaceInfo["rbx"].asFloat())
                        {
                            jvStep1["train_index"] = pPartionInfo->iTrainIndex - 1;
                        }
                        if (!bSpace_front && jvStep1["ltx"].asFloat() < jvSpaceInfo["ltx"].asFloat())
                        {
                            jvStep1["train_index"] = pPartionInfo->iTrainIndex + 1;
                        }
                    }
                }
                else
                {
                    jvStep1["train_index"] = pPartionInfo->iTrainIndex;
                }
                jvFrameInfo["step1"][i] = jvStep1;

                SingleData singleData;
                singleData.iTrainIndex = jvStep1["train_index"].asInt();
                singleData.iTargetType = jvStep1["target_type"].asInt();
                singleData.iClassId = jvStep1["classid"].asInt();
                singleData.fLTX = jvStep1["ltx"].asFloat();
                singleData.fLTY = jvStep1["lty"].asFloat();
                singleData.fRBX = jvStep1["rbx"].asFloat();
                singleData.fRBY = jvStep1["rby"].asFloat();
                vecSingleData.emplace_back(singleData);
            }

            if (!FileUtil::getins()->writeJsonInfo(jvFrameInfo, strJsonPath))
            {
                continue;
            }

            std::shared_ptr<VStep2InputData> pVStep2InputData = std::make_shared<VStep2InputData>();
            pVStep2InputData->strTrainDate = pPartionInfo->strTrainDate;
            pVStep2InputData->strTrainTime = pPartionInfo->strTrainTime;
            pVStep2InputData->vecSingleData = vecSingleData;
            pVStep2InputData->iFrameId = iFrameId;
            pVStep2InputData->cvImage = cv::imread(strImagePath);
            pVStep2InputData->bIsEnd = !jvFrameInfo.isMember("isEnd") || jvFrameInfo["isEnd"].asBool();

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