import torch
import torchvision.transforms as transforms
from torchvision import models
from torch import nn
import torch.nn.functional as F
import gradio as gr
import json


# Get cpu or gpu device for training.
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

model_path = './resnet18_flower.pth'
num_classes = 5
label_name_list = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']

# step1: 创建并加载模型
resnet18 = models.resnet18(pretrained=False)
num_ftrs = resnet18.fc.in_features
resnet18.fc = nn.Linear(num_ftrs, num_classes)
resnet18 = resnet18.to(device)

resnet18.load_state_dict(torch.load(model_path,map_location=torch.device('cpu')))

model = resnet18
# model = base_model_vgg13
model.eval()


# step2: 图片转换
train_transform = transforms.Compose([
                    # transforms.RandomRotation(5),
                    # transforms.RandomHorizontalFlip(),
                    transforms.Resize((224,224)),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                    ])



def predict(inp):
  # 定义预处理变换
  transform = train_transform

  # 加载图片并进行预处理
  # image = Image.open(image_path)
  image = transform(inp).unsqueeze(0).to(device)

  # step3：使用模型进行预测
  with torch.no_grad():
      output = model(image)
      
  # step4：数据后处理
  # 计算预测概率
  pred_score = nn.functional.softmax(output[0], dim=0)
  pred_score = pred_score.cpu().numpy()

  # 获取预测结果
  pred_index = torch.argmax(output, dim=1).item()
  pred_label = label_name_list[pred_index]
  
  # 转为json字符串格式
  result_dict = {'pred_score':str(max(pred_score)),'pred_index':str(pred_index),'pred_label':pred_label }
  result_json = json.dumps(result_dict)

  return result_json


demo = gr.Interface(fn=predict, 
             inputs=gr.Image(type="pil"),
             outputs="text",
             examples=["./592px-Red_sunflower.jpg"],
             )
             
# demo.launch(debug=True)
demo.launch()