Update app.py

app.py

@@ -4,6 +4,7 @@ from PIL import Image
 import numpy as np
 import matplotlib.pyplot as plt
 import gradio as gr
+import os
 
 # Load pretrained Mask R-CNN model
 model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
@@ -26,8 +27,24 @@ COCO_INSTANCE_CATEGORY_NAMES = [
     'hair drier', 'toothbrush'
 ]
 
+# Function to load example image or use uploaded image
+def load_image(uploaded_image, example_image):
+    if uploaded_image is not None:
+        return uploaded_image
+    elif example_image and example_image != "None":
+        example_path = os.path.join("examples", example_image)
+        if os.path.exists(example_path):
+            return Image.open(example_path).convert("RGB")
+        else:
+            raise FileNotFoundError(f"Example image {example_path} not found.")
+    else:
+        raise ValueError("Please upload an image or select an example image.")
+
 # Detection and segmentation function
-def segment_objects(image, threshold=0.5):
+def segment_objects(uploaded_image, example_image, threshold=0.5):
+    # Load the image (either uploaded or example)
+    image = load_image(uploaded_image, example_image)
+    
     transform = torchvision.transforms.ToTensor()
     img_tensor = transform(image).unsqueeze(0)
 
@@ -77,11 +94,16 @@ interface = gr.Interface(
     fn=segment_objects,
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
+        gr.Dropdown(
+            choices=["None", "example1.jpg", "example2.jpg", "example3.jpg", "example4.jpg"],
+            value="None",
+            label="Select Example Image"
+        ),
         gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Confidence Threshold")
     ],
     outputs=gr.Image(type="filepath", label="Segmented Output"),
     title="Mask R-CNN Instance Segmentation",
-    description="Upload an image to detect and segment objects using a pretrained Mask R-CNN model (TorchVision)."
+    description="Upload an image or select an example image to detect and segment objects using a pretrained Mask R-CNN model (TorchVision)."
 )
 
 if __name__ == "__main__":