Update app.py

app.py

@@ -12,12 +12,18 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
 processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
-def segment_image(input_image, x, y):
+def segment_image(input_image, segment_anything):
     # Convert input_image to PIL Image
     input_image = Image.fromarray(input_image)
     
-    # Prepare inputs
-    inputs = processor(input_image, input_points=np.array([[x, y]]), return_tensors="pt").to(device)
+    if segment_anything:
+        # Segment everything in the image
+        inputs = processor(input_image, return_tensors="pt").to(device)
+    else:
+        # Use the center of the image as a point prompt
+        height, width = input_image.size
+        center_point = [[width // 2, height // 2]]
+        inputs = processor(input_image, input_points=[center_point], return_tensors="pt").to(device)
     
     # Generate masks
     with torch.no_grad():
@@ -29,15 +35,19 @@ def segment_image(input_image, x, y):
         inputs["original_sizes"].cpu(),
         inputs["reshaped_input_sizes"].cpu()
     )
-    scores = outputs.iou_scores
     
     # Convert mask to numpy array
-    mask = masks[0][0].numpy()
+    if segment_anything:
+        # Combine all masks
+        combined_mask = np.any(masks[0].numpy() > 0.5, axis=0)
+    else:
+        # Use the first mask
+        combined_mask = masks[0][0].numpy() > 0.5
     
     # Overlay the mask on the original image
     result_image = np.array(input_image)
     mask_rgb = np.zeros_like(result_image)
-    mask_rgb[mask > 0.5] = [255, 0, 0]  # Red color for the mask
+    mask_rgb[combined_mask] = [255, 0, 0]  # Red color for the mask
     result_image = cv2.addWeighted(result_image, 1, mask_rgb, 0.5, 0)
     
     return result_image
@@ -47,12 +57,11 @@ iface = gr.Interface(
     fn=segment_image,
     inputs=[
         gr.Image(type="numpy"),
-        gr.Slider(minimum=0, maximum=1000, step=1, label="X coordinate"),
-        gr.Slider(minimum=0, maximum=1000, step=1, label="Y coordinate")
+        gr.Checkbox(label="Segment Everything")
     ],
     outputs=gr.Image(type="numpy"),
     title="Segment Anything Model (SAM) Image Segmentation",
-    description="Enter X and Y coordinates of the object you want to segment."
+    description="Upload an image and choose whether to segment everything or use a center point."
 )
 
 # Launch the interface