Update app.py

app.py

@@ -18,11 +18,13 @@ 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)
@@ -32,21 +34,30 @@ def soft_alpha(mask_uint8, blur_radius=10):
     return (blurred.astype(np.float32) / 255.0).clip(0.0, 1.0)
 
 def isolate_with_click(image: Image.Image, evt: gr.SelectData):
+    """
+    image: PIL image uploaded
+    evt: click coordinates from Gradio (x,y)
+    """
     img_rgb = np.array(image.convert("RGB"))
     predictor.set_image(img_rgb)
 
+    # SAM expects input points as numpy array [[x,y]]
     input_point = np.array([[evt.index[0], evt.index[1]]])
-    input_label = np.array([1])
+    input_label = np.array([1])  # 1 = foreground
 
-    masks, scores, _ = predictor.predict(
+    masks, scores, logits = predictor.predict(
         point_coords=input_point,
         point_labels=input_label,
         multimask_output=True
     )
 
+    # Pick the highest score mask
     best_mask = masks[np.argmax(scores)].astype(np.uint8) * 255
+
+    # Soft alpha
     alpha = soft_alpha(best_mask, blur_radius=BLUR_RADIUS)
 
+    # Crop to bounding box
     ys, xs = np.where(best_mask == 255)
     if len(xs) == 0 or len(ys) == 0:
         return None
@@ -65,17 +76,15 @@ def isolate_with_click(image: Image.Image, evt: gr.SelectData):
 # --- 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 = gr.Image(type="pil", label="Isolated cutout (RGBA)")
-
     inp.select(isolate_with_click, inputs=[inp], outputs=out)
 
-    # Examples section at the bottom
+    # Add demo example at the bottom
     gr.Examples(
-        examples=["demo.png"],  
+        examples=["demo.png"],   # make sure demo.png is in your repo
         inputs=inp,
         label="Try with demo image"
     )
 
-demo.launch()
+demo.launch(share=True)