Update app.py

app.py

@@ -37,14 +37,16 @@ def apply_mask(image: Image.Image, prompt: str, color: str) -> Image.Image:
 
     # Get the binary mask from predictions
     mask = preds.sigmoid().detach().cpu().numpy()
-    mask = (mask > 0.5).astype(np.uint8)
+    mask = cv2.resize(mask[0], (image_np.shape[1], image_np.shape[0]))
+    mask = cv2.GaussianBlur(mask, (15, 15), 0)
+    mask_bin = (mask > 0.4).astype(np.uint8)
 
     # Convert image to RGBA
     image_np = np.array(image.convert("RGBA"))
 
     # Resize mask to match image size
     mask_resized = cv2.resize(mask, (image_np.shape[1], image_np.shape[0]))
-    mask_3d = np.stack([mask_resized] * 4, axis=-1)  # Extend mask to 3D
+    mask_3d = np.stack([mask_bin] * 4, axis=-1)  # Extend mask to 3D
 
     # Convert the color string to an RGB tuple
     color_rgb = parse_color(color)
@@ -52,8 +54,11 @@ def apply_mask(image: Image.Image, prompt: str, color: str) -> Image.Image:
 
     # Create an overlay with the selected color
     overlay = np.zeros_like(image_np, dtype=np.uint8)
-    overlay[:] = overlay_color
-
+    # overlay[:] = overlay_color
+    masked_image = image_np.copy()
+    masked_image[mask_bin == 1] = (
+        0.5 * masked_image[mask_bin == 1] + 0.5 * overlay[mask_bin == 1]).astype(np.units)
+    )
     # Apply the mask to the image
     masked_image = np.where(mask_3d == 1, overlay, image_np)
     return Image.fromarray(masked_image)