// Created by luozhiwang ([email protected])
// Date: 2020/8/27

#include "yolov5.h"

void initInputParams(common::InputParams &inputParams){
    inputParams.ImgH = 640;
    inputParams.ImgW = 640;
    inputParams.ImgC = 3;
    inputParams.BatchSize = 1;
    inputParams.HWC = false;
    inputParams.IsPadding = true;
    inputParams.InputTensorNames = std::vector<std::string>{"images"};
    inputParams.OutputTensorNames = std::vector<std::string>{"output", "692", "693"};
    inputParams.pFunction = [](const unsigned char &x){return static_cast<float>(x) /255;};
}

void initTrtParams(common::TrtParams &trtParams){
    trtParams.ExtraWorkSpace = 0;
    trtParams.FP32 = true;
    trtParams.FP16 = false;
    trtParams.Int32 = false;
    trtParams.Int8 = false;
    trtParams.worker = 4;
    trtParams.MaxBatch = 100;
    trtParams.MinTimingIteration = 1;
    trtParams.AvgTimingIteration = 2;
    trtParams.CalibrationTablePath = "/work/tensorRT-7/data/yolo5Int8.calibration";
    trtParams.CalibrationImageDir = "/data/dataset/coco/images/train2017";
    trtParams.OnnxPath = "/work/tensorRT-7/data/onnx/yolov5x.onnx";
    trtParams.SerializedPath = "/work/tensorRT-7/data/onnx/yolov5x.serialized";
}

std::vector<common::Anchor> initAnchors(){
    std::vector<common::Anchor> anchors;
    common::Anchor anchor;
    // 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90,  156,198,  373,326
    anchor.width = 10;
    anchor.height = 13;
    anchors.emplace_back(anchor);
    anchor.width = 16;
    anchor.height = 30;
    anchors.emplace_back(anchor);
    anchor.width = 32;
    anchor.height = 23;
    anchors.emplace_back(anchor);
    anchor.width = 30;
    anchor.height = 61;
    anchors.emplace_back(anchor);
    anchor.width = 62;
    anchor.height = 45;
    anchors.emplace_back(anchor);
    anchor.width = 59;
    anchor.height = 119;
    anchors.emplace_back(anchor);
    anchor.width = 116;
    anchor.height = 90;
    anchors.emplace_back(anchor);
    anchor.width = 156;
    anchor.height = 198;
    anchors.emplace_back(anchor);
    anchor.width = 373;
    anchor.height = 326;
    anchors.emplace_back(anchor);
    return anchors;
}

void initDetectParams(common::DetectParams &yoloParams){
    yoloParams.Strides = std::vector<int> {8, 16, 32};
    yoloParams.Anchors = initAnchors();
    yoloParams.AnchorPerScale = 3;
    yoloParams.NumClass = 80;
    yoloParams.NMSThreshold = 0.5;
    yoloParams.PostThreshold = 0.6;
}

int main(int argc, char **argv){
    common::InputParams inputParams;
    common::TrtParams trtParams;
    common::DetectParams yoloParams;
    initInputParams(inputParams);
    initTrtParams(trtParams);
    initDetectParams(yoloParams);

    Yolov5 *yolov5 = new Yolov5(inputParams, trtParams, yoloParams);
    yolov5->initSession(0);
    StreamProcess streamProcess(yolov5, 20);
    const char* video_path {"./data/image/test.mp4"};
    const char* render_path {".//data/image/render.mp4"};
    streamProcess.run(video_path, render_path);

    return 0;
}