Hu Chunming / vpt_ascend

/*
 * File: pedestrian_vehicle_trespass.cpp
 * Created Date: Tuesday February 22nd 2022
 * Author: yangzilong (yangzilong@objecteye.com)
 * Description:
 * -----
 * Last Modified: Tuesday, 22nd February 2022 4:38:48 pm
 * Modified By: yangzilong (yangzilong@objecteye.com>)
 * -----
 * Copyright 2022
 */

#include "./pedestrian_vehicle_trespass.h"
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/opencv.hpp"
#include <cmath>
#include "../decoder/interface/DeviceMemory.hpp"
#include "../common/logger.hpp"
#include "../ai_platform/mvpt_process_assist.h"


namespace ai_engine_module {
namespace pedestrian_vehicle_trespass {

#define MAX_TRACE_NUM 5
#define MINIMUM_DISTANCE 10

static std::set<algorithm_type_t> algor_type_list_ = {
    algorithm_type_t::PEDESTRIAN_TRESPASS,
    algorithm_type_t::VEHICLE_TRESPASS,
};

// void show_algorthim_result(sy_img img, box_t cur_box, const obj_key_t &key);
// ############################################################ //
// !                Auxiliary Function                        ! //
// ############################################################ //

std::set<det_class_label_t> algor_type_to_det_label_set(const algorithm_type_t &algor_type) {
  if (algorithm_type_t::PEDESTRIAN_TRESPASS == algor_type) {
    return {det_class_label_t::HUMAN};
  } else if (algorithm_type_t::VEHICLE_TRESPASS == 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 is_valid_box(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));
}

bool is_valid_box(const box_t &box, const algorithm_type_t &algor_type,
                  const task_param_manager::algo_param_type_t_ *params_ptr = nullptr) {
  if (!params_ptr)
    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;

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

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

  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 is_valid_box(box.cls, algor_type);
}

/* 获取指定任务的算法配置参数 */
const task_param_manager::algo_param_type_t_ *get_task_param_ptr(const task_id_t &task_id,
                                                                 const algorithm_type_t &algor_type,
                                                                 task_param_manager *const task_param_manager) {
  if (!task_param_manager)
    return nullptr;

  auto &&algor_param_wrap = task_param_manager->get_task_other_param(task_id, algor_type);
  if (!algor_param_wrap)
    LOG_ERROR("{} is nullptr when get algor param from task_param", task_id);
  return algor_param_wrap;
}

/* 按照目标任务 检索指定任务的算法配置列表 */
std::vector<algorithm_type_t> task_id_to_algorithm_type_seq(const task_id_t &task_id,
                                                            task_param_manager *const task_param) {
  std::vector<algorithm_type_t> seq;
  auto &&algor_map = task_param->get_task_other_param(task_id);
  for (auto iter = algor_map->begin(); iter != algor_map->end(); ++iter) {
    // LOG_TRACE("task id is {} algor type is {}", task_id, int(iter->first));
    if (algor_type_list_.count(iter->first) > 0)
      seq.emplace_back(iter->first);
  }
  return seq; // N(RVO)
}

// ############################################################ //
// !                    Class Member                          ! //
// ############################################################ //

PedestrianVehicleTrespass::PedestrianVehicleTrespass() : task_param_manager_(nullptr) {
  if (!task_param_manager_)
    task_param_manager_ = task_param_manager::getInstance();
}

PedestrianVehicleTrespass::~PedestrianVehicleTrespass() = default;

/* 根据目标id 获取该目标的算法分析结果 */
vector<result_data_t> PedestrianVehicleTrespass::get_results_by_id(const obj_key_t &id, bool do_erase) {

  vector<result_data_t> res;
  auto it = obj_to_alarm_boxes_.find(id);
  if (it == obj_to_alarm_boxes_.end()) {
    // printf("cant find %s\n", id.task_id.c_str());
    return res;
  }

  res = it->second;
  if (do_erase)
    obj_to_alarm_boxes_.erase(id);

  return res;
}

/* 目标在禁区内外的辅助判断函数 */
bool PedestrianVehicleTrespass::in_rect_analysis(const obj_key_t &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);

