Spaces:
Sleeping
Sleeping
| import os | |
| import subprocess | |
| # --- Ensure SAM checkpoint is present --- | |
| SAM_CHECKPOINT = "sam_vit_h_4b8939.pth" | |
| if not os.path.exists(SAM_CHECKPOINT): | |
| print("Downloading SAM checkpoint...") | |
| subprocess.run([ | |
| "wget", | |
| "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth" | |
| ]) | |
| # ---------------------------------------- | |
| 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 | |
| # -------------- | |
| # Load SAM once | |
| 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]) # foreground | |
| 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, 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))) | |
| 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) | |
| 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.") | |
| 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") | |
| inp.select(isolate_with_click, inputs=[inp], outputs=[out_cutout, out_overlay]) | |
| # Demo example at the bottom | |
| gr.Examples( | |
| examples=["demo.png"], # make sure demo.png is in your repo | |
| inputs=inp, | |
| label="Try with demo image" | |
| ) | |
| demo.launch(share=True) | |