Update app.py

app.py

@@ -1,29 +1,47 @@
-import gradio as gr
-from transformers import ViTFeatureExtractor, ViTForImageClassification
-from PIL import Image
-import torch
-
-# Load pre-trained model and feature extractor
-model_name = "google/vit-base-patch16-224"
-feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
-model = ViTForImageClassification.from_pretrained(model_name)
-
-# Define the prediction function
-def classify_image(img):
-    inputs = feature_extractor(images=img, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
-    predicted_label = model.config.id2label[predicted_class_idx]
-    return predicted_label
-
-# Build the Gradio interface
-interface = gr.Interface(fn=classify_image,
-                         inputs=gr.Image(type="pil"),
-                         outputs="text",
-                         title="Image Classification with ViT",
-                         description="Upload an image and classify it using Vision Transformer (ViT)")
-
-# Launch the app
-interface.launch()
+import gradio as gr
+import torch
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from PIL import Image
+
+# 1. Load a modern SOTA model (DINOv2)
+# This model is faster and more accurate than the original ViT
+model_name = "facebook/dinov2-base-imagenet1k-1-layer"
+processor = AutoImageProcessor.from_pretrained(model_name)
+model = AutoModelForImageClassification.from_pretrained(model_name)
+
+def classify_image(img):
+    # 2. Pre-process image
+    inputs = processor(images=img, return_tensors="pt")
+    
+    # 3. Inference
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # 4. Convert logits to probabilities (0% to 100%)
+    logits = outputs.logits
+    probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
+    
+    # 5. Extract top 5 results for a better UI
+    top5_prob, top5_indices = torch.topk(probabilities, 5)
+    
+    # Create a dictionary of {Label: Probability} for Gradio's Label component
+    confidences = {
+        model.config.id2label[idx.item()]: float(prob) 
+        for prob, idx in zip(top5_prob, top5_indices)
+    }
+    
+    return confidences
+
+# 6. Build a modern UI
+demo = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    # gr.Label automatically creates a beautiful bar chart for probabilities
+    outputs=gr.Label(num_top_classes=5, label="Predictions"),
+    title="Next-Gen Image Classification",
+    description="Running on **Meta's DINOv2** foundation model. Upload any image to see the top 5 predicted categories.",
+    theme="soft" 
+)
+
+if __name__ == "__main__":
+    demo.launch()