File size: 1,203 Bytes
8e1d2f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
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() |