  obj_key_t tmp_id = {0, id.task_id, id.algor_type};
  cv::Mat &dst_mat = trespass_regions[tmp_id];
  // printf("id=%d center: %d %d ", id.obj_id, center_x, center_y);
  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]) {
    // printf(" true\n");
    return true; // 进入禁区
  } else {
    // printf(" false\n");
    return false; // 未进入禁区
  }
}

/* 根据用户输入的点 初始化禁区区域mask */
void PedestrianVehicleTrespass::pedestrianvehicletrespass_init_region(const string &task_id,
                                                                      const algorithm_type_t algor_type,
                                                                      const int width, const int height) {
  obj_key_t obj_key = {0, task_id, algor_type};

  auto &&algor_type_seq = task_id_to_algorithm_type_seq(task_id, task_param_manager_);

  const task_param_manager::algo_param_type_t_ *param_ptr_wrap =
      get_task_param_ptr(task_id, algor_type, task_param_manager_);
  if (!param_ptr_wrap) {
    printf("cant find %s algorthim params\n", task_id.c_str());
    return;
  }

  using param_type_t = algor_config_param_trespass_basic;
  const param_type_t *param_ptr = (param_type_t *)(param_ptr_wrap->algor_param);

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

  std::vector<cv::Point> contour;
  contour.reserve(param_ptr->points_count);
  // printf("region points count: %d\n", param_ptr->points_count);
  for (int idx = 0; idx < param_ptr->points_count; idx++) {
    contour.emplace_back(param_ptr->points[idx].x_, param_ptr->points[idx].y_);
    // printf(" =point:%d x:%d y:%d\n", idx, param_ptr->points[idx].x_, param_ptr->points[idx].y_);
  }

  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 = trespass_regions[obj_key];
  cv::cvtColor(src, trespass_regions[obj_key], cv::COLOR_BGR2GRAY);
  cv::threshold(trespass_regions[obj_key], trespass_regions[obj_key], 100, 255, cv::THRESH_BINARY);
}

