import gradio as gr
from fastai.vision.all import load_learner, PILImage, ToTensor, Normalize
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import numpy as np

learn = load_learner("model.pkl")
labels = learn.dls.vocab

# Build inference transform manually — bypasses fasttransform bug
tfms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    ),
])

def predict_image(img):
    img = img.convert("RGB")
    tensor = tfms(img).unsqueeze(0)  # add batch dim
    
    learn.model.eval()
    with torch.no_grad():
        out = learn.model(tensor)
        probs = F.softmax(out[0], dim=0)
    
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=5),
    title="🌸 Flower Classifier",
    description="Upload a flower photo to classify it using a fastai model.",
)

demo.launch()