#ifndef YOLOV5_COMMON_H_
#define YOLOV5_COMMON_H_

#include <fstream>
#include <map>
#include <sstream>
#include <vector>
#include <tuple>
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
#include "yololayer.h"

using namespace std;
using namespace nvinfer1;

#define LOCATIONS 4

struct alignas(float) stDetection { //alignas(float)
    float bbox[LOCATIONS];
    float class_conf;
    float clear_conf;
    int class_id;
    int clear_id;
};

struct alignas(float) yolov5Detection { //alignas(float)
    float bbox[LOCATIONS];
    float class_conf;
    int class_id;
};


//HSV2BGR
std::tuple<uint8_t, uint8_t, uint8_t> hsv2bgr(float h, float s, float v);

//HSV2RGB
std::tuple<uint8_t, uint8_t, uint8_t> hsv2rgb(float h, float s, float v);

//随机颜色生成
std::tuple<uint8_t, uint8_t, uint8_t> random_color(int id);

//center x, y, w, h to x1,y1,x2,y2
void xywh2xyxy(float *xywh, float * xyxy);

//获取区域框1
cv::Rect get_rect(cv::Mat& img, float bbox[4]);

//获取区域框2
cv::Rect get_rectangle(cv::Mat& img, float bbox[4]);

//计算IOU
float iou(float lbox[4], float rbox[4]);

//比较两者置信度
bool cmp(const Yolo::Detection& a, const Yolo::Detection& b);

//非极大值抑制,默认0.5
void nms(std::vector<Yolo::Detection>& res, float *output, float conf_thresh, float nms_thresh);

//获取缩放比例
float GetResizeRatio(unsigned int img_width, unsigned int img_height, unsigned int model_width, unsigned int model_height);  

//坐标还原
void UpperVertexResetLocation(float fResizeRatio, unsigned int orig_width, unsigned int orig_height, stDetection &detection);  //左上顶点补边方式坐标还原
void CenterResetLocation(float fResizeRatio, unsigned int orig_width, unsigned int orig_height, unsigned int input_w, unsigned int input_h, stDetection &detection);  //中心补边方式坐标还原

//yolov5带清晰度后处理
void yolov5_clearity_decode_opencv_nms(std::vector<stDetection>& res, float *output, int outSize, unsigned int detSize, unsigned int classNum, unsigned int clearNum, float confThresh, float nmsThresh);

//yolov5后处理
void yolov5_decode_opencv_nms(std::vector<stDetection>& res, float *output, int outSize, unsigned int detSize, unsigned int classNum, float confThresh, float nmsThresh);

// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file);

IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::string lname, float eps) ;

ILayer* convBlock(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname);

ILayer* focus(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int inch, int outch, int ksize, std::string lname);

ILayer* bottleneck(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, bool shortcut, int g, float e, std::string lname);

ILayer* bottleneckCSP(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname);

ILayer* C3(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname);

ILayer* SPP(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int k1, int k2, int k3, std::string lname);

ILayer* SPPF(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int k, std::string lname);

std::vector<std::vector<float>> getAnchors(std::map<std::string, Weights>& weightMap, std::string lname);

IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, std::string lname, std::vector<IConvolutionLayer*> dets);

#endif