#include "postprocess.h"


std::tuple<uint8_t, uint8_t, uint8_t> hsv2bgr(float h, float s, float v)
{
    const int h_i = static_cast<int>(h * 6);
    const float f = h * 6 - h_i;
    const float p = v * (1 - s);
    const float q = v * (1 - f*s);
    const float t = v * (1 - (1 - f) * s);
    float r, g, b;
    switch (h_i) {
    case 0:r = v; g = t; b = p;break;
    case 1:r = q; g = v; b = p;break;
    case 2:r = p; g = v; b = t;break;
    case 3:r = p; g = q; b = v;break;
    case 4:r = t; g = p; b = v;break;
    case 5:r = v; g = p; b = q;break;
    default:r = 1; g = 1; b = 1;break;}
    return std::make_tuple(static_cast<uint8_t>(b * 255), static_cast<uint8_t>(g * 255), static_cast<uint8_t>(r * 255));
}

std::tuple<uint8_t, uint8_t, uint8_t> hsv2rgb(float h, float s, float v)
{
    const int h_i = static_cast<int>(h * 6);
    const float f = h * 6 - h_i;
    const float p = v * (1 - s);
    const float q = v * (1 - f*s);
    const float t = v * (1 - (1 - f) * s);
    float r, g, b;
    switch (h_i) {
    case 0:r = v; g = t; b = p;break;
    case 1:r = q; g = v; b = p;break;
    case 2:r = p; g = v; b = t;break;
    case 3:r = p; g = q; b = v;break;
    case 4:r = t; g = p; b = v;break;
    case 5:r = v; g = p; b = q;break;
    default:r = 1; g = 1; b = 1;break;}
    return std::make_tuple(static_cast<uint8_t>(r * 255), static_cast<uint8_t>(g * 255), static_cast<uint8_t>(b * 255));
}

std::tuple<uint8_t, uint8_t, uint8_t> randomColor(int id)
{
    float h_plane = ((((unsigned int)id << 2) ^ 0x937151) % 100) / 100.0f;;
    float s_plane = ((((unsigned int)id << 3) ^ 0x315793) % 100) / 100.0f;
    // return hsv2bgr(h_plane, s_plane, 1);
    return hsv2rgb(h_plane, s_plane, 1);
}

//坐标转换
void xywh2xyxy(float *xywh, float * xyxy)
{
    xyxy[0] = (float)(xywh[0] - xywh[2] / 2);
    xyxy[1] = (float)(xywh[1] - xywh[3] / 2);
    xyxy[2] = (float)(xywh[0] + xywh[2] / 2);
    xyxy[3] = (float)(xywh[1] + xywh[3] / 2);
}

//获取区域框1
cv::Rect getRect(cv::Mat& img, unsigned int uiModelWidth, unsigned int uiModelHeight, float fBbox[4]) 
{
    float l, r, t, b;
    float r_w = uiModelWidth / (img.cols * 1.0);
    float r_h = uiModelHeight / (img.rows * 1.0);
    if(r_h > r_w){
        l = fBbox[0] - fBbox[2] / 2.f;
        r = fBbox[0] + fBbox[2] / 2.f;
        t = fBbox[1] - fBbox[3] / 2.f - (uiModelHeight - r_w * img.rows) / 2;
        b = fBbox[1] + fBbox[3] / 2.f - (uiModelHeight - r_w * img.rows) / 2;
        l = l / r_w;
        r = r / r_w;
        t = t / r_w;
        b = b / r_w;
    }else{
        l = fBbox[0] - fBbox[2] / 2.f - (uiModelWidth - r_h * img.cols) / 2;
        r = fBbox[0] + fBbox[2] / 2.f - (uiModelWidth - r_h * img.cols) / 2;
        t = fBbox[1] - fBbox[3] / 2.f;
        b = fBbox[1] + fBbox[3] / 2.f;
        l = l / r_h;
        r = r / r_h;
        t = t / r_h;
        b = b / r_h;
    }
    return cv::Rect(round(l), round(t), round(r - l), round(b - t));
}

//获取区域框2
cv::Rect getRectangle(cv::Mat& img, float fBbox[4]) 
{
    return cv::Rect(round(fBbox[0]), round(fBbox[1]), round(fBbox[2]-fBbox[0]), round(fBbox[3]-fBbox[1]));
}

//计算IOU
float iou(float fLbox[4], float fRbox[4])
{
    float interBox[] = {
        (std::max)(fLbox[0] - fLbox[2] / 2.f , fRbox[0] - fRbox[2] / 2.f), //left
        (std::min)(fLbox[0] + fLbox[2] / 2.f , fRbox[0] + fRbox[2] / 2.f), //right
        (std::max)(fLbox[1] - fLbox[3] / 2.f , fRbox[1] - fRbox[3] / 2.f), //top
        (std::min)(fLbox[1] + fLbox[3] / 2.f , fRbox[1] + fRbox[3] / 2.f), //bottom
    };

    if (interBox[2] > interBox[3] || interBox[0] > interBox[1])
        return 0.0f;

    float fInterBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]);
    return fInterBoxS / (fLbox[2] * fLbox[3] + fRbox[2] * fRbox[3] - fInterBoxS);
}

