import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
from timm import create_model
import json

with open('class_names.json', 'r') as json_file:
    class_mapping = json.load(json_file)


# 加载模型
def load_model(model_path):
    model = create_model('resnet18', pretrained=False, num_classes=len(class_mapping))
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
    model.eval()
    return model

model = load_model("res18_nabird555_acc596.pth")

# 定义图像预处理
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    return transform(image).unsqueeze(0)

# 定义推理函数
def classify_image(image):
    image = preprocess_image(image)
    with torch.no_grad():
        outputs = model(image)
        
        _, predicted_class = torch.max(outputs, 1)
        predicted_class_idx = predicted_class.item()
        predicted_class_name = class_mapping[str(predicted_class_idx)]

    return predicted_class_name

# 创建 Gradio 接口
title = "Bird Species Classifier"
description = "Upload an image of a bird, and the model will predict its species."

interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title=title,
    description=description,
)

if __name__ == "__main__":
    interface.launch()