#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <SimpleLog.h>
#include <Filesystem.h>
#include <YOLOV5.h>

int main()
{
    // 创建YOLOV5检测器
    migraphxSamples::DetectorYOLOV5 detector;
    migraphxSamples::InitializationParameterOfDetector initParamOfDetectorYOLOV5;
    initParamOfDetectorYOLOV5.configFilePath = CONFIG_FILE;
    migraphxSamples::ErrorCode errorCode=detector.Initialize(initParamOfDetectorYOLOV5);
    if(errorCode!=migraphxSamples::SUCCESS)
    {
        LOG_ERROR(stdout, "fail to initialize detector!\n");
        exit(-1);
    }
    LOG_INFO(stdout, "succeed to initialize detector\n");

    // 读取测试图像
    std::vector<cv::Mat> srcImages;
    cv::String folder = "../Resource/Images/DynamicPics";
    std::vector<cv::String> imagePathList;
    cv::glob(folder,imagePathList);
    for (int i = 0; i < imagePathList.size(); ++i)
    {
        cv:: Mat srcImage=cv::imread(imagePathList[i], 1);
        srcImages.push_back(srcImage); 
    }

    // 设置动态推理shape
    std::vector<std::vector<std::size_t>> inputShapes;
    inputShapes.push_back({1,3,416,416});
    inputShapes.push_back({1,3,608,608});

    for (int i = 0; i < srcImages.size(); ++i)
    {
        // 推理
        std::vector<migraphxSamples::ResultOfDetection> predictions;
        double time1 = cv::getTickCount();
        detector.Detect(srcImages[i], inputShapes[i], predictions);
        double time2 = cv::getTickCount();
        double elapsedTime = (time2 - time1)*1000 / cv::getTickFrequency();
        LOG_INFO(stdout, "inference image%d time:%f ms\n", i, elapsedTime);

        // 获取推理结果
        LOG_INFO(stdout,"========== Detection Image%d Results ==========\n", i);
        for(int j=0;j<predictions.size();++j)
        {
            migraphxSamples::ResultOfDetection result=predictions[j];
            cv::rectangle(srcImages[i],result.boundingBox,cv::Scalar(0,255,255),2);

            std::string label = cv::format("%.2f", result.confidence);
            label = result.className + " " + label;
            int left = predictions[j].boundingBox.x;
            int top = predictions[j].boundingBox.y;
            int baseLine;
            cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
            top = max(top, labelSize.height);
            cv::putText(srcImages[i], label, cv::Point(left, top-10), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 255), 2);

            LOG_INFO(stdout,"box:%d %d %d %d,label:%d,confidence:%f\n",predictions[j].boundingBox.x,
            predictions[j].boundingBox.y,predictions[j].boundingBox.width,predictions[j].boundingBox.height,predictions[j].classID,predictions[j].confidence);
        }
        std::string imgName = cv::format("Result%d.jpg", i);
        cv::imwrite(imgName, srcImages[i]);
        LOG_INFO(stdout,"Detection results have been saved to ./Result%d.jpg\n", i);
    }

    return 0;
}