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

MODEL_ID  = "jaqen79/retail_images_classification_v1"
processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model     = AutoModelForImageClassification.from_pretrained(MODEL_ID)

def predict(image: Image.Image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    probs  = outputs.logits.softmax(dim=-1).tolist()[0]
    labels = [model.config.id2label[i] for i in range(len(probs))]
    return {labels[i]: float(probs[i]) for i in range(len(probs))}

demo = gr.Interface(
    fn=predict,
    inputs=gr.components.Image(type="pil"),
    outputs=gr.components.Label(num_top_classes=5),
    title="Retail Image classification using fine-tuned ViT",
    description="Upload an image and the model returns the classes with probabilities."
)

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