Update app.py

app.py

@@ -52,8 +52,11 @@ def segment_image(input_image, text_input):
     image_bytes = base64.b64decode(input_image)
     image = Image.open(BytesIO(image_bytes))
 
-    # Convert the image to a NumPy array
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)  # Convert to RGB color space
+    # Convert the image to RGB color mode
+    image = image.convert("RGB")
+
+    # Convert the image to a numpy array
+    image = np.array(image)
 
     mask_generator = SamAutomaticMaskGenerator(sam)
     masks = mask_generator.generate(image)
@@ -62,13 +65,19 @@ def segment_image(input_image, text_input):
 
     for i, mask_dict in enumerate(masks):
         mask_data = (mask_dict['segmentation'] * 255).astype(np.uint8)
-        segmented_region = cv2.bitwise_and(image, image, mask=mask_data)
+
+        # Create a mask with the same shape as the original image
+        mask = np.zeros_like(image)
+        mask[:, :] = mask_data[:, :, np.newaxis]
+
+        # Apply the mask to the original image
+        segmented_region = cv2.bitwise_and(image, mask)
 
         x, y, w, h = map(int, mask_dict['bbox'])
         cropped_region = segmented_region[y:y+h, x:x+w]
 
         # Convert to base64 image
-        _, buffer = cv2.imencode(".png", cv2.cvtColor(cropped_region, cv2.COLOR_RGB2BGR))  # Convert back to BGR for encoding
+        _, buffer = cv2.imencode(".png", cropped_region)
         segmented_image_base64 = base64.b64encode(buffer).decode()
 
         # Calculate similarity for the segmented image