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 "../ai_platform/mvpt_process_assist.h"



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

/* 根据用户输入的点 初始化禁区区域mask */
bool RegionLeave::add_task(const string &task_id, const int width, const int height) {

  auto &&params_ptr = m_task_param_manager->get_task_other_param(task_id, m_eType);
  if (!params_ptr)
  {
    LOG_ERROR("{} is nullptr when get algor param from task_param", task_id);
    return false;
  }

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

  auto *algor_params_ptr = (algor_config_param_trespass_basic *)(params_ptr->algor_param);
  if (algor_params_ptr == nullptr)
  {
    return false;
  }

  cv::Mat src(height, width, CV_8UC3);
  src.setTo(0);

  std::vector<cv::Point> contour;
  for (int idx = 0; idx < algor_params_ptr->points_count; idx++) {
    contour.emplace_back(algor_params_ptr->points[idx].x_, algor_params_ptr->points[idx].y_);
  }

  if(contour.size() <= 0){
    return false;
  }

  std::vector<std::vector<cv::Point>> contours;
  contours.push_back(contour);

  cv::polylines(src, contours, true, cv::Scalar(255, 255, 255), 2, 8); // 第2个参数可以采用contour或者contours，均可
  cv::fillPoly(src, contours, cv::Scalar(255, 255, 255)); // fillPoly函数的第二个参数是二维数组

  cv::Mat &dst_mat = region_map[task_id];
  cv::cvtColor(src, region_map[task_id], cv::COLOR_BGR2GRAY);
  cv::threshold(region_map[task_id], region_map[task_id], 100, 255, cv::THRESH_BINARY);

  return true;
}

/* 目标在禁区内外的辅助判断函数 */
bool RegionLeave::in_rect_analysis(string task_id, const box_t &cur_bos) {
  int center_x = int((cur_bos.left + cur_bos.right) / 2.0);
  int center_y = int((cur_bos.top + cur_bos.bottom) / 2.0);

  cv::Mat &dst_mat = region_map[task_id];
  if (dst_mat.data[center_y * dst_mat.cols + center_x] && dst_mat.data[center_y * dst_mat.cols + center_x + 1] &&
      dst_mat.data[center_y * dst_mat.cols + center_x - 1] && dst_mat.data[(center_y + 1) * dst_mat.cols + center_x] &&
      dst_mat.data[(center_y - 1) * dst_mat.cols + center_x]) {
    return true; // 进入禁区
  } else {
    return false; // 未进入禁区
  }
}

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 (!snapshot_legal_inarea(params_ptr->basic_param->algor_valid_rect, box.left, box.top, box.right, box.bottom))
    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 it = region_map.find(task_id);
    if (it == region_map.end()) {
      continue;
    }

    auto &&params_ptr = m_task_param_manager->get_task_other_param(task_id, m_eType);
    if (!params_ptr && !params_ptr->basic_param)
    {
      LOG_ERROR("{} is nullptr when get algor param from task_param", task_id);
      continue;
    }

    bool b_src_saved = false;

    string sep = "/";

    auto& basic_param = params_ptr->basic_param;
    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); // 如果不满足条件 非 合法框 依然删除
      }

      //! add or update.
      if (in_rect_analysis(task_id, unique_box)) {
        // 禁区内
        // 新加
        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);
  }
}

void RegionLeave::release_task(string task_id){
  auto it = region_map.find(task_id);
  if (it != region_map.end())
  {
    region_map.erase(it);
  }
}