import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
import torch
from PIL import Image

# Carica feature extractor e modello dal tuo Hub
MODEL_ID = "jaqen79/retail_images_classification_v1"
extractor = AutoFeatureExtractor.from_pretrained(MODEL_ID)
model = AutoModelForImageClassification.from_pretrained(MODEL_ID)

def predict(image: Image.Image):
    # Preprocess
    inputs = extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    probs = outputs.logits.softmax(dim=-1).tolist()[0]
    # Assumi che model.config.id2label esista
    labels = [model.config.id2label[i] for i in range(len(probs))]
    # Ritorna dizionario label→probabilità
    return {labels[i]: float(probs[i]) for i in range(len(probs))}

# Interfaccia Gradio
demo = gr.Interface(
    fn=predict,
    inputs=gr.components.Image(type="pil"), # Changed to gr.components.Image
    outputs=gr.components.Label(num_top_classes=5), # Changed to gr.components.Label
    title="Vision Transformer Demo",
    description="Carica un'immagine e il modello ritorna le classi con probabilità."
)

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