Update app.py

app.py

@@ -13,53 +13,63 @@ 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."
+    try:
+        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).convert("RGB")
+        
+        # Store original size
+        original_size = input_image.size
+        if not original_size or 0 in original_size:
+            return None, "Invalid image size. Please upload a different image."
+        
+        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
+            width, height = original_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():
+            outputs = model(**inputs)
+        
+        # Post-process masks
+        masks = processor.image_processor.post_process_masks(
+            outputs.pred_masks.cpu(),
+            inputs["original_sizes"].cpu(),
+            inputs["reshaped_input_sizes"].cpu()
+        )
+        
+        # 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)
+        else:
+            # Use the first mask
+            combined_mask = masks[0][0].numpy() > 0.5
+        
+        # Ensure mask is 2D
+        if combined_mask.ndim > 2:
+            combined_mask = combined_mask.squeeze()
+        
+        # Resize mask to match original image size
+        combined_mask = cv2.resize(combined_mask.astype(np.uint8), (original_size[0], 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, "Segmentation completed successfully."
     
-    # 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
-        width, height = 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():
-        outputs = model(**inputs)
-    
-    # Post-process masks
-    masks = processor.image_processor.post_process_masks(
-        outputs.pred_masks.cpu(),
-        inputs["original_sizes"].cpu(),
-        inputs["reshaped_input_sizes"].cpu()
-    )
-    
-    # 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)
-    else:
-        # 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, "Segmentation completed successfully."
+    except Exception as e:
+        return None, f"An error occurred: {str(e)}"
 
 # Create Gradio interface
 iface = gr.Interface(