import json
import torch
import torch.nn as nn
from torchvision import models
from torchvision.transforms import Compose, Resize, ToTensor, Normalize
import gradio as gr

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# 1) Load class names from your saved file
with open("classes.json", "r", encoding="utf-8") as f:
    class_names = json.load(f)

# 2) Build the model architecture (no downloading on the server)
model = models.efficientnet_b0(weights=None)
num_ftrs = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_ftrs, len(class_names))

# 3) Load your trained weights
state_dict = torch.load("best_efficientnetb0.pt", map_location=DEVICE)
model.load_state_dict(state_dict)
model.to(DEVICE)
model.eval()

# 4) Same preprocessing as validation/testing
val_transform = Compose([
    Resize((224, 224)),
    ToTensor(),
    Normalize(mean=[0.485, 0.456, 0.406],
              std=[0.229, 0.224, 0.225]),
])

@torch.no_grad()
def predict(image):
    x = val_transform(image).unsqueeze(0).to(DEVICE)
    logits = model(x)
    probs = torch.softmax(logits, dim=1)[0].cpu().tolist()
    # show top-5 nicely
    conf = {class_names[i]: float(probs[i]) for i in range(len(class_names))}
    return conf

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload a dog image"),
    outputs=gr.Label(num_top_classes=5, label="Prediction"),
    title="Dog Breed Classifier (EfficientNet-B0)",
    description="Upload a dog photo and the model predicts the breed with confidence."
)

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