#ifndef _MODEL_H_
#define _MODEL_H_

#include <iostream>

using namespace std;

#define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
#define DEVICE 0  // GPU id
#define BATCH_SIZE 1    //batch size
#define MAX_WORKSPAXE_SIZE (16 * (1 << 20)) //工作空间大小
#define MAX_IMAGE_INPUT_SIZE_THRESH 3000*3000
#define INPUT_CHANNEL 3     //输入通道数
#define LOCATIONS 4

//yolov5 目标检测

// == step 1 👇 yolov5 ==
// -- 车号/属性 --
#define STEP1_INPUT_BLOB_NAME "images"   //输入层名称
#define STEP1_OUTPUT_BLOB_NAME "output"   //输出层名称

#define STEP1_INPUT_CHANNEL 3     //输入通道数
#define STEP1_BATCH_SIZE 1        //batch size

#define STEP1_POS_CONF 5    //置信度加坐标
#define STEP1_LOCATIONS 4 //坐标数目

#define STEP1_NMS_THRESH 0.5     //step1 非极大值抑制阈值
#define STEP1_SCORE_THRESH 0.9   //step1 置信度(得分)阈值

#define STEP1_CLASS_NUM 19    //step1 分类数
#define STEP1_CLEAR_NUM 21    //step1 清晰度个数

#define STEP1_INPUT_H  960        //step1 输入图像高
#define STEP1_INPUT_W  960       //step1 输入图像宽
#define STEP1_BATCH_SIZE 1
#define STEP1_OUTPUT_HISTOGRAM_N 5    //不能超过139167

#define STEP1_BBOX_SIZE1  (STEP1_INPUT_H/32)
#define STEP1_BBOX_SIZE2  (STEP1_INPUT_H/16)
#define STEP1_BBOX_SIZE3  (STEP1_INPUT_H/8)

#define STEP1_DET_SIZE (STEP1_CLASS_NUM + STEP1_CLEAR_NUM + STEP1_POS_CONF)

#define STEP1_INPUT_SIZE    INPUT_CHANNEL*STEP1_INPUT_H*STEP1_INPUT_W      //input
#define STEP1_OUTPUT_SIZE   INPUT_CHANNEL*(STEP1_BBOX_SIZE1*STEP1_BBOX_SIZE1+STEP1_BBOX_SIZE2*STEP1_BBOX_SIZE2+STEP1_BBOX_SIZE3*STEP1_BBOX_SIZE3)*(STEP1_CLEAR_NUM+STEP1_CLASS_NUM+STEP1_POS_CONF)       //output
// -- 集装箱 --
#define STEP1_CONTAINER_INPUT_BLOB_NAME "images"   //输入层名称
#define STEP1_CONTAINER_OUTPUT_BLOB_NAME "output"   //输出层名称
#define STEP1_CONTAINER_INPUT_CHANNEL 3     //输入通道数
#define STEP1_CONTAINER_BATCH_SIZE 1        //batch size
#define STEP1_CONTAINER_POS_CONF 5    //置信度加坐标
#define STEP1_CONTAINER_LOCATIONS 4 //坐标数目
#define STEP1_CONTAINER_NMS_THRESH 0.5     //step1 非极大值抑制阈值
#define STEP1_CONTAINER_SCORE_THRESH 0.9   //step1 置信度(得分)阈值
#define STEP1_CONTAINER_CLASS_NUM 19    //step1 分类数
#define STEP1_CONTAINER_CLEAR_NUM 21    //step1 清晰度个数
#define STEP1_CONTAINER_INPUT_H  960        //step1 输入图像高
#define STEP1_CONTAINER_INPUT_W  960       //step1 输入图像宽
#define STEP1_CONTAINER_BATCH_SIZE 1
#define STEP1_CONTAINER_BBOX_SIZE1  (STEP1_CONTAINER_INPUT_H/32)
#define STEP1_CONTAINER_BBOX_SIZE2  (STEP1_CONTAINER_INPUT_H/16)
#define STEP1_CONTAINER_BBOX_SIZE3  (STEP1_CONTAINER_INPUT_H/8)
#define STEP1_CONTAINER_DET_SIZE (STEP1_CONTAINER_CLASS_NUM + STEP1_CONTAINER_CLEAR_NUM + STEP1_CONTAINER_POS_CONF)
#define STEP1_CONTAINER_INPUT_SIZE    INPUT_CHANNEL*STEP1_CONTAINER_INPUT_H*STEP1_CONTAINER_INPUT_W      //input
#define STEP1_CONTAINER_OUTPUT_SIZE   INPUT_CHANNEL*(STEP1_CONTAINER_BBOX_SIZE1*STEP1_CONTAINER_BBOX_SIZE1+STEP1_CONTAINER_BBOX_SIZE2*STEP1_CONTAINER_BBOX_SIZE2+STEP1_CONTAINER_BBOX_SIZE3*STEP1_CONTAINER_BBOX_SIZE3)*(STEP1_CONTAINER_CLEAR_NUM+STEP1_CONTAINER_CLASS_NUM+STEP1_CONTAINER_POS_CONF)       //output

