gsd_var/plugins/gsd_InferPlugin.cc

/**
 *
 *  gsd_InferPlugin.cc
 *
 */

#include "gsd_InferPlugin.h"


using namespace drogon;
using namespace gsd;

#define CLIP(x) x < 0 ? 0 : (x > 1 ? 1 : x)

/**
 * @description: 初始化
 * @param {Value} &config
 * @return {*}
 */
void InferPlugin::initAndStart(const Json::Value &config)
{
    /// Initialize and start the plugin
    InfoL; 
    if(m_InferTools == nullptr) m_InferTools = InferTools::getPtr();
    auto *plugin = app().getPlugin<CenterGroup>();
    if(m_InferTools->Init(plugin->getInferEngine()) != false) {
        InfoL << "初始化成功" << endl;
    }else{
        WarnL << "初始化失败" << endl;
    }
    if(this->pool == nullptr) this->pool = std::make_shared<ThreadPool>(5,ThreadPool::PRIORITY_HIGHEST, false);
}

/**
 * @description: 关闭
 * @return {*}
 */
void InferPlugin::shutdown() 
{
    /// Shutdown the plugin
    InfoL;
    this->Destroy();
}

/**
 * @description: 启动任务
 * @return {*}
 */
void InferPlugin::StartTask(){
    InfoL;
}

/**
 * @description: 创建任务
 * @param {HttpRequestPtr&} req
 * @return {*}
 */
std::future<std::pair<bool, std::string>> InferPlugin::CreateTask(const drogon::HttpRequestPtr& req, InferPluginTaskType type){
    std::shared_ptr<std::packaged_task<std::pair<bool, std::string>()>> task = std::make_shared<std::packaged_task<std::pair<bool, std::string>()>>([&, req, type]{
        std::string images_base64 = req->getParameter("image");
        std::string datas = req->getParameter("datas");
        std::pair<bool, std::string> result(false, "");
        if(images_base64 == "" || datas == "") return result;
        std::string decode_str = base64_decode(images_base64);
        vector<uint8_t> img_data(decode_str.begin(), decode_str.end());
        std::shared_ptr<cv::Mat> img = std::make_shared<cv::Mat>();
        *img = cv::imdecode(cv::Mat(img_data), cv::IMREAD_COLOR);
        CNStreamInferData::Ptr data = std::make_shared<CNStreamInferData>();
        
        try
        {
            if(type == InferPluginTaskType::CHECKT_TASK){
                if(data->jsonToObject(datas)) {
                    return this->CheckTask(img, data);
                }
                else return this->CheckTask(img);
            }else if(type == InferPluginTaskType::CLASSIFY_TASK){
                if(data->jsonToObject(datas)) {
                    return this->CheckTask(img, data);
                }
                else return this->CheckTask(img);
            }
        }
        catch(const std::exception& e)
        {
            ErrorL << e.what() << endl;
        }
    });
    this->pool->async([task](){
        if(task->valid()){
            (*task)();
        }
    });
    this->pool->start();
    return task->get_future();
}

/**
 * @description: 推理任务
 * @param {Mat&} img
 * @param {CNStreamInferData&} data
 * @return {*}
 */
std::pair<bool, std::string> InferPlugin::CheckTask(std::shared_ptr<cv::Mat> img, CNStreamInferData::Ptr data){
    CNStreamInferData::Ptr result = std::make_shared<CNStreamInferData>();
    this->m_InferTools->Inference(img, result);
    return this->getInfoResult(img, data, result);
}

/**
 * @description: 获取结果
 * @param {CNStreamInferData&} data1
 * @param {CNStreamInferData&} data2
 * @return {*}
 */
