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


void MIGraphXSamplesUsage(char *programName)
{
    printf("Usage : %s <index> \n", programName);
    printf("index:\n");
    printf("\t 0) DeepLabV3 Single Image Sample.\n");
    printf("\t 1) DeepLabV3 Multiple Image Sample.\n");
}


int main(int argc,char *argv[]){

    if (argc < 2 || argc > 2){
        MIGraphXSamplesUsage(argv[0]);
        return -1;
    }
    if (!strncmp(argv[1], "-h", 2)){
        MIGraphXSamplesUsage(argv[0]);
        return 0;
    }

    migraphxSamples::InitializationParameterOfSegmentation initParamOfSegmentationUnet;
    std::vector<cv::Mat> srcImages;
    switch(*argv[1]){
        case '0':
            srcImages.push_back(cv::imread("../Resource/Images/000001.jpg", 1));
            initParamOfSegmentationUnet.loadMode = 1;                                         //加载单个图片
            break;
        case '1':
            std::string folderPath = "../Resource/Images/";
            std::string imageExt = "*.jpg";
            std::vector<cv::String> imagePaths;
            cv::glob(folderPath + imageExt, imagePaths, false);
            for(const auto& path : imagePaths){
                srcImages.push_back(cv::imread(path, 1));
            }
            initParamOfSegmentationUnet.loadMode = 0;                                        //加载多个图片
            break;
    }

    migraphxSamples::DeepLabV3 deeplabv3;
    initParamOfSegmentationUnet.configFilePath = CONFIG_FILE;
    migraphxSamples::ErrorCode errorCode       = deeplabv3.Initialize(initParamOfSegmentationUnet);
    if(errorCode != migraphxSamples::SUCCESS){
        LOG_ERROR(stdout, "fail to initialize DeepLabV3!\n");
        exit(-1);
    }
    LOG_INFO(stdout, "succeed to initialize DeepLabV3\n");


    // 推理
    std::vector<cv::Mat> maskImages;
    deeplabv3.Segmentation(srcImages, maskImages);
    LOG_INFO(stdout, "========== Segmentation Results ==========\n");
    LOG_INFO(stdout, "Segmentation results have been saved to build directory\n");
    for(int i = 0;i<maskImages.size();i++){
        std::string fileName ="Result_"+std::to_string(i+1)+".jpg";
        cv::imwrite(fileName, maskImages[i]);
    }
    

    return 0;

}