Update app.py

app.py

@@ -13,15 +13,21 @@ model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
 processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
 def segment_image(input_image, segment_anything):
+    if input_image is None:
+        return None, "Please upload an image before submitting."
+    
     # Convert input_image to PIL Image
     input_image = Image.fromarray(input_image)
     
+    # Store original size
+    original_size = input_image.size
+    
     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
+        width, height = input_image.size
         center_point = [[width // 2, height // 2]]
         inputs = processor(input_image, input_points=[center_point], return_tensors="pt").to(device)
     
@@ -36,7 +42,7 @@ def segment_image(input_image, segment_anything):
         inputs["reshaped_input_sizes"].cpu()
     )
     
-    # Convert mask to numpy array
+    # Convert mask to numpy array and resize to match original image
     if segment_anything:
         # Combine all masks
         combined_mask = np.any(masks[0].numpy() > 0.5, axis=0)
@@ -44,22 +50,28 @@ def segment_image(input_image, segment_anything):
         # Use the first mask
         combined_mask = masks[0][0].numpy() > 0.5
     
+    # Resize mask to match original image size
+    combined_mask = cv2.resize(combined_mask.astype(np.uint8), original_size[::-1]) > 0
+    
     # Overlay the mask on the original image
     result_image = np.array(input_image)
     mask_rgb = np.zeros_like(result_image)
     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
+    return result_image, "Segmentation completed successfully."
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=segment_image,
     inputs=[
-        gr.Image(type="numpy"),
+        gr.Image(type="numpy", label="Upload an image"),
         gr.Checkbox(label="Segment Everything")
     ],
-    outputs=gr.Image(type="numpy"),
+    outputs=[
+        gr.Image(type="numpy", label="Segmented Image"),
+        gr.Textbox(label="Status")
+    ],
     title="Segment Anything Model (SAM) Image Segmentation",
     description="Upload an image and choose whether to segment everything or use a center point."
 )