//比较两者置信度
bool confCmp(const Detection& a, const Detection& b) 
{
    return a.fClassConf > b.fClassConf;
}

//获取缩放比例
float GetResizeRatio(unsigned int uiImgWidth, unsigned int uiImgHeight, unsigned int uiModelWidth, unsigned int uiModelHeight)
{
    float fRatioW = static_cast<float>(uiImgWidth) / uiModelWidth;
    float fRatioH = static_cast<float>(uiImgHeight) / uiModelHeight;

    return (fRatioW - fRatioH > 1e-5) ? fRatioW : fRatioH;
}

//左上顶点补边方式坐标还原
void UpperVertexResetLocation(float fResizeRatio, unsigned int uiOrigWidth, unsigned int uiOrigHeight, Detection &detection)
{
    for(int i=0; i<4; i++){
        detection.fBbox[i] = detection.fBbox[i] * fResizeRatio;
    }

    detection.fBbox[0] = (detection.fBbox[0] < uiOrigWidth) ? detection.fBbox[0] : uiOrigWidth;
    detection.fBbox[1] = (detection.fBbox[1] < uiOrigHeight) ? detection.fBbox[1] : uiOrigHeight;
    detection.fBbox[2] = (detection.fBbox[2] < uiOrigWidth) ? detection.fBbox[2] : uiOrigWidth;
    detection.fBbox[3] = (detection.fBbox[3] < uiOrigHeight) ? detection.fBbox[3] : uiOrigHeight;
}

//中心补边方式坐标还原
void CenterResetLocation(float fResizeRatio, unsigned int uiOrigWidth, unsigned int uiOrigHeight, unsigned int uiInputWidth, unsigned int uiInputHeight, Detection &detection)
{
    int w, h, x, y;
    float r_w = uiInputWidth / (uiOrigWidth*1.0);
    float r_h = uiInputHeight / (uiOrigHeight*1.0);
    if (r_h > r_w) {
        w = uiInputWidth;
        h = r_w * uiOrigHeight;
        x = 0;
        y = (uiInputHeight - h) / 2;
        detection.fBbox[1] -= y;
        detection.fBbox[3] -= y;
    } else {
        w = r_h * uiOrigWidth;
        h = uiInputHeight;
        x = (uiInputWidth - w) / 2;
        y = 0;
        detection.fBbox[0] -= x;
        detection.fBbox[2] -= x;
    }

    for(int i=0; i<4; i++){
        detection.fBbox[i] = detection.fBbox[i] * fResizeRatio;
    }

    detection.fBbox[0] = (detection.fBbox[0] < uiOrigWidth) ? detection.fBbox[0] : uiOrigWidth;
    detection.fBbox[1] = (detection.fBbox[1] < uiOrigHeight) ? detection.fBbox[1] : uiOrigHeight;
    detection.fBbox[2] = (detection.fBbox[2] < uiOrigWidth) ? detection.fBbox[2] : uiOrigWidth;
    detection.fBbox[3] = (detection.fBbox[3] < uiOrigHeight) ? detection.fBbox[3] : uiOrigHeight;
}

