/*********版权所有（C）2022，武汉高德红外股份有限公司***************
* 文件名称：YOLO.h
* 文件标识：yolo目标识别算法
* 内容摘要：yolo目标识别算法
* 其它说明：
* 当前版本：V1.0
* 创建作者：04046wcw
* 创建日期：20220408
* 
* 修改记录1：
*    修改日期：
*    版 本 号：
*    修 改 人：
*    修改内容： 
* 修改记录2：… 
*******************************************************************/
#ifndef YOLO_H
#define YOLO_H

#pragma once

#include "Arith_CommonDef.h"
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <vector>
#include <string>
#include <thread>
using std::string;


#ifdef _WIN32
#include <windows.h>
#else
#include <unistd.h>
#endif

#include <ctime>
#include <iostream>

void SLEEP(int ms);
using cv::dnn::Net;


struct Detection
{
	int class_id;
	float confidence;
	cv::Rect box;
};

struct YoloParam
{
	float score_threshold;
	float nms_threshold;
	float conf_threshold;
	bool  bUseCuda;
};



class YOLO_ONNX_Detector
{

public:
	YOLO_ONNX_Detector(){};
	virtual ~YOLO_ONNX_Detector(){};

	virtual int detect(cv::Mat& srcimg) = 0;

	virtual void SetParam(YoloParam para){m_para = para;};

	virtual std::vector<Detection>& getTargetVec(){return m_targetArray;};

	virtual std::string getClassName(int ID){return g_YoloClassName[ID];};
	
protected:
	cv::dnn::Net m_YoloNet;
	std::vector<std::string> g_YoloClassName;
	std::vector<Detection> m_targetArray;
	YoloParam m_para;
	std::string NetPath;
};


class  yolox:public YOLO_ONNX_Detector
{
public:
	yolox(std::string modelpath, std::string namesConfig, YoloParam para);
	int detect(cv::Mat& srcimg);

private:
	const int stride[3] = { 8, 16, 32 };
	//const string classesFile = "coco.names";   
	const int input_shape[2] = { 640, 640 };   //// height, width
	const double mean[3] = { 0.485, 0.456, 0.406 };
	const double std[3] = { 0.229, 0.224, 0.225 };

	int num_class;

	cv::Mat resize_image(cv::Mat srcimg, float* scale);
	void normalize(cv::Mat& srcimg);
	int get_max_class(float* scores);
};

class yolo5 :public YOLO_ONNX_Detector
{
	const int INPUT_WIDTH = 1024;
	const int INPUT_HEIGHT = 1024;

public:
	yolo5(std::string modelpath, std::string namesConfig, YoloParam para);
	int detect(cv::Mat& srcimg);

private:
	cv::Mat keepRatio(const cv::Mat &source);

};















#endif