Hu Chunming / src

#include "face_det_ai_engine.h"
#include <cuda.h>
#include <cuda_runtime.h>
#include "cuda_kernels.h"
#include "mvpt_process_assist.h"

#include "../helpers/logger.hpp"
#include "../helpers/cuda_helper.h"

face_det_ai_engine::~face_det_ai_engine()
{
	fd_release(&handle);
	clear();
}

int face_det_ai_engine::init_ai_engine(const facedet_ai_engine_param &ai_param/*, person_det_algorthim_cache * cache*/)
{
	max_batchsize = ai_param.sdk_param.max_batch_size_detect;
	frame_count = 0;
	skips = 6;
	is_open = true;

	task_param_manager_ = task_param_manager::getInstance();

	return fd_init(&handle, ai_param.sdk_param);
}

int face_det_ai_engine::ai_engine_process_batch(sy_img *image_data_array, const int &image_size, std::set<std::string> &task_ids,
	std::vector<onelevel_det_result> &result, std::vector<std::vector<int>> &deleteObjectID)
{
	map<string, map<algo_type, task_param_manager::algo_param_type_t_*>> && algor_param = task_param_manager_->get_task_other_params();

	const int total_batchsize = image_size;
	if (total_batchsize > 0 /*&& task_batchsize != 0*/)
	{
		fd_result *fd_result_ = new fd_result[total_batchsize];
		for (int i = 0; i < total_batchsize; ++i)
			fd_result_[i].info = new fd_info[50];

		// int stride = min(total_batchsize, max_batchsize);  // min(best_batchsize, max_batchsize));

		int stride = max_batchsize;
		int steps = (total_batchsize + stride - 1) / stride;


		for (int c = 0; c < steps; ++c)
		{
			int offset = c * max_batchsize;
			const int batchsize = (c == steps - 1) ? (total_batchsize - offset) : stride;
			LOG_TRACE("[face_detection]: image_size: {} model_batch_size: {}, step: [{}/{}] offset: {} batchsize: {}", image_size, max_batchsize, c, steps, offset, batchsize);
			auto flag = fd_detect_batch(handle, image_data_array + offset, SY_FORMAT_BGR888, batchsize, fd_result_ + offset);
		}

		// 属性检测使用人脸检测的原图，不需要切图
		int cur_index = 0;
		int img_index = 0;

		vector <vector< vector <float>>> detectResult(total_batchsize);  // sort
		auto task_id_iter = task_ids.cbegin();
		for (int c = 0; c < total_batchsize; ++c)
		{
			task_param_manager::algo_param_type_t_* cur_task_params = algor_param[*task_id_iter++][algorithm_type_t::FACE_SNAPSHOT];
			for (int i = 0; i < fd_result_[c].count; ++i)
			{
				if (!snapshot_legal_inarea(cur_task_params->basic_param->algor_valid_rect,
					fd_result_[c].info[i].face_position.left_, fd_result_[c].info[i].face_position.top_,
					fd_result_[c].info[i].face_position.left_ + fd_result_[c].info[i].face_position.width_,
					fd_result_[c].info[i].face_position.top_ + fd_result_[c].info[i].face_position.height_)
					|| fd_result_[c].info[i].face_pos_score < ((algor_config_param_snapshot *)cur_task_params->algor_param)->threshold)
				{
					continue;
				}

				vector <float> obj;
				obj.push_back(fd_result_[c].info[i].face_position.left_);
				obj.push_back(fd_result_[c].info[i].face_position.top_);
				obj.push_back(fd_result_[c].info[i].face_position.left_ + fd_result_[c].info[i].face_position.width_); //right
				obj.push_back(fd_result_[c].info[i].face_position.top_ + fd_result_[c].info[i].face_position.height_); //bottom
				obj.push_back(fd_result_[c].info[i].score);
				obj.push_back(1);		//统一index值为1
				//存入关键点信息
				for(int j = 0; j < FACIALFEAPOINTSIZE; ++j)
				{
					obj.push_back(fd_result_[c].info[i].facial_fea_point[j].x_);
					obj.push_back(fd_result_[c].info[i].facial_fea_point[j].y_);
				}

				//-added by zsh 添加姿态角信息------------------------------
				obj.push_back(fd_result_[c].info[i].roll);
				obj.push_back(fd_result_[c].info[i].yaw);
				obj.push_back(fd_result_[c].info[i].pitch);
				// cout << fabs(fd_result_[c].info[i].roll) << " " << fabs(fd_result_[c].info[i].yaw) << " " << fabs(fd_result_[c].info[i].pitch) << endl;
				//--------------------------------------------------------
				detectResult[c].push_back(obj);

#if 0
				if (fd_result_[img_index].count > 1)
				{
					//选择居中且靠上的人脸作为唯一的结果
					float min_dis = numeric_limits<float>::max();
					int min_index = 0;
					float person_center_x = (float)(cur_persondet_result[c]->obj[i].right - cur_persondet_result[c]->obj[i].left) / 2.0;
					float person_center_y = (float)(cur_persondet_result[c]->obj[i].bottom - cur_persondet_result[c]->obj[i].top) / 6.0;
					for (int j = 0; j < fd_result_[img_index].count; ++j)
					{
						float cx = (float)fd_result_[img_index].info[j].face_position.left_ + (float)(fd_result_[img_index].info[j].face_position.width_) / 2.0;
						float cy = (float)fd_result_[img_index].info[j].face_position.top_ + (float)(fd_result_[img_index].info[j].face_position.height_) / 2.0;
						float dis = (person_center_x - cx) * (person_center_x - cx) + (person_center_y - cy) * (person_center_y - cy);
						if (dis < min_dis)
						{
							min_dis = dis;
							min_index = j;
						}
					}
					//姿态角控制
					if (fabs(fd_result_[img_index].info[min_index].roll) < pose_thresld[c] && fabs(fd_result_[img_index].info[min_index].yaw) < pose_thresld[c] && fabs(fd_result_[img_index].info[min_index].pitch) < pose_thresld[c])
					{
						cur_res.count = 1;
						cur_res.info = new fd_info[1];
						memcpy(&cur_res.info[0], &fd_result_[img_index].info[min_index], sizeof(fd_info));
					}
					else
					{
						cur_res.info = new fd_info[1];
						cur_res.count = 0;
					}

				}
				else if (fd_result_[img_index].count == 1 && fabs(fd_result_[img_index].info[0].roll) < pose_thresld[c] && fabs(fd_result_[img_index].info[0].yaw) < pose_thresld[c] && fabs(fd_result_[img_index].info[0].pitch) < pose_thresld[c]) //姿态角控制
				{
					cur_res.count = 1;
					cur_res.info = new fd_info[1];
					memcpy(&cur_res.info[0], &fd_result_[img_index].info[0], sizeof(fd_info));
				}
				else
				{
					cur_res.info = new fd_info[1];
					cur_res.count = 0;
				}
				_fd_result[vec_ids[c]].push_back(cur_res);

				for (int j = 0; j < cur_res.count; ++j)
				{
					++cur_index;
				}
#endif

			}
		}

// added by zsh
#if 0
	for (int b = 0; b < total_batchsize; b++)
	{
		// printf("batch: %d, %d %d %d \n", b, image_data_array[b].c_ , image_data_array[b].h_ , image_data_array[b].w_);
		int data_size = image_data_array[b].c_ * image_data_array[b].h_ * image_data_array[b].w_;
		unsigned char *imgdata = new unsigned char[data_size];
		cudaMemcpy(imgdata, image_data_array[b].data_, sizeof(unsigned char) * data_size, cudaMemcpyDeviceToHost);
		cv::Mat big_img = cv::Mat(image_data_array[b].h_, image_data_array[b].w_, CV_8UC3, imgdata);

		for (int c = 0; c < detectResult[b].size(); c++)
		{
			// printf("%d %d %d %d\n", detectResult[b][c][0], detectResult[b][c][1], detectResult[b][c][2] - detectResult[b][c][0],
			// 	detectResult[b][c][3] - detectResult[b][c][1]);
			cv::rectangle(big_img, cv::Rect(detectResult[b][c][0], detectResult[b][c][1], detectResult[b][c][2] - detectResult[b][c][0],
				detectResult[b][c][3] - detectResult[b][c][1]), cv::Scalar(158, 52, 254), 3, 1, 0);
		}
		std::string file_path = "res/test/";
		auto time_now = std::chrono::system_clock::now();
		std::string cur_timestamp_us =  std::to_string(std::chrono::duration_cast<std::chrono::microseconds>(time_now.time_since_epoch()).count());
		std::string img_filename = file_path + cur_timestamp_us + "_" + std::to_string(detectResult[b].size()) + ".jpg";
		cv::imwrite(img_filename, big_img);
		delete[] imgdata;
	}
#endif

		//跟踪
		int real_index = 0;
		std::set<std::string>::iterator iter = task_ids.cbegin();
		while (iter != task_ids.cend())
		{
			if (!task_trackers[*iter].tracker.GetState())
				continue;
			bool isUseDet = true;
			vector<int> delete_ids;
			const float maxLen = std::sqrt(image_data_array[real_index].w_ * image_data_array[real_index].w_ + image_data_array[real_index].h_ * image_data_array[real_index].h_); //-modified by zsh 220719
			// printf("maxLen: %f\n", maxLen);
			for (int j = 0; j < task_trackers[*iter].fusion_interval; ++j)
			{
				if (j == 0)
				{
					// int objCount = task_trackers[*iter].tracker.update(isUseDet, /*save lk = */true, detectResult[real_index], result[real_index].obj, deleteObjectID[real_index]);
					//-modified by zsh 220719
					int objCount = task_trackers[*iter].tracker.update_v2(isUseDet, /*save lk = */true, /*center_dist = */true, maxLen, detectResult[real_index], result[real_index].obj, deleteObjectID[real_index]);
					result[real_index].obj_count = objCount;
					vector<vector<float>>().swap(detectResult[real_index]);
					detectResult[real_index].clear();
					isUseDet = false;
				}
				else
				{
					onelevel_det_result unresult;
					// unresult.obj_count = task_trackers[*iter].tracker.update(isUseDet, true, detectResult[real_index], unresult.obj, deleteObjectID[real_index]);
					//-modified by zsh 220719
					unresult.obj_count = task_trackers[*iter].tracker.update_v2(isUseDet, true, true, maxLen, detectResult[real_index], unresult.obj, deleteObjectID[real_index]);
				}

			}

			++real_index;
			++iter;
		}

// added by zsh
#if 0
	for (int b = 0; b < total_batchsize; b++)
	{
		int data_size = image_data_array[b].c_ * image_data_array[b].h_ * image_data_array[b].w_;
		unsigned char *imgdata = new unsigned char[data_size];
		cudaMemcpy(imgdata, image_data_array[b].data_, sizeof(unsigned char) * data_size, cudaMemcpyDeviceToHost);
		cv::Mat big_img = cv::Mat(image_data_array[b].h_, image_data_array[b].w_, CV_8UC3, imgdata);

		for (int c = 0; c < result[b].obj_count; c++)
		{
			cv::putText(big_img, std::to_string(result[b].obj[c].id), cv::Point(result[b].obj[c].left, result[b].obj[c].top-5), cv::FONT_HERSHEY_SIMPLEX, 0.75, cv::Scalar(255,0,0),2,8); 
			cv::rectangle(big_img, cv::Rect(result[b].obj[c].left, result[b].obj[c].top, result[b].obj[c].right - result[b].obj[c].left,
				result[b].obj[c].bottom - result[b].obj[c].top), cv::Scalar(158, 52, 254), 3, 1, 0);
		}
		std::string file_path = "res/track/";
		auto time_now = std::chrono::system_clock::now();
		std::string cur_timestamp_us =  std::to_string(std::chrono::duration_cast<std::chrono::microseconds>(time_now.time_since_epoch()).count());
		std::string img_filename = file_path + cur_timestamp_us + "_" + std::to_string(result[b].obj_count) + ".jpg";
		cv::imwrite(img_filename, big_img);
		delete[] imgdata;
	}
#endif

		vector <vector< vector <float>>>().swap(detectResult);  // free memory.

		if (fd_result_)
		{
			for (int i = 0; i < total_batchsize; ++i)
			{
				delete[] fd_result_[i].info;
				fd_result_[i].info = nullptr;
			}
			delete[] fd_result_;
			fd_result_ = nullptr;
		}
	}
	return 0;
}