//非极大值抑制,默认0.5
void yolov5ClearDecodeOpenCVNms(std::vector<ClearDetection>& vecRes, float *fOutput, unsigned int uiOutSize, unsigned int uiDetSize, unsigned int uiClassNum, unsigned int uiClearNum, float fConfThresh = 0.5, float fNmsThresh = 0.4) 
{
    //1.筛选出第一轮结果，根据conf > conf_thresh
    std::vector<int> vecFilterList;
    for(int i = 0; i < uiOutSize /uiDetSize ; ++i){    //锚框数量(遍历锚框筛去置信度过于低的)
        if(fOutput[uiDetSize * i + 4] > fConfThresh){
            vecFilterList.emplace_back(i);    // 记录下所有置信度大于conf_thresh的锚框
        }
    }
    if(vecFilterList.size() == 0) return;

    //2.查找剩余锚框中每个类得分最高值并记录标签
    std::vector<ClearDetection> vecResult;
    std::vector<cv::Rect> vecBoxes;
    std::vector<float> vecScores;
    for(int i : vecFilterList){
        float* pClassConfidence = &fOutput[uiDetSize * i + 5];    //类别
        float fClassConf = *pClassConfidence++;
        int iClassLabel = 0;
        for(int j = 1; j <uiClassNum; ++j, ++pClassConfidence){   //N个类别中查找置信度最高的,并记录标签
            if(*pClassConfidence > fClassConf){
                fClassConf = *pClassConfidence;
                iClassLabel = j;
            }
        }

        float* pClearConfidence = &fOutput[uiDetSize * i + 5 + uiClassNum];    //清晰度
        float fClearConf  = *pClearConfidence++;
        int iClearLabel = 0;
        for(int n = 1; n <uiClearNum; ++n, ++pClearConfidence){   //N个清晰度中查找置信度最高的,并记录标签
            if(*pClearConfidence > fClearConf){
                fClearConf = *pClearConfidence;
                iClearLabel = n;
            }
        }

        float xywh[4] = {fOutput[i*uiDetSize + 0], fOutput[i*uiDetSize + 1], fOutput[i*uiDetSize + 2], fOutput[i*uiDetSize + 3]};
        float xyxy[4];
        xywh2xyxy(xywh, xyxy);  //转换后的坐标
        ClearDetection clearDetection;
        for(int n = 0; n<4; n++){
            clearDetection.detection.fBbox[n] = xyxy[n];
        }
        clearDetection.detection.fClassConf = fOutput[uiDetSize * i + 4] * fClassConf;
        clearDetection.detection.iClassId = iClassLabel;
        clearDetection.fClearConf = fOutput[uiDetSize * i + 4] * fClearConf;
        clearDetection.iClearId = iClearLabel;

        //总得分大于某阈值才进行NMS,此处阈值可以乘某系数!
        if(clearDetection.detection.fClassConf>fConfThresh) {
            cv::Rect tempRect;
            tempRect.x = xyxy[0];
            tempRect.y = xyxy[1];
            tempRect.width = std::abs(xyxy[2] - xyxy[0]);
            tempRect.height = std::abs(xyxy[3] - xyxy[1]);
            
            vecBoxes.push_back(tempRect);
            vecScores.push_back(clearDetection.detection.fClassConf);
            vecResult.push_back(clearDetection);
        }
    }

    std::vector<int> vecIdx; //保留nms后框的索引
    cv::dnn::NMSBoxes(vecBoxes, vecScores, fConfThresh, fNmsThresh, vecIdx);

    for (std::size_t i = 0; i < vecIdx.size(); i++){
        vecRes.push_back(vecResult[vecIdx[i]]);
    }
}


//非极大值抑制,默认0.5
void yolov5DecodeOpenCVNms(std::vector<Detection>& vecRes, float *fOutput, unsigned int uiOutSize, unsigned int uiDetSize, unsigned int uiClassNum, float fConfThresh = 0.5, float fNmsThresh = 0.4) 
{
    //1.筛选出第一轮结果，根据conf > conf_thresh
    std::vector<int> vecFilterList;
    for(int i = 0; i < uiOutSize /uiDetSize ; ++i){    //锚框数量(遍历锚框筛去置信度过于低的)
        if(fOutput[uiDetSize * i + 4] > fConfThresh){
            vecFilterList.emplace_back(i);    // 记录下所有置信度大于conf_thresh的锚框
        }
    }
    if(vecFilterList.size() == 0) return;

    //2.查找剩余锚框中每个类得分最高值并记录标签
    std::vector<Detection> vecResult;
    std::vector<cv::Rect> vecBoxes;
    std::vector<float> vecScores;
    for(int i : vecFilterList){
        float* pClassConfidence = &fOutput[uiDetSize * i + 5];    //类别
        float fClassConf = *pClassConfidence++;
        int iClassLabel = 0;
        for(int j = 1; j <uiClassNum; ++j, ++pClassConfidence){   //N个类别中查找置信度最高的,并记录标签
            if(*pClassConfidence > fClassConf){
                fClassConf = *pClassConfidence;
                iClassLabel = j;
            }
        }

        float xywh[4] = {fOutput[i*uiDetSize + 0], fOutput[i*uiDetSize + 1], fOutput[i*uiDetSize + 2], fOutput[i*uiDetSize + 3]};
        float xyxy[4];
        xywh2xyxy(xywh, xyxy);  //转换后的坐标
        Detection detection;
        for(int n = 0; n<4; n++){
            detection.fBbox[n] = xyxy[n];
        }
        detection.fClassConf = fOutput[uiDetSize * i + 4] * fClassConf;
        detection.iClassId = iClassLabel;

        //总得分大于某阈值才进行NMS,此处阈值可以乘某系数!
        if(detection.fClassConf>fConfThresh) {
            cv::Rect tempRect;
            tempRect.x = xyxy[0];
            tempRect.y = xyxy[1];
            tempRect.width = std::abs(xyxy[2] - xyxy[0]);
            tempRect.height = std::abs(xyxy[3] - xyxy[1]);
            
            vecBoxes.push_back(tempRect);
            vecScores.push_back(detection.fClassConf);
            vecResult.push_back(detection);
        }
    }

    std::vector<int> vecIdx; //保留nms后框的索引
    cv::dnn::NMSBoxes(vecBoxes, vecScores, fConfThresh, fNmsThresh, vecIdx);

    for (std::size_t i = 0; i < vecIdx.size(); i++){
        vecRes.push_back(vecResult[vecIdx[i]]);
    }
}