// == step 1 👆 yolov5 ==
// == step 2 👇 yolov5 ==
// -- 车号/属性 --
#define STEP2_INPUT_BLOB_NAME "images"   //输入层名称
#define STEP2_OUTPUT_BLOB_NAME "output"   //输出层名称

#define STEP2_INPUT_CHANNEL 3     //输入通道数
#define STEP2_BATCH_SIZE 1        //batch size

#define STEP2_POS_CONF 5    //置信度加坐标
#define STEP2_LOCATIONS 4 //坐标数目

#define STEP2_NMS_THRESH 0.5      //step2 非极大值抑制阈值
#define STEP2_SCORE_THRESH 0.6    //step2 置信度(得分)阈值

#define STEP2_CLASS_NUM 47    //step2 分类数
#define STEP2_CLEAR_NUM 5     //step2 清晰度个数

#define STEP2_INPUT_H  608        //step2 输入图像高
#define STEP2_INPUT_W  608       //step2 输入图像宽
#define STEP2_BATCH_SIZE 1
#define STEP2_OUTPUT_HISTOGRAM_N 5    //step2 不能超过22743

#define STEP2_BBOX_SIZE1  (STEP2_INPUT_H/32)
#define STEP2_BBOX_SIZE2  (STEP2_INPUT_H/16)
#define STEP2_BBOX_SIZE3  (STEP2_INPUT_H/8)

#define STEP2_DET_SIZE (STEP2_CLASS_NUM + STEP2_CLEAR_NUM + STEP2_POS_CONF)

