Update app.py

app.py

@@ -29,14 +29,14 @@ css = """
 }
 """
 
-def predict(img):
+def predict(img_path):
     """Process single image and return prediction"""
-    if img is None:
-        return "Please select or upload an image"
+    if not img_path:
+        return "Please select or upload an image first"
     
     try:
-        image = Image.open(img).convert('RGB')
-        tensor = preprocess(image).unsqueeze(0)  # Add batch dimension
+        image = Image.open(img_path).convert('RGB')
+        tensor = preprocess(image).unsqueeze(0)
         
         with torch.inference_mode():
             outputs = model(tensor)
@@ -49,38 +49,51 @@ def predict(img):
 
 with gr.Blocks(title="Animal Classifier", css=css) as demo:
     gr.Markdown("## 🐾 Animal Classifier")
-    gr.Markdown("Click example images below or upload your own")
+    gr.Markdown("Select an image below or upload your own, then click Classify")
+    
+    # Store current image path
+    current_image = gr.State()
     
     with gr.Row():
-        input_image = gr.Image(type="filepath", label="Selected Image")
-        output_label = gr.Label(label="Prediction")
+        with gr.Column():
+            image_preview = gr.Image(label="Selected Image", type="filepath")
+            upload_btn = gr.UploadButton("Upload Custom Image", file_types=["image"])
+            classify_btn = gr.Button("Classify 🚀", variant="primary")
+        result = gr.Textbox(label="Prediction", lines=3)
     
-    # Example gallery with click handling
+    # Example gallery at bottom
     with gr.Row(variant="panel"):
         examples_gallery = gr.Gallery(
             value=example_images,
-            label="Example Images (Click to Predict)",
+            label="Example Images (Click to Select)",
             columns=7,
             height=120,
             allow_preview=False,
             elem_classes=["centered-examples"]
         )
-    
-    # Handle example image clicks
+
+    # Handle image selection from examples
     def select_example(evt: gr.SelectData):
         return example_images[evt.index]
     
     examples_gallery.select(
         fn=select_example,
-        outputs=input_image,
+        outputs=[image_preview, current_image],
         show_progress=False
     )
-    
-    # Handle predictions for both upload and example clicks
-    input_image.change(
+
+    # Handle custom uploads
+    upload_btn.upload(
+        fn=lambda file: (file.name, file.name),
+        inputs=upload_btn,
+        outputs=[image_preview, current_image]
+    )
+
+    # Handle classification
+    classify_btn.click(
         fn=predict,
-        inputs=input_image,
-        outputs=output_label
+        inputs=current_image,
+        outputs=result
     )
 
 if __name__ == "__main__":