import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import gradio as gr

DEVICE = torch.device("cpu")
NUM_CLASSES = 64
MODEL_PATH = "resnet18_animals_64.pth"
CLASS_NAMES_FILE = "classes.txt"

with open(CLASS_NAMES_FILE, "r") as f:
    class_names = [line.strip() for line in f.readlines()]

def load_model():
    model = models.resnet18(weights=None)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, NUM_CLASSES)
    
    model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
    model.to(DEVICE)
    model.eval()
    return model

model = load_model()

image_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def predict(image):
    img_tensor = image_transforms(image).unsqueeze(0)
    img_tensor = img_tensor.to(DEVICE)
    
    with torch.no_grad():
        outputs = model(img_tensor)
        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
    
    top3_prob, top3_catid = torch.topk(probabilities, 3)
    
    confidences = {}
    for i in range(top3_prob.size(0)):
        confidences[class_names[top3_catid[i]]] = float(top3_prob[i])
    
    return confidences

title = "Animal Classifier 🐾"
description = "Nhận diện 64 loài động vật sử dụng mô hình ResNet-18."
examples = []

interface = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Label(num_top_classes=3),
    title=title,
    description=description
)

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