#define STEP2_INPUT_SIZE    INPUT_CHANNEL*STEP2_INPUT_H*STEP2_INPUT_W      //input
#define STEP2_OUTPUT_SIZE   INPUT_CHANNEL*(STEP2_BBOX_SIZE1*STEP2_BBOX_SIZE1+STEP2_BBOX_SIZE2*STEP2_BBOX_SIZE2+STEP2_BBOX_SIZE3*STEP2_BBOX_SIZE3)*(STEP2_CLEAR_NUM+STEP2_CLASS_NUM+STEP2_POS_CONF)       //output
// -- 集装箱 --
#define STEP2_CONTAINER_INPUT_BLOB_NAME "images"   //输入层名称
#define STEP2_CONTAINER_OUTPUT_BLOB_NAME "output"   //输出层名称
#define STEP2_CONTAINER_INPUT_CHANNEL 3     //输入通道数
#define STEP2_CONTAINER_BATCH_SIZE 1        //batch size
#define STEP2_CONTAINER_POS_CONF 5    //置信度加坐标
#define STEP2_CONTAINER_LOCATIONS 4 //坐标数目
#define STEP2_CONTAINER_NMS_THRESH 0.5     //step1 非极大值抑制阈值
#define STEP2_CONTAINER_SCORE_THRESH 0.9   //step1 置信度(得分)阈值
#define STEP2_CONTAINER_CLASS_NUM 7    //step1 分类数
#define STEP2_CONTAINER_CLEAR_NUM 21    //step1 清晰度个数
#define STEP2_CONTAINER_INPUT_H  960        //step1 输入图像高
#define STEP2_CONTAINER_INPUT_W  960       //step1 输入图像宽
#define STEP2_CONTAINER_BATCH_SIZE 1
#define STEP2_CONTAINER_BBOX_SIZE1  (STEP2_CONTAINER_INPUT_H/32)
#define STEP2_CONTAINER_BBOX_SIZE2  (STEP2_CONTAINER_INPUT_H/16)
#define STEP2_CONTAINER_BBOX_SIZE3  (STEP2_CONTAINER_INPUT_H/8)
#define STEP2_CONTAINER_DET_SIZE (STEP2_CONTAINER_CLASS_NUM + STEP2_CONTAINER_CLEAR_NUM + STEP2_CONTAINER_POS_CONF)
#define STEP2_CONTAINER_INPUT_SIZE    INPUT_CHANNEL*STEP2_CONTAINER_INPUT_H*STEP2_CONTAINER_INPUT_W      //input
#define STEP2_CONTAINER_OUTPUT_SIZE   INPUT_CHANNEL*(STEP2_CONTAINER_BBOX_SIZE1*STEP2_CONTAINER_BBOX_SIZE1+STEP2_CONTAINER_BBOX_SIZE2*STEP2_CONTAINER_BBOX_SIZE2+STEP2_CONTAINER_BBOX_SIZE3*STEP2_CONTAINER_BBOX_SIZE3)*(STEP2_CONTAINER_CLEAR_NUM+STEP2_CONTAINER_CLASS_NUM+STEP2_CONTAINER_POS_CONF)       //output

// == step 2 👆 yolov5 ==


// -- 来车检测 retinanet --
#define RETINANET_CLASSIFY_INPUT_CHANNEL 3  //输入通道数
#define RETINANET_CLASSIFY_BATCH_SIZE 1     //batch size

#define RETINANET_CLASSIFY_INPUT_H  537     // 模型输入高
#define RETINANET_CLASSIFY_INPUT_W  925     // 模型输入宽

#define RETINANET_CLASSIFY_INPUT_SIZE   (RETINANET_CLASSIFY_BATCH_SIZE*RETINANET_CLASSIFY_INPUT_CHANNEL*RETINANET_CLASSIFY_INPUT_W*RETINANET_CLASSIFY_INPUT_H)  //input
#define RETINANET_CLASSIFY_OUTPUT_SIZE  2   //output

#define RETINANET_CLASSIFY_INPUT_BLOB_NAME  "data"    //输入层名称
#define RETINANET_CLASSIFY_OUTPUT_BLOB_NAME  "train"   //输出层名称


typedef struct alignas(float) _Detection{
    float fBbox[LOCATIONS];
    float fClassConf;
    int iClassId;
}Detection;

typedef struct alignas(float) _ClearDetection{
    Detection detection;
    float fClearConf;
    int iClearId;
}ClearDetection;

typedef struct _ModelParam{
    unsigned int uiClassNum;

    unsigned int uiDetSize;
    float fScoreThreshold;
    float fNmsThreshold;
}ModelParam;

typedef struct _ClearModelParam{
    unsigned int uiClearNum;
    ModelParam modelParam;
}ClearModelParam;

typedef struct _ModelInfo{
    unsigned int uiModelWidth;
    unsigned int uiModelHeight;
    unsigned int uiInputSize;
    unsigned int uiOutputSize;
    unsigned int uiChannel;
    unsigned int uiBatchSize;
    std::string strInputBlobName;
    std::string strOutputBlobName;
}ModelInfo;

typedef struct _ClearModelInfo{
    ClearModelParam clearModelParam;
    ModelInfo modelInfo;
}ClearModelInfo;

typedef struct _CommonModelInfo{
    ModelParam modelParam;
    ModelInfo modelInfo;
}CommonModelInfo;

#endif //END OF _MODEL_H_