Update app.py

app.py

@@ -1,25 +1,32 @@
-import gradio as gr
-from transformers import pipeline
-from PIL import Image
-
-# Load an image classification pipeline
-classifier = pipeline("image-classification", model="google/vit-base-patch16-224")
-
-def classify_image(img, top_k=3):
-    if img is None:
-        return {"Error": 1.0}
-    results = classifier(img, top_k=top_k)
-    # Return as {label: score} for Gradio Label component
-    return {r["label"]: float(r["score"]) for r in results}
-
-# Gradio interface
-demo = gr.Interface(
-    fn=classify_image,
-    inputs=[gr.Image(type="pil", label="Upload Image"), gr.Slider(1, 5, value=3, label="Top K Predictions")],
-    outputs=gr.Label(num_top_classes=5, label="Predictions"),
-    title="Image Classification App",
-    description="Upload an image and the model will predict the top objects in it."
-)
-
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+from transformers import pipeline
+from PIL import Image
+
+# Use ResNet-50 model (1000 common ImageNet categories like dog, cat, car, etc.)
+classifier = pipeline("image-classification", model="microsoft/resnet-50")
+
+def classify_image(img, top_k=3):
+    """
+    Takes an uploaded image, runs classification, 
+    and returns the top_k labels with confidence scores.
+    """
+    if img is None:
+        return {"Error": 1.0}
+    
+    results = classifier(img, top_k=top_k)
+    return {r["label"]: float(r["score"]) for r in results}
+
+# Gradio interface
+demo = gr.Interface(
+    fn=classify_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Slider(1, 5, value=3, step=1, label="Top K Predictions")
+    ],
+    outputs=gr.Label(num_top_classes=5, label="Predictions"),
+    title="Image Classification App",
+    description="Upload an image and the model will predict the top objects in it."
+)
+
+if __name__ == "__main__":
+    demo.launch()