Update app.py

app.py

@@ -33,12 +33,37 @@ def soft_alpha(mask_uint8, blur_radius=10):
     blurred = cv2.GaussianBlur(mask_uint8, (0,0), sigmaX=blur_radius, sigmaY=blur_radius)
     return (blurred.astype(np.float32) / 255.0).clip(0.0, 1.0)
 
+def make_overlay(img_rgb: np.ndarray, mask_uint8: np.ndarray) -> Image.Image:
+    """Return a purple-tinted overlay on the selected region."""
+    # Ensure 0/255 mask and uint8 image
+    mask = (mask_uint8 > 127).astype(np.uint8)
+    overlay = img_rgb.copy()
+
+    # Create a purple layer
+    purple = np.zeros_like(img_rgb, dtype=np.uint8)
+    purple[..., 0] = 180  # R
+    purple[..., 1] = 0    # G
+    purple[..., 2] = 180  # B
+
+    # Blend only on selected pixels
+    sel = mask.astype(bool)
+    blended = cv2.addWeighted(overlay[sel], 0.6, purple[sel], 0.4, 0)
+    overlay[sel] = blended
+
+    return Image.fromarray(overlay)
+
 def isolate_with_click(image: Image.Image, evt: gr.SelectData):
+    # Guard: if no image or no click event, return original image and a subtle info overlay
+    if image is None or evt is None:
+        return None, None
+
     img_rgb = np.array(image.convert("RGB"))
     predictor.set_image(img_rgb)
 
-    input_point = np.array([[evt.index[0], evt.index[1]]])
-    input_label = np.array([1])  # foreground
+    # SAM expects input points as numpy array [[x,y]]
+    x, y = evt.index  # (x, y) from Gradio click
+    input_point = np.array([[x, y]], dtype=np.float32)
+    input_label = np.array([1], dtype=np.int32)  # 1 = foreground
 
     masks, scores, _ = predictor.predict(
         point_coords=input_point,
@@ -46,12 +71,28 @@ def isolate_with_click(image: Image.Image, evt: gr.SelectData):
         multimask_output=True
     )
 
-    best_mask = masks[np.argmax(scores)].astype(np.uint8) * 255
+    # If SAM didn't return masks, show original plus a faint marker
+    if masks is None or len(masks) == 0:
+        # Create a simple marker overlay to indicate click
+        overlay = img_rgb.copy()
+        cv2.circle(overlay, (int(x), int(y)), 12, (180, 0, 180), thickness=3)
+        return None, Image.fromarray(overlay)
+
+    # Pick the highest score mask
+    best_idx = int(np.argmax(scores))
+    best_mask = masks[best_idx].astype(np.uint8) * 255
+
+    # Soft alpha for RGBA cutout
     alpha = soft_alpha(best_mask, blur_radius=BLUR_RADIUS)
 
+    # Crop to bounding box (with small pad)
     ys, xs = np.where(best_mask == 255)
     if len(xs) == 0 or len(ys) == 0:
-        return None, None
+        # If mask is empty, still return overlay with click marker
+        overlay = img_rgb.copy()
+        cv2.circle(overlay, (int(x), int(y)), 12, (180, 0, 180), thickness=3)
+        return None, Image.fromarray(overlay)
+
     x0, x1 = xs.min(), xs.max()
     y0, y1 = ys.min(), ys.max()
     pad = int(max(img_rgb.shape[:2]) * 0.02)
@@ -61,30 +102,28 @@ def isolate_with_click(image: Image.Image, evt: gr.SelectData):
     fg_rgb = img_rgb[y0:y1+1, x0:x1+1]
     fg_alpha = alpha[y0:y1+1, x0:x1+1]
 
-    rgba = np.dstack((fg_rgb, (fg_alpha * 255).astype(np.uint8)))
+    # Compose RGBA correctly once
+    rgba = np.dstack((fg_rgb, (fg_alpha * 255.0).astype(np.uint8)))
     cutout = Image.fromarray(rgba)
 
-    # Build overlay preview (purple tint on mask)
-    overlay = img_rgb.copy()
-    tint = np.array([180, 0, 180], dtype=np.uint8)  # purple
-    sel = best_mask == 255
-    overlay[sel] = (0.6 * overlay[sel] + 0.4 * tint).astype(np.uint8)
-    overlay_img = Image.fromarray(overlay)
+    # Build and return the purple overlay on the original image
+    overlay_img = make_overlay(img_rgb, best_mask)
 
     return cutout, overlay_img
 
 # --- Gradio UI ---
 with gr.Blocks() as demo:
-    gr.Markdown("### SAM Object Isolation\nUpload an image, then click on the object to isolate it.")
+    gr.Markdown("### SAM Object Isolation\nUpload an image (or pick the demo), then click on the object to isolate it. The right panel shows a purple overlay of the mask.")
     inp = gr.Image(type="pil", label="Upload image", interactive=True)
     out_cutout = gr.Image(type="pil", label="Isolated cutout (RGBA)")
     out_overlay = gr.Image(type="pil", label="Segmentation overlay preview")
 
+    # Click-to-segment
     inp.select(isolate_with_click, inputs=[inp], outputs=[out_cutout, out_overlay])
 
-    # Demo example at the bottom
+    # Demo example at the bottom that populates the upload image
     gr.Examples(
-        examples=["demo.png"],   # make sure demo.png is in your repo
+        examples=["demo.png"],   # ensure demo.png is in your repo
         inputs=inp,
         label="Try with demo image"
     )