void face_det_ai_engine::clear()
{
	for (auto it = _fd_result.begin(); it != _fd_result.end();)
	{
		for (auto &fd : it->second)
		{
			delete[] fd.info;
			fd.info = nullptr;
		}
		_fd_result.erase(it++);
	}
}

void face_det_ai_engine::operator_tracker(std::string task_id, const int &oper, int m_fusion_interval)
{
	//if (get_status())
	{
		switch (oper)
		{
		case ADDTASK:
		{
			LOG_INFO("face: tracker add task {}", task_id.c_str());
			task_tracker t;
			t.task_id = task_id;

			t.fusion_interval = m_fusion_interval;
			task_trackers.insert(std::make_pair(task_id, t));
			break;
		}
		case PAUSETASK:
		{
			auto iter = task_trackers.find(task_id);
			if (iter != task_trackers.end())
				iter->second.tracker.Pause();
			break;
		}
		case RESTARTTASK:
		{
			auto iter = task_trackers.find(task_id);
			if (iter != task_trackers.end())
				iter->second.tracker.ReSet();
			break;
		}
		case FINISHTASK:
		{
			LOG_INFO("face: tracker finish task {}", task_id.c_str());
			auto iter = task_trackers.find(task_id);
			if (iter != task_trackers.end())
			{
				task_trackers.erase(task_id);
			}
			break;
		}
		default:
			break;
		}
	}
}


void face_det_ai_engine::finish_task(std::string task_id)
{
	auto iter = id_frame_count.find(task_id);
	if (iter != id_frame_count.end())
	{
		id_frame_count.erase(task_id);
	}
	frame_count = 0;
}