Hu Chunming / vpt_ascend_arm

#include "./RegionLeave.h"
#include "opencv2/highgui.hpp"
#include "opencv2/imgproc.hpp"
#include "opencv2/opencv.hpp"
#include <cmath>

#include "../helpers/gen_json.hpp"
#include "../helpers/img_util.h"


namespace ai_engine_module {

std::set<det_class_label_t> algor_type_to_det_label_set(const algorithm_type_t &algor_type) {
  if (algorithm_type_t::HUMAN_LEAVE_REGION == algor_type || algorithm_type_t::HUMAN_REGION_DISMISS == algor_type
   || algorithm_type_t::HUMAN_REGION_FAST_MOVING == algor_type) {
    return {det_class_label_t::HUMAN};
  } else if (algorithm_type_t::VEHICLE_LEAVE_REGION == algor_type) {
    return {
        det_class_label_t::LARGE_CAR, det_class_label_t::MEDIUM_BUS, det_class_label_t::SMALL_CAR,
        det_class_label_t::TRUCK,     det_class_label_t::TRACTOR,
    };
  } else {
    return {};
  }
}

/* 是否是有效目标框的辅助判断函数 */
bool check_obj_cls(const int &cls, const algorithm_type_t &algor_type) {
  return algor_type_to_det_label_set(algor_type).count(static_cast<det_class_label_t>(cls));
}

RegionLeave::RegionLeave(){
  m_task_param_manager = task_param_manager::getInstance();
}

RegionLeave::~RegionLeave()
{
  if (m_save_util)
  {
    delete m_save_util;
    m_save_util = nullptr;
  }
  
}
void RegionLeave::init(int devId, algorithm_type_t eType){

  m_devId = devId;
  m_eType = eType;

  m_save_util = new save_snapshot_reprocessing(m_devId);
}

#ifdef POST_USE_RABBITMQ
void RegionLeave::set_callback(callback_t cb) {
  m_save_util->set_callback(cb);
}
#endif

bool RegionLeave::is_valid_box(string task_id, const box_t &box, const algorithm_type_t &algor_type) {
  auto &&params_ptr = m_task_param_manager->get_task_other_param(task_id, algor_type);
  if (!params_ptr)
  {
    LOG_ERROR("{} is nullptr when get algor param from task_param", task_id);
    return false;
  }

  if (params_ptr->basic_param == nullptr)
    return false;

  if (params_ptr->algor_param == nullptr)
    return false;

  auto *algor_params_ptr = (algor_config_param_trespass_basic *)(params_ptr->algor_param);

  if (box.width() == 0 || box.height() == 0)
    return false;

  if (box.score < algor_params_ptr->conf_threshold || box.width() < algor_params_ptr->minmum_width ||
      box.height() < algor_params_ptr->minmum_height)
    return false;

  return check_obj_cls(box.cls, algor_type);
}

/* 非法闯入禁区的 算法判断函数 */
void RegionLeave::process(std::vector<DeviceMemory*>& vec_gpuMem, const std::vector<onelevel_det_result> &det_results, const vector<vector<int>> &delete_objs) 
{
  if (det_results.size() <= 0) {
    return ;
  }

  for (int i = 0; i < det_results.size(); i++){
    DeviceMemory* gpuMem = vec_gpuMem[i];
    string task_id = gpuMem->getId();

    // 删除 已经删除的目标
    for (auto obj_idx : delete_objs[i]) {
        OBJ_KEY obj_key{task_id, obj_idx};

        if (obj_to_position_.find(obj_key) != obj_to_position_.end()) {
          obj_to_position_.erase(obj_key);
        }
    }

    auto &&params_ptr = m_task_param_manager->get_task_other_param(task_id, m_eType);
    if (!params_ptr || !params_ptr->basic_param)
    {
      continue;
    }

    auto& basic_param = params_ptr->basic_param;
    if (basic_param == nullptr || basic_param->adapt_param == nullptr) {
      continue;
    }
    universal_algor_adapt_param *adapt_param = basic_param->adapt_param;

    string sep = "/";

    std::string cur_src_ts = std::to_string(helpers::timer::get_timestamp<std::chrono::milliseconds>());
    std::string origin_file_path = basic_param->result_folder + sep + task_id  + "_origin_" + cur_src_ts + ".jpg";

    /* 依次判断检测目标框 是否有非法闯入 判断逻辑：之前帧在禁区外 当前帧进入禁区 */
    auto& one_result = det_results[i];
    std::vector<box_t> boxes;
    for (unsigned c = 0; c < one_result.obj_count; ++c) {

      auto obj_c = one_result.obj[c];

      box_t unique_box;
      unique_box.id = obj_c.id;
      unique_box.cls = obj_c.index;
      unique_box.top = obj_c.top;
      unique_box.left = obj_c.left;
      unique_box.right = obj_c.right;
      unique_box.bottom = obj_c.bottom;
      unique_box.score = obj_c.confidence;

      OBJ_KEY obj_key{task_id, obj_c.id};

      if (!is_valid_box(task_id, unique_box, m_eType)){
        obj_to_position_.erase(obj_key); // 如果不满足条件 非 合法框 依然删除
      }

      sy_point center;
      center.x_ = (obj_c.right + obj_c.left)/ 2 ;
      center.y_ = (obj_c.bottom + obj_c.top) / 2;
      if (common::isInPolygon(adapt_param->points, adapt_param->points_count, center)) {
        // 禁区内
        // 新加
        obj_to_position_[obj_key] = unique_box;
      } else {                      
        // 更新                                    
        if (obj_to_position_.find(obj_key) != obj_to_position_.end()) // 之前在禁区内，可报警
        {
          box_t box;
          box.left = obj_c.left;
          box.right = obj_c.right;
          box.top = obj_c.top;
          box.bottom = obj_c.bottom;
          box.score = obj_c.confidence;
          boxes.push_back(box);

          // 报完警清除记录
          obj_to_position_.erase(obj_key);
        }
      }
    }

    if (boxes.size() <= 0)
    {
      continue;
    }

    int algorithm_type = (int)m_eType;
    string json_str = helpers::gen_json::gen_boxes_json(task_id, algorithm_type, boxes, origin_file_path);

    ImgSaveInfo info_origin;
    info_origin.img_info = VPCUtil::vpc_devMem2vpcImg(gpuMem);
    info_origin.file_path = origin_file_path;
    info_origin.json_str = json_str;
    m_save_util->reprocessing_process_wo_locus_async(info_origin);
  }
}

}  // namespace ai_engine_module