import torch
import torchvision.models as models
from PIL import Image
import numpy as np
from torchvision import transforms

model_resnet = models.resnet50(pretrained=True).cuda()

# model_resnet.load_state_dict(torch.load("resnet_Epoch_4_Top1_99.75845336914062.pkl"))
model_resnet.eval()
# 自己選擇任意一張圖片，并將它的路徑寫在open方法里，用來讀取圖像，我這里路徑就是‘111.jpg'了
image = Image.open("111.jpg").convert('RGB')

tf = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
])

img = tf(image)
img = img.unsqueeze(dim=0)
print(img.shape)
input = torch.rand(1, 3, 224, 224).cuda()

traced_script_module_resnet = torch.jit.trace(model_resnet, input)

output = traced_script_module_resnet(img.cuda())
print(output.shape)
pred = torch.argmax(output, dim=1)
print(pred)
traced_script_module_resnet.save("model_resnet_jit_cuda.pt")

#include <torch/script.h> // One-stop header.
#include <opencv2/opencv.hpp>
#include <iostream>
#include <memory>

//https://pytorch.org/tutorials/advanced/cpp_export.html

std::string image_path = "../../111.jpg";

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

	// Deserialize the ScriptModule from a file using torch::jit::load().
	//std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("../../model_resnet_jit.pt");
	using torch::jit::script::Module;
	Module module = torch::jit::load("../../model_resnet_jit_cuda.pt");
	module.to(at::kCUDA);

	//assert(module != nullptr);
	//std::cout << "ok\n";

	//輸入圖像
	auto image = cv::imread(image_path, cv::ImreadModes::IMREAD_COLOR);
	cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
	cv::Mat image_transfomed;
	cv::resize(image, image_transfomed, cv::Size(224, 224));

	// 轉換為Tensor
	torch::Tensor tensor_image = torch::from_blob(image_transfomed.data,
		{ image_transfomed.rows, image_transfomed.cols,3 }, torch::kByte);
	tensor_image = tensor_image.permute({ 2,0,1 });
	tensor_image = tensor_image.toType(torch::kFloat);
	tensor_image = tensor_image.div(255);
	tensor_image = tensor_image.unsqueeze(0);
	tensor_image = tensor_image.to(at::kCUDA);

	// 網(wǎng)絡前向計算
	at::Tensor output = module.forward({ tensor_image }).toTensor();
	//std::cout << "output:" << output << std::endl;

	auto prediction = output.argmax(1);
	std::cout << "prediction:" << prediction << std::endl;

	int maxk = 3;
	auto top3 = std::get<1>(output.topk(maxk, 1, true, true));

	std::cout << "top3: " << top3 << '\n';

	std::vector<int> res;
	for (auto i = 0; i < maxk; i++) {
		res.push_back(top3[0][i].item().toInt());
	}
	for (auto i : res) {
		std::cout << i << " ";
	}
	std::cout << "\n";

	system("pause");
}