Create app.py

app.py

@@ -0,0 +1,65 @@
+import gradio as gr
+import numpy as np
+from PIL import Image
+import cv2
+import torch
+from segment_anything import sam_model_registry, SamPredictor
+
+# --- CONFIG ---
+SAM_CHECKPOINT = "sam_vit_h_4b8939.pth"
+SAM_MODEL_TYPE = "vit_h"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+BLUR_RADIUS = 10
+DEMO_IMAGE = "demo.png"   # put a demo image in the repo
+# --------------
+
+sam = sam_model_registry[SAM_MODEL_TYPE](checkpoint=SAM_CHECKPOINT)
+sam.to(device=DEVICE)
+predictor = SamPredictor(sam)
+
+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 isolate_with_click(image: Image.Image, evt: gr.SelectData):
+    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])
+
+    masks, scores, _ = predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        multimask_output=True
+    )
+
+    best_mask = masks[np.argmax(scores)].astype(np.uint8) * 255
+    alpha = soft_alpha(best_mask, blur_radius=BLUR_RADIUS)
+
+    ys, xs = np.where(best_mask == 255)
+    if len(xs) == 0 or len(ys) == 0:
+        return None
+    x0, x1 = xs.min(), xs.max()
+    y0, y1 = ys.min(), ys.max()
+    pad = int(max(img_rgb.shape[:2]) * 0.02)
+    x0 = max(0, x0 - pad); x1 = min(img_rgb.shape[1]-1, x1 + pad)
+    y0 = max(0, y0 - pad); y1 = min(img_rgb.shape[0]-1, y1 + pad)
+
+    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)))
+    return Image.fromarray(rgba)
+
+with gr.Blocks() as demo:
+    gr.Markdown("### SAM Object Isolation\nUpload an image or use the demo below, then click on the object to isolate it.")
+
+    inp = gr.Image(type="pil", label="Upload image", interactive=True)
+    demo_img = gr.Image(value=DEMO_IMAGE, type="pil", label="Demo image (click to try)", interactive=True)
+    out = gr.Image(type="pil", label="Isolated cutout (RGBA)")
+
+    inp.select(isolate_with_click, inputs=[inp], outputs=out)
+    demo_img.select(isolate_with_click, inputs=[demo_img], outputs=out)
+
+demo.launch()