#include <torch/script.h>
#include <torch/torch.h>
#include <opencv2/opencv.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/core.hpp>
#include <iostream>
#include <memory>
#include <string>
#include <vector>
 
int main(int argc, const char *argv[]) {
    torch::DeviceType device_type;
    device_type = torch::kCUDA;  //torch::kCUDA  and torch::kCPU
    torch::Device device(device_type, 0);
    
    std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("model.pt");
     
    assert(module != nullptr);
    std::cout << "load model ok\n";
    module->to(device);
 
    auto image = cv::imread("284193,c89b000bf74ce6e.jpg");
    cv::Mat image_transfomed;
    cv::resize(image, image_transfomed, cv::Size(1024, 720),cv::INTER_AREA);
    cv::cvtColor(image_transfomed, image_transfomed, cv::COLOR_BGR2RGB);
 
    torch::Tensor tensor_image = torch::from_blob(image_transfomed.data, {image_transfomed.rows, image_transfomed.cols,3},torch::kByte).to(device);//hwc
    //ToTensor
    tensor_image = tensor_image.permute({ 2,0,1 });//chw
    tensor_image = tensor_image.toType(torch::kFloat);
    tensor_image = tensor_image.div(255.0);
    //normalize
    tensor_image[0] = tensor_image[0].sub_(0.485).div_(0.229);
    tensor_image[1] = tensor_image[1].sub_(0.456).div_(0.224);
    tensor_image[2] = tensor_image[2].sub_(0.406).div_(0.225);
    //add batch dimension
    tensor_image = tensor_image.unsqueeze(0);
    std::vector<torch::jit::IValue> inputs;
    inputs.emplace_back(tensor_image);
 
    auto t = (double) cv::getTickCount();
    at::Tensor output = module->forward(inputs).toTensor().cpu();
    at::IntList sizes = output.sizes();//shape
    //get result
    auto count = output[0][0].sum();
    t = (double) cv::getTickCount() - t;
    printf("count:%f,execution time = %gs\n",count,t / cv::getTickFrequency());
    inputs.pop_back();
    return 0;
}

import PIL.Image as Image
import torch
from torchvision import transforms
torch.backends.cudnn.benchmark = True
import torch
 
transform=transforms.Compose([
                       transforms.Resize((720, 1024)),
                       transforms.ToTensor(),
                       transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
                   ])
 
org = Image.open("284193,c89b000bf74ce6e.jpg")
img = transform(org).cuda()
traced_script_module = torch.jit.load('model.pt')
output = traced_script_module(img.unsqueeze(0))
print(output.sum())