import gradio as gr
import torch
import json
from PIL import Image
import torchvision.transforms as transforms
from fastai.vision.all import *

with open("clases.json") as f:
    clases = json.load(f)

model = torch.load("full_model.pth", map_location="cpu", weights_only=False)
model.eval()

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

def clasificar_zona(imagen):
    img = Image.fromarray(imagen).convert("RGB")
    tensor = tfms(img).unsqueeze(0)
    with torch.no_grad():
        probs = torch.softmax(model(tensor), dim=1)[0]
    return dict(zip(clases, map(float, probs)))

demo = gr.Interface(
    fn=clasificar_zona,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=4),
    title="🔐 Security Room Classifier",
    description="Identifica en qué zona de la casa fue tomada la imagen.",
)
demo.launch(share=True)