Update app.py

app.py

@@ -1,31 +1,26 @@
 import gradio as gr
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from transformers import AutoImageProcessor, 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)
+MODEL_ID  = "jaqen79/retail_images_classification_v1"
+processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+model     = AutoModelForImageClassification.from_pretrained(MODEL_ID)
 
 def predict(image: Image.Image):
-    # Preprocess
-    inputs = extractor(images=image, return_tensors="pt")
+    inputs = processor(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
+    probs  = outputs.logits.softmax(dim=-1).tolist()[0]
     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à."
+    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__":