VTrain/engine/SaveDebugImageEngine/SaveDebugImageEngine.cpp

//
// Created by nvidia on 24-10-14.
//

#include "SaveDebugImageEngine.h"

using namespace ai_matrix;

SaveDebugImageEngine::SaveDebugImageEngine() {}

SaveDebugImageEngine::~SaveDebugImageEngine() {}

APP_ERROR SaveDebugImageEngine::Init()
{

    strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";

    this->baseConfig_ = Config::getins()->getBaseConfig();
    vecCompressionParams_.push_back(cv::IMWRITE_JPEG_QUALITY);  //选择jpeg
    vecCompressionParams_.push_back(this->iPicQuality_);        //图片质量

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

APP_ERROR SaveDebugImageEngine::DeInit()
{
    LogInfo << "SaveDebugImageEngine DeInit ok";
    return APP_ERR_OK;
}

APP_ERROR SaveDebugImageEngine::Process()
{
    int iRet = APP_ERR_OK;
    while (!isStop_)
    {
        std::shared_ptr<void> pvoidd = nullptr;
        inputQueMap_[strPort0_]->pop(pvoidd);
        if (nullptr == pvoidd)
        {
            usleep(1000);
            continue;
        }

        if (!this->baseConfig_.bTestModel)
        {
            usleep(1000);
            continue;
        }

        std::shared_ptr<VStep2OutputData> pVStep2OutputData = std::static_pointer_cast<VStep2OutputData>(pvoidd);

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

        std::string strJsonPath = strDataDir + std::to_string(pVStep2OutputData->iFrameId) + ".json";
        std::string strImagePath = strDataDir + std::to_string(pVStep2OutputData->iFrameId) + ".jpg";

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

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

        std::string strDirection = "";
        switch (jvFrameInfo["direction"].asInt()) {
            case DIRECTION_LEFT:
                strDirection = "LEFT";
                break;
            case DIRECTION_RIGHT:
                strDirection = "RIGHT";
                break;
            default:
                strDirection = "...";
                break;
        }

        std::string strTrainStage = "";
        switch (jvFrameInfo["trainStage"].asInt()) {
            case MONITOR_MODEL_TRAIN_HEAD:
                strTrainStage = "train_head";
                break;
            case MONITOR_MODEL_HEAD_FIRST:
                strTrainStage = "head_and_carriage";
                break;
            case MONITOR_MODEL_NO_TRAIN:
                strTrainStage = "no_train";
                break;
            case MONITOR_MODEL_TRAIN_TAIL:
                strTrainStage = "train_tail";
                break;
            case MONITOR_MODEL_TRAIN_BODY:
                strTrainStage = "carriage";
                break;
            default:
                strTrainStage = "";
                break;
        }

        std::string strTrainStatus = "";
        switch (jvFrameInfo["train_status"].asInt()) {
            case TRAINSTATUS_RUN:
                strTrainStatus = "run";
                break;
            case TRAINSTATUS_STOP:
            case TRAINSTATUS_BACK:
                strTrainStatus = "Waiting";
                break;
            case TRAINSTATUS_NO:
            default:
                strTrainStatus = "";
                break;
        }

        cv::Mat image = cv::imread(strImagePath);

        if (image.empty())
        {
            LogWarn << "未能读取到需要标注的图像:" << strImagePath;
            continue;
        }

        std::stringstream ss;
        std::vector<std::string> vecTitle;
        vecTitle.emplace_back("FrameID:" + to_string(pVStep2OutputData->iFrameId));
        vecTitle.emplace_back("Direction:[" + strDirection + "]");
        vecTitle.emplace_back("TrainStage:[" + strTrainStage + "]");
        vecTitle.emplace_back("Status:[" + strTrainStatus + "]");
        vecTitle.emplace_back("TimeStamp:" + jvFrameInfo["timeStamp"].asString());

        int base_line = 0;
        cv::Point text_org(15, 45); // 文本起始点
        for (int i = 0; i < vecTitle.size(); ++i)
        {
            cv::Size text_size = cv::getTextSize(vecTitle.at(i), cv::FONT_HERSHEY_SIMPLEX, 1, 2, 0);
            cv::Point baseline_loc(text_org);
            baseline_loc.y += base_line + text_size.height;
            cv::putText(image, vecTitle.at(i), baseline_loc, cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 0, 255), 2);
            // 手动调整x位置，为下一个单词留出空间
            text_org.x += text_size.width + 10;
            text_org.x = text_org.x > IMAGE_WIDTH ? 15 : text_org.x;
            base_line++;
        }

        for (auto it = pVStep2OutputData->vecStep2ResultData.cbegin(); it != pVStep2OutputData->vecStep2ResultData.cend(); ++it)
        {
            bool bIsDivideSpace = false;
            cv::Scalar cvScalar = {0, 255, 0, 255};
            if (it->iClassId >= C_TRAIN_SPACE && it->iClassId <= U_TRAIN_SPACE && it->iClassId != W_TRAIN_NUM)
            {
                cvScalar = {0, 255, 255, 255};
                bIsDivideSpace = true;
            }
            float centerX = it->fLTX + (it->fRBX - it->fLTX)/2;
            float centerY = it->fLTY + (it->fRBY - it->fLTY)/2;
            cv::rectangle(image,
                          cv::Point(it->fLTX, it->fLTY),
                          cv::Point(it->fRBX, it->fRBY),
                          cvScalar, 2);
            cv::line(image,
                     cv::Point(centerX, it->fLTY-30), cv::Point(centerX, it->fRBY+30),
                     cvScalar, 1);
//            int iLine = 1;
            cv::Point linePoint(it->fLTX+10, it->fRBY+10);
            for (auto lineT = it->vecTransInfo.cbegin(); lineT != it->vecTransInfo.cend(); ++lineT)
            {
                cv::Size text_size = cv::getTextSize(lineT->strTmpResult, cv::FONT_HERSHEY_SIMPLEX, 1, 2, 0);
                linePoint.y += text_size.height + 5;
                cv::putText(image,
                            lineT->strTmpResult,
                            linePoint,
                            cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 0, 180), 2);
//                ++iLine;
            }

            linePoint.y += 35;
            if (!bIsDivideSpace)
            {
                cv::putText(image,
                            "train_index:" + std::to_string(it->iTrainIndex),
                            linePoint,
                            cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 0, 180), 2);
            }
        }

        if (!cv::imwrite(strImagePath, image, this->vecCompressionParams_))
        {
            LogError << "图片存储失败：" << strImagePath;
        }
    }
    return APP_ERR_OK;
}