Update app.py

app.py

@@ -30,9 +30,6 @@ def parse_color(color_str):
     raise ValueError(f"Invalid color format: {color_str}. Use hex like '#ff0000', color name like 'red', or rgba format.")
 
 def apply_mask(image: Image.Image, prompt: str, color: str) -> Image.Image:
-        # Convert image to RGBA
-    image_np = np.array(image.convert("RGBA"))
-    
     # Process the input image and prompt
     inputs = processor(text=prompt, images=image, return_tensors="pt")
     outputs = model(**inputs)
@@ -40,14 +37,14 @@ 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 = 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)
+    mask = (mask > 0.5).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_bin] * 4, axis=-1)  # Extend mask to 3D
+    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
 
     # Convert the color string to an RGB tuple
     color_rgb = parse_color(color)
@@ -56,11 +53,7 @@ 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
-    
-    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)
@@ -78,4 +71,4 @@ iface = gr.Interface(
     description="Upload an image, input a prompt (e.g., 'person', 'sky'), and pick a mask color."
 )
 
-iface.launch()
+iface.launch()