/* 非法闯入禁区的 算法判断函数 */
bool PedestrianVehicleTrespass::update_mstreams(const std::vector<task_id_t> &tasks_id, vector<DeviceMemory*> det_input_images,
                                                const std::vector<onelevel_det_result> &det_results,
                                                const vector<vector<int>> &delete_objs) {
  //! check.
  if (tasks_id.empty() || tasks_id.size() != det_results.size())
    return false;

  vector<obj_key_t> alarm_objs;
  //! loop.
  unsigned stream_idx = 0U;
  std::map<task_id_t, unsigned> task_id_to_stream_idx;
  for (auto task_id_iter = tasks_id.begin(); task_id_iter != tasks_id.end(); ++task_id_iter, ++stream_idx) {
    const auto &task_id = *task_id_iter;
    task_id_to_stream_idx[task_id] = stream_idx;
    // printf("\nbegin judge: %s\n", task_id.c_str());

    /* 判断该路任务 是否开启该算法 */
    auto &&algor_type_seq = task_id_to_algorithm_type_seq(task_id, task_param_manager_);
    if (algor_type_seq.empty())
      continue;

    // 删除 已经删除的目标
    for (auto obj_idx : delete_objs[stream_idx]) {
      for (auto algor_type : algor_type_seq) {
        obj_key_t obj_key{obj_idx, task_id, algor_type};

        if (obj_to_position_.find(obj_key) != obj_to_position_.end()) {
          // printf("delete obj: %s %d\n", task_id.c_str(), obj_idx);
          obj_to_position_.erase(obj_key);
        }
      }

      // delete obj
    }

    /* 依次判断检测目标框 是否有非法闯入 判断逻辑：之前帧在禁区外 当前帧进入禁区 */
    auto &det_result = det_results[stream_idx];
    // printf("det count: %d\n", det_result.obj_count);
    for (unsigned box_idx = 0; box_idx < det_result.obj_count; ++box_idx) {
      auto &box = det_result.obj[box_idx];

      box_t unique_box{};
      {
        unique_box.id = box.id;
        unique_box.cls = box.index;

        unique_box.top = box.top;
        unique_box.left = box.left;
        unique_box.right = box.right;
        unique_box.bottom = box.bottom;
        unique_box.score = box.confidence;
      }

      //! loop algor
      for (auto algor_type : algor_type_seq) {
        obj_key_t obj_key{box.id, task_id, algor_type};

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

        //! add or update.
        if (!in_rect_analysis(obj_key, unique_box)) // 禁区外
        {
          // printf("push in: %s %d\n", task_id.c_str(), box.id);
          obj_to_position_[obj_key] = unique_box;
        }

        else {                                                          // 进入禁区
          if (obj_to_position_.find(obj_key) != obj_to_position_.end()) // 之前在禁区外，可报警
          {
            obj_to_position_[obj_key] = unique_box;

            alarm_objs.emplace_back(obj_key);
          }

        }
      }
    }
  }

  /* 针对需要报警的目标 缓存当前大图&抠图 等待报警返回 */
  VPCUtil* pVpcUtil = VPCUtil::getInstance();
  for (const auto &obj_key : alarm_objs) {
    auto &&it = task_id_to_stream_idx.find(obj_key.task_id);
    if (it == task_id_to_stream_idx.end())
      continue;

    // 221009 byzsh记录10条即可--------------------------------------------------------------------------------------
    if (obj_to_alarm_boxes_.find(obj_key) != obj_to_alarm_boxes_.end() && obj_to_alarm_boxes_[obj_key].size() >= 10)
      continue;
    //--------------------------------------------------------------------------------------------------------------
    const unsigned stream_idx = it->second;
    auto &src_img = det_input_images[stream_idx];

    result_data_t result_data;
    result_data.box = obj_to_position_[obj_key]; 

    // 原图
    vpc_img_info src_img_info = VPCUtil::vpc_devMem2vpcImg(src_img);
    result_data.origin_img_desc = src_img_info.pic_desc;

    // 抠图
    int width = src_img->getWidth();
    int height = src_img->getHeight();
    video_object_info obj;
    strcpy(obj.task_id, obj_key.task_id.c_str());
    obj.object_id = obj_key.obj_id;
    obj.left = clip(result_data.box.left, 0, width);
    obj.top = clip(result_data.box.top, 0, height);
    obj.right = clip(result_data.box.right, 0, width);
    obj.bottom = clip(result_data.box.bottom, 0, height);

    vpc_img_info img_info = pVpcUtil->crop(src_img, obj);
    if(img_info.pic_desc != nullptr){
      result_data.roi_img_desc = img_info.pic_desc;
      obj_to_alarm_boxes_[obj_key].emplace_back(std::move(result_data));
    }
    
    obj_to_position_.erase(obj_key);
  }

  return true;
}

/* 辅助函数 显示算法结果 */
// void show_algorthim_result(sy_img img, box_t cur_box, const obj_key_t &key) {
//   const unsigned size = img.c_ * img.h_ * img.w_;
//   auto *img_data = new unsigned char[size];

//   CHECK(cudaMemcpy(img_data, img.data_, size * sizeof(unsigned char), cudaMemcpyDeviceToHost));
//   cv::Mat show_image(img.h_, img.w_, CV_8UC3, img_data);

//   std::vector<cv::Point> contour;
//   /*
//       contour.push_back(cv::Point(200, 200));
//       contour.push_back(cv::Point(600, 200));
//       contour.push_back(cv::Point(600, 500));
//       contour.push_back(cv::Point(200, 500));
//   */
//   contour.emplace_back(500, 500);
//   contour.emplace_back(1500, 500);
//   contour.emplace_back(1500, 900);
//   contour.emplace_back(500, 900);

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

//   cv::polylines(show_image, contours, true, cv::Scalar(255, 255, 255), 2,
//                 8); // 第2个参数可以采用contour或者contours，均可

//   cv::rectangle(show_image, cv::Point(cur_box.left, cur_box.top), cv::Point(cur_box.right, cur_box.bottom),
//                 cv::Scalar(0, 250, 0), 2, 8);

//   char filename[256];
//   sprintf(filename, "res_image/%s_%ld.jpg", key.task_id.c_str(), key.obj_id);
//   cv::imwrite(filename, show_image);
//   printf("\n *****ERROR finish save %s %ld\n", key.task_id.c_str(), key.obj_id);
// }

} // namespace pedestrian_vehicle_trespass
} // namespace ai_engine_module