#include "HumanParsing.h"

#include <opencv2/opencv.hpp>
static string head[5] = { "长发", "短发", "光头", "帽子", "其他" };
static string head_color[9] = { "黑", "白", "灰", "红", "蓝", "黄", "绿", "多色", "其他" };
static string eye[2] = { "未戴眼镜", "戴眼镜" };
static string mouth[2] = { "未戴口罩", "戴口罩" };
//string up[9] = { "T恤/背心", "衬衫",  "毛衣", "羽绒服", "大衣/风衣", "外套/夹克/西服", "连衣裙", "无上衣", "其他" };
static string up[9] = { "T恤", "衬衫", "毛衣", "羽绒服", "大衣", "外套", "连衣裙", "无上衣", "其他" };
static string up_color[9] = { "黑", "白", "灰", "红", "蓝", "黄", "绿", "多色", "其他" };
static string clothing_text[5] = { "纯色", "碎花", "条纹", "格子", "其他" };
static string down[6] = { "长裤", "短裤", "长裙", "短裙", "连衣裙", "其他" };
static string down_color[9] = { "黑", "白", "灰", "红", "蓝", "黄", "绿", "多色", "其他" };
static string bao[5] = { "无包", "单肩包", "背包/双肩包", "手提包/手提物", "拉杆箱/小推车" };
static string sex[3] = { "男", "女", "其他" };
static string age[5] = { "小孩", "青年", "中年", "老年", "其他" };
static string viewpoint[3] = { "正面", "背面", "侧面" };
static string dasan[2] = { "未打伞", "打伞" };
static string child[2] = { "未抱小孩", "抱小孩" };
static string personstate[5] = { "行走", "奔跑", "蹲坐", "推车", "其他" };


const int FIR_INDEX_SIZE = 16;
const int SEC_INDEX_SIZE[FIR_INDEX_SIZE] = { 5, 9, 2, 2, 9, 9, 5, 6, 9, 5, 3, 5, 3, 2, 2, 5 };



int HumanParsing_Init(void *&handle, int gpuid)
{


	hp_param param;
    param.mode = DEVICE_GPU;
    param.gpuid = gpuid;
	param.engine = ENGINE_TENSORRT;
	param.max_batch = 20;
	param.serialize_file = "./serialize_file/HP";
	if (hp_init(&handle, param) != 0)
	{
		cout << "Init HP Failed!" << endl;
		return FAILED;
	}
 
	return SUCCESS;

}




#include <fstream>
int HumanParsing_Process(void * handle,  sy_img * batch_img, int batch_size, hp_analysis_res*& result)
{

 
	for (int i = 0; i < batch_size; i++)
	{
		if (batch_img[i].data_ == NULL)
			cout << i << " data null" << endl;
	}

	//cout << "begin HumanParsingProcessBatch" << endl;
	hp_batch(handle, batch_img, batch_size, result);

	for (int b = 0; b < batch_size; b++)
	{
		hp_analysis_res &cur_result = result[b];

		for (int i = 0; i < FIR_INDEX_SIZE; i++)
		{
			if (cur_result.res_objs[i].res_index >= SEC_INDEX_SIZE[i])
			{
				cur_result.res_objs[i].res_index = 0;
				
			}
		}
	}

	//cout << "end HumanParsingProcessBatch" << endl;


    //for (int b = 0; b < batch_size; b++)
    //{
    //    ctools_result &cur_result = result[b];
    //    int big_class_count = cur_result.obj_count_;
    //    ctools_obj_result &index_score = cur_result.obj_results_[big_class_count];

    //    if (index_score.data_count_ != FIR_INDEX_SIZE * 2)
    //    {
    //        printf("wrong result.\n");
    //        break;
    //    }

    //    //for (int i = 0; i < FIR_INDEX_SIZE; i++)
    //    //{
    //    //    //int resIndex = result[j][i].index;
    //    //    int resIndex = index_score.data_[i * 2];

    //    //    switch (i)
    //    //    {
    //    //        case 0:
    //    //            cout << head[resIndex] << " ";
    //    //            break;
    //    //        case 1:
    //    //            cout << head_color[resIndex] << " ";
    //    //            break;
    //    //        case 2:
    //    //            cout << eye[resIndex] << " ";
    //    //            break;
    //    //        case 3:
    //    //            cout << mouth[resIndex] << " ";
    //    //            break;
    //    //        case 4:
    //    //            cout << up[resIndex] << " ";
    //    //            break;
    //    //        case 5:
    //    //            cout << up_color[resIndex] << " ";
    //    //            break;
    //    //        case 6:
    //    //            cout << clothing_text[resIndex] << " ";
    //    //            break;
    //    //        case 7:
    //    //            cout << down[resIndex] << " ";
    //    //            break;
    //    //        case 8:
    //    //            cout << down_color[resIndex] << " ";
    //    //            break;
    //    //        case 9:
    //    //            cout << bao[resIndex] << " ";
    //    //            break;
    //    //        case 10:
    //    //            cout << sex[resIndex] << " ";
    //    //            break;
    //    //        case 11:
    //    //            cout << age[resIndex] << " ";
    //    //            break;
    //    //        case 12:
    //    //            cout << viewpoint[resIndex] << " ";
    //    //            break;
    //    //        case 13:
    //    //            cout << dasan[resIndex] << " ";
    //    //            break;
    //    //        case 14:
    //    //            cout << child[resIndex] << " ";
    //    //            break;
    //    //        case 15:
    //    //            cout << personstate[resIndex] << " ";
    //    //            break;
    //    //        default:
    //    //            break;
    //    //    }

    //    //}
    //    //cout << endl << endl;
    //}

	return SUCCESS;
}

int HumanParsing_Release(void *& handle)
{
	if(handle)
	hp_release(&handle);
	return SUCCESS;
}