Update app.py

app.py

@@ -26,92 +26,61 @@ css = """
     height: 100px !important;
     width: 100px !important;
     object-fit: cover !important;
-}"""
+}
+"""
 
+def predict(img):
+    """Process single image and return prediction"""
+    if img is None:
+        return "Please select or upload an image"
+    
+    try:
+        image = Image.open(img).convert('RGB')
+        tensor = preprocess(image).unsqueeze(0)  # Add batch dimension
+        
+        with torch.inference_mode():
+            outputs = model(tensor)
+            _, pred = torch.max(outputs, 1)
+        
+        return classes[pred.item()]
+    
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 with gr.Blocks(title="Animal Classifier", css=css) as demo:
     gr.Markdown("## 🐾 Animal Classifier")
-    gr.Markdown("Upload multiple animal images to get predictions!")
-    
-    # Store uploaded and example file paths
-    all_files_state = gr.State([])
-    
-    with gr.Row():
-        inputs = gr.File(file_count="multiple", file_types=["image"], label="Upload Animal Images")
-        submit = gr.Button("Classify 🚀", variant="primary")
+    gr.Markdown("Click example images below or upload your own")
     
     with gr.Row():
-        gallery = gr.Gallery(label="Upload Preview", columns=4)
-        outputs = gr.Textbox(label="Predictions", lines=5)
+        input_image = gr.Image(type="filepath", label="Selected Image")
+        output_label = gr.Label(label="Prediction")
     
     # Example gallery with click handling
     with gr.Row(variant="panel"):
         examples_gallery = gr.Gallery(
             value=example_images,
-            label="Example Images (Click to Add)",
+            label="Example Images (Click to Predict)",
             columns=7,
-            height=244,
+            height=120,
             allow_preview=False,
             elem_classes=["centered-examples"]
         )
     
-    # Update state when files are uploaded
-    def update_state(new_files):
-        return [f.name for f in new_files] if new_files else []
+    # Handle example image clicks
+    def select_example(evt: gr.SelectData):
+        return example_images[evt.index]
     
-    inputs.change(update_state, inputs, all_files_state)
-    
-    # Handle example selection - FIXED WITH VALIDATION
-    def add_example(example_index, current_files):
-        try:
-            # Handle Gradio's selection format
-            if isinstance(example_index, list):
-                if not example_index:  # Empty selection
-                    return current_files
-                selected_idx = example_index[0]
-            else:
-                selected_idx = example_index
-            
-            # Validate index range
-            if 0 <= selected_idx < len(example_images):
-                selected_path = example_images[selected_idx]
-                return current_files + [selected_path]
-            return current_files
-        except Exception as e:
-            print(f"Error handling example selection: {str(e)}")
-            return current_files
-    
-    # Corrected gallery select handler
     examples_gallery.select(
-        add_example,
-        [all_files_state],
-        all_files_state,
+        fn=select_example,
+        outputs=input_image,
         show_progress=False
     )
     
-    # Update gallery preview
-    def update_gallery(files):
-        return files if files else []
-    
-    all_files_state.change(update_gallery, all_files_state, gallery)
-    
-    # Modified prediction function
-    def predict(files):
-        if not files:
-            return ""
-        try:
-            batch = torch.stack([preprocess(Image.open(img).convert('RGB')) for img in files])
-            with torch.inference_mode():
-                outputs = model(batch)
-                _, preds = torch.max(outputs, 1)
-            return ", ".join([classes[p] for p in preds.cpu().numpy()])
-        except Exception as e:
-            return f"Error: {str(e)}"
-    
-    submit.click(
+    # Handle predictions for both upload and example clicks
+    input_image.change(
         fn=predict,
-        inputs=all_files_state,
-        outputs=outputs
+        inputs=input_image,
+        outputs=output_label
     )
 
 if __name__ == "__main__":