#include <opencv2/opencv.hpp>
#include <ncnn/gpu.h>
#include <ncnn/net.h>

//模型路徑
std::string celeba_model = "models/Celeba.bin";
std::string celeba_param = "models/Celeba.param";
std::string face_v1_model = "models/FacePointV1.bin";
std::string face_v1_param = "models/FacePointV1.param";
std::string face_v2_model = "models/FacePointV2.bin";
std::string face_v2_param = "models/FacePointV2.param";

//讀取模型
int readModels(ncnn::Net& ncnn_net, std::string param_path, std::string model_path, bool use_gpu = true);

//推理，target_w，target_h為推理尺寸，越往大的改，細節(jié)保留就越多，也越吃GPU算力
int animeInference(const ncnn::Net& ncnn_net, const cv::Mat& cv_src, cv::Mat& cv_dst,int target_w = 512, int target_h = 512)
{
    cv::Mat cv_backup = cv_src.clone();
    const int w = cv_src.cols;
    const int h = cv_src.rows;

    const float mean_vals[3] = { 127.5f, 127.5f,  127.5f };
    const float norm_vals[3] = { 1 / 127.5f, 1 / 127.5f, 1 / 127.5f };
    ncnn::Mat in = ncnn::Mat::from_pixels_resize(cv_backup.data, ncnn::Mat::PIXEL_BGR2RGB, w, h, target_w, target_h);
    in.substract_mean_normalize(mean_vals, norm_vals);
    ncnn::Mat out;
    
    ncnn::Extractor ex = ncnn_net.create_extractor();

    ex.input("input", in);
    ex.extract("out", out);
    
    cv::Mat result(out.h, out.w, CV_32FC3);
    for (int i = 0; i < out.c; i++)
    {
        float* out_data = out.channel(i);
        for (int h = 0; h < out.h; h++)
        {
            for (int w = 0; w < out.w; w++)
            {
                result.at<cv::Vec3f>(h, w)[2 - i] = out_data[h * out.h + w];
            }
        }
    }
    cv::Mat result8U(out.h, out.w, CV_8UC3);
    result.convertTo(result8U, CV_8UC3, 255.0/2, 255.0/2);
    result8U.copyTo(cv_dst);
    cv::resize(cv_dst, cv_dst, cv_src.size());
    return 0;
}

//調(diào)用，三種風(fēng)格
int styletransferResult(const cv::Mat& cv_src, std::vector<cv::Mat>& cv_dsts, ncnn::Net& face_v1_net, ncnn::Net& face_v2_net, ncnn::Net& celeba_net)
{
    animeInference(face_v1_net, cv_src, cv_dsts.at(0));
    animeInference(face_v2_net, cv_src, cv_dsts.at(1));
    animeInference(celeba_net, cv_src, cv_dsts.at(2));

}

int main(int argc, char** argv)
{
	//圖像路徑
    std::string path = "images";
    std::vector<std::string> filenames;
    cv::glob(path, filenames, false);

    ncnn::Net face_v1_net, face_v2_net, celeba_net;

    readModels(face_v1_net, face_v1_param, face_v1_model);
    readModels(face_v2_net, face_v1_param, face_v1_model);
    readModels(celeba_net, celeba_param, celeba_model);

    for(auto v : filenames)
    {

        cv::Mat cv_src = cv::imread(v, 1);

        std::vector<cv::Mat> cv_dsts(3);

        double start = static_cast<double>(cv::getTickCount());
        styletransferResult(cv_src, cv_dsts, face_v1_net, face_v2_net, celeba_net);
        double time = ((double)cv::getTickCount() - start) / cv::getTickFrequency();
        std::cout << "time:" << time << "(s)" << std::endl;
 

        cv_src.push_back(cv_dsts.at(0));
        cv_dsts.at(1).push_back(cv_dsts.at(2));

        cv::Mat des;
        des.create(cv_src.rows, cv_src.cols*2, cv_src.type());
        cv::Mat r1 = des(cv::Rect(0, 0, cv_src.cols, cv_src.rows));
        cv_src.copyTo(r1);

        cv::Mat r2 = des(cv::Rect(cv_src.cols, 0, cv_src.cols,cv_src.rows));
        cv_dsts.at(1).copyTo(r2);

        imshow("style", des);
        cv::imwrite(std::to_string(j)+".jpg", des);
        
      // cv::waitKey();
    }
    return 0;
}

int readModels(ncnn::Net& ncnn_net,std::string param_path,std::string model_path, bool use_gpu)
{
    bool has_gpu = false;

#if NCNN_VULKAN
    ncnn::create_gpu_instance();
    has_gpu = ncnn::get_gpu_count() > 0;
#endif

    bool to_use_gpu = has_gpu && use_gpu;
    ncnn_net.opt.use_vulkan_compute = to_use_gpu;

    int rp = ncnn_net.load_param(param_path.c_str());

    int rb = ncnn_net.load_model(model_path.c_str());

    if (rp < 0 || rb < 0)
    {
        return 1;
    }
    return 0;
}