add example images for segmentation

app.py

@@ -50,19 +50,27 @@ def segment(image_files):
 
 # components for Gradion interface
 def update_gallery(images):
-    print(f"Type in update_gallery: {type(images[0])}")
-    gallery_data = []
+    print(f"Images received: {images}")
 
+    gallery_data = []
     if images:
-        segmented_images = segment(images)  # Process all images in one batch
+        segmented_images = segment(images)  # Process images
+
+        for i, image_path in enumerate(images):
+            try:
+                image = Image.open(image_path).convert("RGB")  # Load original image
+                
+                segmented_mask = segmented_images[i].to(dtype=torch.float32, device="cpu")
 
-        for i, image in enumerate(images):
-            segmented_image_pil = transforms.ToPILImage()(segmented_images[i].to(device, dtype=torch.float32))
-            gallery_data.extend([(image, "Original Image"), (segmented_image_pil, "Segmented Image")])
+                segmented_image_pil = transforms.ToPILImage()(segmented_mask)  # Convert to PIL Image
+                
+                gallery_data.extend([(image, "Original Image"), (segmented_image_pil, "Segmented Image")])
+            except Exception as e:
+                print(f"Error processing image {i}: {e}")
+                gallery_data.extend([(image, "Original Image"), (image, f"Error: {e}")])
 
     return gallery_data
 
-# Gradio UI for MEI segmentation
 import gradio as gr
 
 with gr.Blocks() as demo:
@@ -80,5 +88,21 @@ with gr.Blocks() as demo:
              fn=update_gallery,
              inputs=[image_files],
              outputs=[gallery])
+        
+      with gr.Column():
+        example_image = gr.Image(type="filepath", label="MRI Image", visible=False)
+        examples = gr.Examples(examples=["examples/Te-me_0194.jpg", "examples/Te-me_0111.jpg", 
+                                         "examples/Te-me_0295.jpg", "examples/Te-me_0228.jpg",
+                                         "examples/Te-me_0218.jpg", "examples/Te-me_0164.jpg"], 
+                               inputs=[example_image])
+
+      with gr.Column(scale=0):
+
+        example_button = gr.Button("Process Example Image", variant="secondary")
+        example_button.click(
+            fn=lambda img: update_gallery([img]) if img else [],
+            inputs=[example_image],  
+            outputs=[gallery]
+        )
 
 demo.launch()