Train_Identify/nvidia_ascend_engine/common_engine/SaveEngine/SaveImgEngine.cpp

#include "SaveImgEngine.h"


using namespace ai_matrix;

SaveImgEngine::SaveImgEngine() {}

SaveImgEngine::~SaveImgEngine() {}

APP_ERROR SaveImgEngine::Init()
{
    bUseEngine_ = true;
    dataSourceConfig_ = MyYaml::GetIns()->GetDataSourceConfigById(engineId_); //获取摄像机参数
    if (!dataSourceConfig_.bUse)
    {
        bUseEngine_ = false;
        LogWarn << "engineId_:" << engineId_ << " not use engine";
        return APP_ERR_OK;
    }

    strPort0_ = engineName_ + "_" + std::to_string(engineId_) + "_0";
    iPicQuality = MyYaml::GetIns()->GetIntValue("gc_save_pic_quality");
    iPushDirection_ = MyYaml::GetIns()->GetIntValue("gc_push_direction");

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

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

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

APP_ERROR SaveImgEngine::Process()
{
    if (!bUseEngine_)
    {
        LogWarn << "engineId_:" << engineId_ << " not use engine";
        return APP_ERR_OK;
    }
    
    vector<int> compression_params;
	compression_params.push_back(cv::IMWRITE_JPEG_QUALITY);  //选择jpeg
	compression_params.push_back(iPicQuality);              //图片质量
    iDirection_ = DIRECTION_UNKNOWN;

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

        //如果设置了方向，则方向不对直接过滤，但结束帧不能过滤，需流转到后面Engine，保证后面处理正确。
        if (iDirection_ == DIRECTION_UNKNOWN)
        {
            size_t iPos = pSaveImgData->strImgPath.rfind("/");
            if (iPos != std::string::npos)
            {
                std::string strFilePath = pSaveImgData->strImgPath.substr(0, iPos + 1) + "direction.txt";
                Json::Value jvDirectionInfo;
                if (MyUtils::getins()->ReadJsonInfo(jvDirectionInfo, strFilePath, 0))
                {
                    iDirection_ = jvDirectionInfo["direction"].asInt();
                }
            }
        }
        if (!pSaveImgData->bIsEnd)
        {
            if (!(iDirection_ == DIRECTION_UNKNOWN || iDirection_ == iPushDirection_ || iPushDirection_ == DIRECTION_UNKNOWN))
            {
                LogDebug << "image:" << pSaveImgData->strImgPath << "/" << pSaveImgData->strImgName << " no sava continue.";
                continue;
            }
        }
        else
        {
            iDirection_ = DIRECTION_UNKNOWN;
        }

        //图片绝对路径
        if (pSaveImgData->strImgPath.back() != '/')
        {
            pSaveImgData->strImgPath += "/";
        }
        std::string strImgFilePath = pSaveImgData->strImgPath + pSaveImgData->strImgName;
        std::string strTxtFilePath = pSaveImgData->strImgPath + std::to_string(pSaveImgData->iFrameId) + ".txt";

        //1. 图片编码
        uint32_t iWidth = 0;
        uint32_t iHeight = 0;
        int iRate = 0;
        int iStatus = TRAINSTATUS_RUN;

        //2. 创建保存路径
        if (!MyUtils::getins()->CreateDirPath(pSaveImgData->strImgPath))
        {
            LogError << "engineId:" << engineId_ << " CreateDirPath err img:" << strImgFilePath;
            continue;
        }

        // 3.保存图片
        
        if (pSaveImgData->pData != nullptr && pSaveImgData->iSize != 0)
        {
            cv::Mat matBGR(pSaveImgData->iHeight, pSaveImgData->iWidth, CV_8UC3, static_cast<uint8_t *>(pSaveImgData->pData.get()));   //RGB  
            iWidth = pSaveImgData->iWidth;
            iHeight = pSaveImgData->iHeight;
            
            // cv::Mat mtOutImage;
            // cv::cvtColor(cvimg, mtOutImage, cv::COLOR_RGB2BGR);

            cv::imwrite(strImgFilePath, matBGR, compression_params);

            //图片旋转
            if (dataSourceConfig_.iRotate != 0)
            {
                cv::Mat cvRotateImg;
                if (dataSourceConfig_.iRotate == 90)        //顺时针旋转90
                {
                    cv::transpose(matBGR, cvRotateImg);
                    cv::flip(cvRotateImg, cvRotateImg, 1);
                    iWidth = pSaveImgData->iHeight;
                    iHeight = pSaveImgData->iWidth;
                }
                else if (dataSourceConfig_.iRotate == 180)  //顺时针旋转180
                {
                    cv::flip(matBGR, cvRotateImg, -1);
                }
                else if (dataSourceConfig_.iRotate == 270)  //顺时针旋转270
                {
                    cv::transpose(matBGR, cvRotateImg);
                    cv::flip(cvRotateImg, cvRotateImg, 0);
                    iWidth = pSaveImgData->iHeight;
                    iHeight = pSaveImgData->iWidth;
                }
                std::string strRotatePath = pSaveImgData->strImgPath + std::to_string(pSaveImgData->iFrameId) + "_rotate.jpg";
                cv::imwrite(strRotatePath, cvRotateImg, compression_params);
            }

            Json::Value jvFrameInfo;
            jvFrameInfo["timeStamp"] = pSaveImgData->i64TimeStamp;
            jvFrameInfo["status"] = iStatus;
            jvFrameInfo["direction"] = pSaveImgData->iDirection;
            jvFrameInfo["width"] = iWidth;
            jvFrameInfo["height"] = iHeight;
            jvFrameInfo["rate"] = iRate;
            jvFrameInfo["isEnd"] = pSaveImgData->bIsEnd;
            MyUtils::getins()->WriteJsonInfo(jvFrameInfo, strTxtFilePath);
//            LogDebug << "engineId:" << engineId_ << " save success txt:" << strTxtFilePath;
        }
    }
    return APP_ERR_OK;
}