import gradio as gr
import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
from PIL import Image
import json
import os

MODEL_PATH = "best_model.pth"
with open("class_names.json", "r") as f:
    CLASS_NAMES = json.load(f)

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

weights = EfficientNet_V2_S_Weights.IMAGENET1K_V1
model = efficientnet_v2_s(weights=weights)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, len(CLASS_NAMES))
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
model.eval().to(device)

mean = getattr(weights, "meta", {}).get("mean", [0.485, 0.456, 0.406])
std = getattr(weights, "meta", {}).get("std", [0.229, 0.224, 0.225])

transform = transforms.Compose([
    transforms.Resize((384, 384)),
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std),
])

def predict(image):
    image = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        outputs = model(image)
        probs = torch.nn.functional.softmax(outputs, dim=1)[0]
    results = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
    predicted_label = CLASS_NAMES[probs.argmax().item()]
    return predicted_label, results

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Label(label="Prediction"), gr.JSON(label="Confidence Scores")],
    title="StrAI - Cat Identifier",
    description="Upload an image to identify which cat it is."
)

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