std::pair<bool, std::string> InferPlugin::getInfoResult(std::shared_ptr<cv::Mat> img, CNStreamInferData::Ptr data1, CNStreamInferData::Ptr data2){
    static std::shared_ptr<ThreadPool> save_pool = nullptr;
    std::string json;
    data2->objectToJson(json);
    bool result;
    if(data2->Objects.size()) result = true;
    else result = false;
    std::pair<bool, std::string> dates(result, json);
    auto *plugin = app().getPlugin<CenterGroup>();
    // 判断是否保存结果
    if(plugin->getInferSave()){
        if(save_pool == nullptr) save_pool = std::make_shared<ThreadPool>(1,ThreadPool::PRIORITY_HIGHEST, false);
        save_pool->async([img, data1, data2](){
            auto *plugin = app().getPlugin<CenterGroup>();
            std::string dir = plugin->getSavePath();
            DIR *dirp =  opendir(dir.c_str());
            if(dirp == NULL){
                mkdir(dir.c_str(),S_IRWXU | S_IRWXG | S_IRWXO);
            }else{
                closedir(dirp);
            }
            std::string video = dir + utility::uuid::generate() + ".jpg";
            cv::Mat result;
            img->copyTo(result);
            if(data1 != nullptr){
                // MLU220
                for(auto iter : data1->Objects){
                    std::pair<cv::Point, cv::Point> corner = InferPlugin::GetBboxCorner(iter, img->cols, img->rows);
                    cv::Point top_left = corner.first;
                    cv::Point bottom_right = corner.second;
                    cv::Point bottom_left(top_left.x, bottom_right.y);
                    InferPlugin::DrawBox(*img, top_left, bottom_right, cv::Scalar(0,255,255));
                    std::string text = iter.Label + " " + std::to_string(iter.Score);
                    top_left.y = top_left.y - 2;
                    cv::putText(*img, text, top_left, cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
                }
                cv::putText(*img, "MLU220 " + data1->videoPath, cv::Point(0, 15), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
            }
            if(data2 != nullptr){
                // Jetson NX
                for(auto iter : data2->Objects){
                    std::pair<cv::Point, cv::Point> corner = InferPlugin::GetBboxCorner(iter, img->cols, img->rows);
                    cv::Point top_left = corner.first;
                    cv::Point bottom_right = corner.second;
                    cv::Point bottom_left(top_left.x, bottom_right.y);
                    InferPlugin::DrawBox(result, top_left, bottom_right, cv::Scalar(255,255,0));
                    std::string text = iter.Label + " " + std::to_string(iter.Score);
                    top_left.y = top_left.y - 2; 
                    cv::putText(result, text, top_left, cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
                }
                cv::putText(result, "Jetson NX", cv::Point(0, 15), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 4);
            }
            cv::hconcat(*img, result, *img);
            //cv::resize(*img, *img, cv::Size(960, 540));
            cv::imwrite(video, *img);
        });
        save_pool->start();
    }
    return dates;
}

/**
 * @description: Destroy
 * @return {*}
 */
void InferPlugin::Destroy(){
    m_InferTools->Destroy();
}

/**
 * @description: GetBboxCorner
 * @param {*}
 * @return {*}
 */
std::pair<cv::Point, cv::Point> InferPlugin::GetBboxCorner(const InferInfo object, int img_width, int img_height) {
    float x = CLIP(object.BBox.x);
    float y = CLIP(object.BBox.y);
    float w = CLIP(object.BBox.w);
    float h = CLIP(object.BBox.h);
    w = (x + w > 1) ? (1 - x) : w;
    h = (y + h > 1) ? (1 - y) : h;
    cv::Point top_left(x * img_width, y * img_height);
    cv::Point bottom_right((x + w) * img_width, (y + h) * img_height);
    return std::make_pair(top_left, bottom_right);
}

/**
 * @description: CalcThickness
 * @param {int} image_width
 * @param {float} thickness
 * @return {*}
 */
int InferPlugin::CalcThickness(int image_width, float thickness) {
  int result = thickness * image_width / 300;
  if (result <= 0) result = 1;
  return result;
}

/**
 * @description: DrawBox
 * @param {*}
 * @return {*}
 */
void InferPlugin::DrawBox(cv::Mat& image, const cv::Point& top_left, const cv::Point& bottom_right,
                    const cv::Scalar color) {
  float w =bottom_right.x - top_left.x;
  float h =bottom_right.y - top_left.y;
  float box_thickness_ = 1;
  int wscale = w * .25f;
  int hscale = h * .25f;
  cv::Point top_right(bottom_right.x,top_left.y);
  cv::Point bottom_left(top_left.x,bottom_right.y);
  // -- 
  cv::line(image,top_left,cv::Point(top_left.x + wscale,top_left.y),color, CalcThickness(image.cols, box_thickness_));
  // |
  cv::line(image,top_left,cv::Point(top_left.x ,top_left.y + hscale),color, CalcThickness(image.cols, box_thickness_));
  //      --
  cv::line(image,cv::Point(top_right.x - wscale,top_right.y),top_right,color, CalcThickness(image.cols, box_thickness_));
  //        |
  cv::line(image,top_right,cv::Point(top_right.x,top_right.y + hscale),color, CalcThickness(image.cols, box_thickness_));
  //
  // __
  cv::line(image,bottom_left,cv::Point(bottom_left.x+wscale,bottom_left.y),color, CalcThickness(image.cols, box_thickness_));
  //
  //|
  cv::line(image,bottom_left,cv::Point(bottom_left.x,bottom_left.y - hscale),color, CalcThickness(image.cols, box_thickness_));
  //
  //        |
  cv::line(image,bottom_right,cv::Point(bottom_right.x,bottom_right.y - hscale),color, CalcThickness(image.cols, box_thickness_));
  
  cv::line(image,bottom_right,cv::Point(bottom_right.x - wscale,bottom_right.y),color, CalcThickness(image.cols, box_thickness_));
  // cv::rectangle(image, top_left, bottom_right, color, CalcThickness(image.cols, box_thickness_));
}