Update app.py

app.py

@@ -17,7 +17,7 @@ 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
+        # Convert input_image to PIL Image and ensure it's RGB
         input_image = Image.fromarray(input_image).convert("RGB")
         
         # Store original size
@@ -25,11 +25,10 @@ def segment_image(input_image, segment_anything):
         if not original_size or 0 in original_size:
             return None, "Invalid image size. Please upload a different image."
         
+        # Process the 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)
@@ -45,20 +44,20 @@ def segment_image(input_image, segment_anything):
             inputs["reshaped_input_sizes"].cpu()
         )
         
-        # Convert mask to numpy array and resize to match original image
+        # Convert mask to numpy array
         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
+        # Resize mask to match original image size using PIL
+        mask_image = Image.fromarray((combined_mask * 255).astype(np.uint8))
+        mask_image = mask_image.resize(original_size, Image.NEAREST)
+        combined_mask = np.array(mask_image) > 0
         
         # Overlay the mask on the original image
         result_image = np.array(input_image)