release

app-dev.py

@@ -0,0 +1,251 @@
+import os
+import logging
+import hashlib
+import sys
+import traceback
+import copy
+import tempfile
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+import gradio as gr
+from PIL import Image, ImageFilter, ImageChops, ImageDraw
+from huggingface_hub import hf_hub_download
+import spaces
+
+# --- IMPORT YOUR CUSTOM MODULES ---
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+from plm_adapter_lora_with_image_input_only_text_positions import PLMLanguageAdapter
+
+# ----------------- Configuration -----------------
+SAM2_CONFIG = "sam2_hiera_l.yaml"
+BASE_CKPT_NAME = "sam2_hiera_large.pt"
+
+SQUARE_DIM = 1024
+logging.basicConfig(level=logging.INFO)
+
+# Refactored to store specific filenames per model choice
+MODEL_CONFIGS = {
+    "Stage 1": {
+        "repo_id": "aadarsh99/ConvSeg-Stage1",
+        "sam_filename": "fine_tuned_sam2_batched_100000.torch",
+        "plm_filename": "fine_tuned_sam2_batched_plm_100000.torch"
+    },
+    "Stage 2 (grad-acc: 4)": {
+        "repo_id": "aadarsh99/ConvSeg-Stage2",
+        "sam_filename": "fine_tuned_sam2_batched_18000.torch",
+        "plm_filename": "fine_tuned_sam2_batched_plm_18000.torch"
+    },
+    "Stage 2 (grad-acc: 8)": {
+        "repo_id": "aadarsh99/ConvSeg-Stage2",
+        "sam_filename": "fine_tuned_sam2_batched_18000.torch",
+        "plm_filename": "fine_tuned_sam2_batched_plm_18000.torch"
+    }
+}
+
+# Dynamically create cache keys based on config
+MODEL_CACHE = {k: {"sam": None, "plm": None} for k in MODEL_CONFIGS.keys()}
+
+# ----------------- Helper Functions -----------------
+def download_if_needed(repo_id, filename):
+    try:
+        logging.info(f"Checking {filename} in {repo_id}...")
+        return hf_hub_download(repo_id=repo_id, filename=filename)
+    except Exception as e:
+        raise FileNotFoundError(f"Could not find {filename} in {repo_id}. Error: {e}")
+
+def stable_color(key: str):
+    h = int(hashlib.sha256(str(key).encode("utf-8")).hexdigest(), 16)
+    EDGE_COLORS_HEX = ["#3A86FF", "#FF006E", "#43AA8B", "#F3722C", "#8338EC", "#90BE6D"]
+    colors = [tuple(int(c.lstrip("#")[i:i+2], 16) for i in (0, 2, 4)) for c in EDGE_COLORS_HEX]
+    return colors[h % len(colors)]
+
+def make_overlay(rgb: np.ndarray, mask: np.ndarray, key: str = "mask") -> Image.Image:
+    # Convert base to RGBA
+    base = Image.fromarray(rgb.astype(np.uint8)).convert("RGBA")
+    mask_bool = mask > 0
+    color = stable_color(key)
+    
+    # Create fill layer (Semi-transparent)
+    fill_layer = Image.new("RGBA", base.size, color + (0,))
+    fill_alpha = Image.fromarray((mask_bool.astype(np.uint8) * 140), "L")
+    fill_layer.putalpha(fill_alpha)
+
+    # Create stroke/edge layer
+    m = Image.fromarray((mask_bool.astype(np.uint8) * 255), "L")
+    edges = ImageChops.difference(m.filter(ImageFilter.MaxFilter(3)), m.filter(ImageFilter.MinFilter(3)))
+    stroke_layer = Image.new("RGBA", base.size, color + (255,))
+    stroke_layer.putalpha(edges)
+
+    # Composite safely
+    out = Image.alpha_composite(base, fill_layer)
+    out = Image.alpha_composite(out, stroke_layer)
+    
+    return out.convert("RGB")
+
+def ensure_models_loaded(stage_key):
+    global MODEL_CACHE
+    if MODEL_CACHE[stage_key]["sam"] is not None: 
+        return
+    
+    config = MODEL_CONFIGS[stage_key]
+    repo_id = config["repo_id"]
+    
+    logging.info(f"Loading {stage_key} models from {repo_id} into CPU RAM...")
+    
+    # SAM2
+    # Base model is always the same
+    base_path = download_if_needed(repo_id, BASE_CKPT_NAME)
+    model = build_sam2(SAM2_CONFIG, base_path, device="cpu")
+    
+    # Load specific fine-tuned checkpoint
+    final_path = download_if_needed(repo_id, config["sam_filename"])
+    sd = torch.load(final_path, map_location="cpu")
+    model.load_state_dict(sd.get("model", sd), strict=True)
+    
+    # PLM
+    plm_path = download_if_needed(repo_id, config["plm_filename"])
+    plm = PLMLanguageAdapter(
+        model_name="Qwen/Qwen2.5-VL-3B-Instruct", 
+        transformer_dim=model.sam_mask_decoder.transformer_dim, 
+        n_sparse_tokens=0, use_dense_bias=True, use_lora=True, 
+        lora_r=16, lora_alpha=32, lora_dropout=0.05, 
+        dtype=torch.bfloat16, device="cpu"
+    )
+    plm_sd = torch.load(plm_path, map_location="cpu")
+    plm.load_state_dict(plm_sd["plm"], strict=True)
+    plm.eval()
+    
+    MODEL_CACHE[stage_key]["sam"], MODEL_CACHE[stage_key]["plm"] = model, plm
+
+# ----------------- GPU Inference -----------------
+
+@spaces.GPU(duration=120) 
+def run_prediction(image_pil, text_prompt, threshold, stage_choice):
+    if image_pil is None or not text_prompt:
+        return None, None, None
+
+    ensure_models_loaded(stage_choice)
+    sam_model = MODEL_CACHE[stage_choice]["sam"]
+    plm_model = MODEL_CACHE[stage_choice]["plm"]
+    
+    sam_model.to("cuda")
+    plm_model.to("cuda")
+    
+    try:
+        with torch.inference_mode():
+            predictor = SAM2ImagePredictor(sam_model)
+            rgb_orig = np.array(image_pil.convert("RGB"))
+            H, W = rgb_orig.shape[:2]
+            
+            # Padding math
+            scale = SQUARE_DIM / max(H, W)
+            nw, nh = int(W * scale), int(H * scale)
+            top, left = (SQUARE_DIM - nh) // 2, (SQUARE_DIM - nw) // 2
+
+            # Resize & Pad
+            rgb_sq = cv2.resize(rgb_orig, (nw, nh), interpolation=cv2.INTER_LINEAR)
+            rgb_sq = cv2.copyMakeBorder(rgb_sq, top, SQUARE_DIM-nh-top, left, SQUARE_DIM-nw-left, cv2.BORDER_CONSTANT, value=0)
+            
+            predictor.set_image(rgb_sq)
+            image_emb = predictor._features["image_embed"][-1].unsqueeze(0)
+            hi = [lvl[-1].unsqueeze(0) for lvl in predictor._features["high_res_feats"]]
+
+            # PLM adapter
+            with tempfile.NamedTemporaryFile(suffix=".jpg") as tmp:
+                image_pil.save(tmp.name)
+                sp, dp = plm_model([text_prompt], image_emb.shape[2], image_emb.shape[3], [tmp.name])
+
+            # SAM2 Decoding
+            dec = sam_model.sam_mask_decoder
+            dev, dtype = next(dec.parameters()).device, next(dec.parameters()).dtype
+            
+            low, scores, _, _ = dec(
+                image_embeddings=image_emb.to(dev, dtype),
+                image_pe=sam_model.sam_prompt_encoder.get_dense_pe().to(dev, dtype),
+                sparse_prompt_embeddings=sp.to(dev, dtype),
+                dense_prompt_embeddings=dp.to(dev, dtype),
+                multimask_output=True, repeat_image=False,
+                high_res_features=[h.to(dev, dtype) for h in hi]
+            )
+
+            # Postprocess to original dimensions
+            logits = predictor._transforms.postprocess_masks(low, (SQUARE_DIM, SQUARE_DIM))
+            best_idx = scores.argmax().item()
+            logit_crop = logits[0, best_idx, top:top+nh, left:left+nw].unsqueeze(0).unsqueeze(0)
+            logit_full = F.interpolate(logit_crop, size=(H, W), mode="bilinear", align_corners=False)[0, 0]
+            
+            prob = torch.sigmoid(logit_full).float().cpu().numpy()
+
+        # Generate Heatmap
+        heatmap_cv = cv2.applyColorMap((prob * 255).astype(np.uint8), cv2.COLORMAP_JET)
+        heatmap_rgb = cv2.cvtColor(heatmap_cv, cv2.COLOR_BGR2RGB)
+        
+        # Initial Overlay
+        mask = (prob > threshold).astype(np.uint8) * 255
+        overlay = make_overlay(rgb_orig, mask, key=text_prompt)
+        
+        return overlay, Image.fromarray(heatmap_rgb), prob
+
+    except Exception:
+        traceback.print_exc()
+        return None, None, None
+    finally:
+        sam_model.to("cpu")
+        plm_model.to("cpu")
+        torch.cuda.empty_cache()
+
+def update_threshold_ui(image_pil, text_prompt, threshold, cached_prob):
+    """Instant update using CPU only."""
+    if image_pil is None or cached_prob is None:
+        return None
+    rgb_orig = np.array(image_pil.convert("RGB"))
+    mask = (cached_prob > threshold).astype(np.uint8) * 255
+    return make_overlay(rgb_orig, mask, key=text_prompt)
+
+# ----------------- Gradio UI -----------------
+
+with gr.Blocks(title="SAM2 + PLM Segmentation") as demo:
+    prob_state = gr.State()
+    
+    gr.Markdown("# SAM2 + PLM Interactive Segmentation")
+    gr.Markdown("Select a stage, enter a prompt, and run. Adjust the slider for **instant** mask updates.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            text_prompt = gr.Textbox(label="Text Prompt", placeholder="e.g., 'the surgical forceps'")
+            
+            with gr.Row():
+                stage_select = gr.Radio(
+                    choices=list(MODEL_CONFIGS.keys()), 
+                    value="Stage 2 (grad-acc: 8)", 
+                    label="Model Stage"
+                )
+                threshold_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Threshold")
+            
+            run_btn = gr.Button("Run Inference", variant="primary")
+        
+        with gr.Column():
+            out_overlay = gr.Image(label="Segmentation Overlay", type="pil")
+            out_heatmap = gr.Image(label="Probability Heatmap", type="pil")
+
+    # Full Pipeline
+    run_btn.click(
+        fn=run_prediction,
+        inputs=[input_image, text_prompt, threshold_slider, stage_select],
+        outputs=[out_overlay, out_heatmap, prob_state]
+    )
+
+    # Lightweight update on slider move
+    threshold_slider.change(
+        fn=update_threshold_ui,
+        inputs=[input_image, text_prompt, threshold_slider, prob_state],
+        outputs=[out_overlay]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

app.py

@@ -3,51 +3,36 @@ import logging
 import hashlib
 import sys
 import traceback
-import copy
 import tempfile
-
 import cv2
 import numpy as np
 import torch
 import torch.nn.functional as F
 import gradio as gr
-from PIL import Image, ImageFilter, ImageChops, ImageDraw
+from PIL import Image, ImageFilter, ImageChops
 from huggingface_hub import hf_hub_download
 import spaces
 
 # --- IMPORT YOUR CUSTOM MODULES ---
+# Ensure these files are present in your file structure
 from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 from plm_adapter_lora_with_image_input_only_text_positions import PLMLanguageAdapter
 
 # ----------------- Configuration -----------------
+logging.basicConfig(level=logging.INFO)
+
+# Single Model Configuration
+REPO_ID = "aadarsh99/ConvSeg-Stage2"
 SAM2_CONFIG = "sam2_hiera_l.yaml"
 BASE_CKPT_NAME = "sam2_hiera_large.pt"
+FINE_TUNED_SAM = "fine_tuned_sam2_batched_18000.torch"
+FINE_TUNED_PLM = "fine_tuned_sam2_batched_plm_18000.torch"
 
 SQUARE_DIM = 1024
-logging.basicConfig(level=logging.INFO)
-
-# Refactored to store specific filenames per model choice
-MODEL_CONFIGS = {
-    "Stage 1": {
-        "repo_id": "aadarsh99/ConvSeg-Stage1",
-        "sam_filename": "fine_tuned_sam2_batched_100000.torch",
-        "plm_filename": "fine_tuned_sam2_batched_plm_100000.torch"
-    },
-    "Stage 2 (grad-acc: 4)": {
-        "repo_id": "aadarsh99/ConvSeg-Stage2",
-        "sam_filename": "fine_tuned_sam2_batched_18000.torch",
-        "plm_filename": "fine_tuned_sam2_batched_plm_18000.torch"
-    },
-    "Stage 2 (grad-acc: 8)": {
-        "repo_id": "aadarsh99/ConvSeg-Stage2",
-        "sam_filename": "fine_tuned_sam2_batched_18000.torch",
-        "plm_filename": "fine_tuned_sam2_batched_plm_18000.torch"
-    }
-}
 
-# Dynamically create cache keys based on config
-MODEL_CACHE = {k: {"sam": None, "plm": None} for k in MODEL_CONFIGS.keys()}
+# Global Cache
+MODEL_CACHE = {"sam": None, "plm": None}
 
 # ----------------- Helper Functions -----------------
 def download_if_needed(repo_id, filename):
@@ -59,55 +44,49 @@ def download_if_needed(repo_id, filename):
 
 def stable_color(key: str):
     h = int(hashlib.sha256(str(key).encode("utf-8")).hexdigest(), 16)
+    # Bright, distinct colors for overlays
     EDGE_COLORS_HEX = ["#3A86FF", "#FF006E", "#43AA8B", "#F3722C", "#8338EC", "#90BE6D"]
     colors = [tuple(int(c.lstrip("#")[i:i+2], 16) for i in (0, 2, 4)) for c in EDGE_COLORS_HEX]
     return colors[h % len(colors)]
 
 def make_overlay(rgb: np.ndarray, mask: np.ndarray, key: str = "mask") -> Image.Image:
-    # Convert base to RGBA
     base = Image.fromarray(rgb.astype(np.uint8)).convert("RGBA")
     mask_bool = mask > 0
     color = stable_color(key)
     
-    # Create fill layer (Semi-transparent)
+    # Fill layer (Semi-transparent)
     fill_layer = Image.new("RGBA", base.size, color + (0,))
     fill_alpha = Image.fromarray((mask_bool.astype(np.uint8) * 140), "L")
     fill_layer.putalpha(fill_alpha)
 
-    # Create stroke/edge layer
+    # Stroke/Edge layer
     m = Image.fromarray((mask_bool.astype(np.uint8) * 255), "L")
     edges = ImageChops.difference(m.filter(ImageFilter.MaxFilter(3)), m.filter(ImageFilter.MinFilter(3)))
     stroke_layer = Image.new("RGBA", base.size, color + (255,))
     stroke_layer.putalpha(edges)
 
-    # Composite safely
+    # Composite
     out = Image.alpha_composite(base, fill_layer)
     out = Image.alpha_composite(out, stroke_layer)
-    
     return out.convert("RGB")
 
-def ensure_models_loaded(stage_key):
+def ensure_models_loaded():
     global MODEL_CACHE
-    if MODEL_CACHE[stage_key]["sam"] is not None: 
+    if MODEL_CACHE["sam"] is not None: 
         return
     
-    config = MODEL_CONFIGS[stage_key]
-    repo_id = config["repo_id"]
-    
-    logging.info(f"Loading {stage_key} models from {repo_id} into CPU RAM...")
+    logging.info(f"Loading models from {REPO_ID}...")
     
-    # SAM2
-    # Base model is always the same
-    base_path = download_if_needed(repo_id, BASE_CKPT_NAME)
+    # 1. Load SAM2 Base & Fine-tuned weights
+    base_path = download_if_needed(REPO_ID, BASE_CKPT_NAME)
     model = build_sam2(SAM2_CONFIG, base_path, device="cpu")
     
-    # Load specific fine-tuned checkpoint
-    final_path = download_if_needed(repo_id, config["sam_filename"])
-    sd = torch.load(final_path, map_location="cpu")
+    sam_ckpt_path = download_if_needed(REPO_ID, FINE_TUNED_SAM)
+    sd = torch.load(sam_ckpt_path, map_location="cpu")
     model.load_state_dict(sd.get("model", sd), strict=True)
     
-    # PLM
-    plm_path = download_if_needed(repo_id, config["plm_filename"])
+    # 2. Load PLM Adapter
+    plm_path = download_if_needed(REPO_ID, FINE_TUNED_PLM)
     plm = PLMLanguageAdapter(
         model_name="Qwen/Qwen2.5-VL-3B-Instruct", 
         transformer_dim=model.sam_mask_decoder.transformer_dim, 
@@ -119,19 +98,22 @@ def ensure_models_loaded(stage_key):
     plm.load_state_dict(plm_sd["plm"], strict=True)
     plm.eval()
     
-    MODEL_CACHE[stage_key]["sam"], MODEL_CACHE[stage_key]["plm"] = model, plm
+    MODEL_CACHE["sam"] = model
+    MODEL_CACHE["plm"] = plm
+    logging.info("Models loaded successfully.")
 
 # ----------------- GPU Inference -----------------
 
 @spaces.GPU(duration=120) 
-def run_prediction(image_pil, text_prompt, threshold, stage_choice):
+def run_prediction(image_pil, text_prompt, threshold=0.5):
     if image_pil is None or not text_prompt:
         return None, None, None
 
-    ensure_models_loaded(stage_choice)
-    sam_model = MODEL_CACHE[stage_choice]["sam"]
-    plm_model = MODEL_CACHE[stage_choice]["plm"]
+    ensure_models_loaded()
+    sam_model = MODEL_CACHE["sam"]
+    plm_model = MODEL_CACHE["plm"]
     
+    # Move to GPU
     sam_model.to("cuda")
     plm_model.to("cuda")
     
@@ -141,25 +123,26 @@ def run_prediction(image_pil, text_prompt, threshold, stage_choice):
             rgb_orig = np.array(image_pil.convert("RGB"))
             H, W = rgb_orig.shape[:2]
             
-            # Padding math
+            # Smart Resizing & Padding
             scale = SQUARE_DIM / max(H, W)
             nw, nh = int(W * scale), int(H * scale)
             top, left = (SQUARE_DIM - nh) // 2, (SQUARE_DIM - nw) // 2
 
-            # Resize & Pad
             rgb_sq = cv2.resize(rgb_orig, (nw, nh), interpolation=cv2.INTER_LINEAR)
             rgb_sq = cv2.copyMakeBorder(rgb_sq, top, SQUARE_DIM-nh-top, left, SQUARE_DIM-nw-left, cv2.BORDER_CONSTANT, value=0)
             
+            # Image Encoder
             predictor.set_image(rgb_sq)
             image_emb = predictor._features["image_embed"][-1].unsqueeze(0)
             hi = [lvl[-1].unsqueeze(0) for lvl in predictor._features["high_res_feats"]]
 
-            # PLM adapter
+            # PLM Adapter (Text + Image processing)
             with tempfile.NamedTemporaryFile(suffix=".jpg") as tmp:
                 image_pil.save(tmp.name)
+                # Qwen/PLM processes the text prompt here
                 sp, dp = plm_model([text_prompt], image_emb.shape[2], image_emb.shape[3], [tmp.name])
 
-            # SAM2 Decoding
+            # SAM2 Mask Decoder
             dec = sam_model.sam_mask_decoder
             dev, dtype = next(dec.parameters()).device, next(dec.parameters()).dtype
             
@@ -172,19 +155,19 @@ def run_prediction(image_pil, text_prompt, threshold, stage_choice):
                 high_res_features=[h.to(dev, dtype) for h in hi]
             )
 
-            # Postprocess to original dimensions
+            # Post-processing
             logits = predictor._transforms.postprocess_masks(low, (SQUARE_DIM, SQUARE_DIM))
             best_idx = scores.argmax().item()
+            
             logit_crop = logits[0, best_idx, top:top+nh, left:left+nw].unsqueeze(0).unsqueeze(0)
             logit_full = F.interpolate(logit_crop, size=(H, W), mode="bilinear", align_corners=False)[0, 0]
             
             prob = torch.sigmoid(logit_full).float().cpu().numpy()
 
-        # Generate Heatmap
+        # Visuals
         heatmap_cv = cv2.applyColorMap((prob * 255).astype(np.uint8), cv2.COLORMAP_JET)
         heatmap_rgb = cv2.cvtColor(heatmap_cv, cv2.COLOR_BGR2RGB)
         
-        # Initial Overlay
         mask = (prob > threshold).astype(np.uint8) * 255
         overlay = make_overlay(rgb_orig, mask, key=text_prompt)
         
@@ -192,55 +175,102 @@ def run_prediction(image_pil, text_prompt, threshold, stage_choice):
 
     except Exception:
         traceback.print_exc()
-        return None, None, None
+        raise gr.Error("Inference failed. Please check logs.")
     finally:
+        # Cleanup memory
         sam_model.to("cpu")
         plm_model.to("cpu")
         torch.cuda.empty_cache()
 
 def update_threshold_ui(image_pil, text_prompt, threshold, cached_prob):
-    """Instant update using CPU only."""
+    """Real-time update using CPU only (no GPU quota usage)."""
     if image_pil is None or cached_prob is None:
         return None
     rgb_orig = np.array(image_pil.convert("RGB"))
     mask = (cached_prob > threshold).astype(np.uint8) * 255
     return make_overlay(rgb_orig, mask, key=text_prompt)
 
-# ----------------- Gradio UI -----------------
+# ----------------- UI Styling & Layout -----------------
+
+custom_css = """
+h1 {
+    text-align: center;
+    display: block;
+}
+.subtitle {
+    text-align: center;
+    font-size: 1.1em;
+    margin-bottom: 20px;
+}
+"""
+
+theme = gr.themes.Soft(
+    primary_hue="blue",
+    neutral_hue="slate",
+).set(
+    button_primary_background_fill="*primary_600",
+    button_primary_background_fill_hover="*primary_700",
+)
 
-with gr.Blocks(title="SAM2 + PLM Segmentation") as demo:
+with gr.Blocks(theme=theme, css=custom_css, title="ConvSeg-Net Demo") as demo:
     prob_state = gr.State()
     
-    gr.Markdown("# SAM2 + PLM Interactive Segmentation")
-    gr.Markdown("Select a stage, enter a prompt, and run. Adjust the slider for **instant** mask updates.")
-    
+    # Header
+    gr.Markdown("# 🧩 Conversational Image Segmentation")
+    gr.Markdown(
+        "<div class='subtitle'>Grounding abstract concepts and physics-based reasoning into pixel-accurate masks.<br>"
+        "Powered by <b>SAM2 + Qwen2.5-VL</b></div>"
+    )
+
     with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="pil", label="Input Image")
-            text_prompt = gr.Textbox(label="Text Prompt", placeholder="e.g., 'the surgical forceps'")
+        # --- Left Column: Inputs ---
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="pil", label="Input Image", height=400)
             
-            with gr.Row():
-                stage_select = gr.Radio(
-                    choices=list(MODEL_CONFIGS.keys()), 
-                    value="Stage 2 (grad-acc: 8)", 
-                    label="Model Stage"
+            with gr.Group():
+                text_prompt = gr.Textbox(
+                    label="Conversational Prompt", 
+                    placeholder="e.g., Segment the object that is prone to rolling...",
+                    lines=2
                 )
-                threshold_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Threshold")
+                gr.Markdown("💡 **Tip:** The model works best when prompts start with **'Segment the...'**")
             
-            run_btn = gr.Button("Run Inference", variant="primary")
-        
-        with gr.Column():
-            out_overlay = gr.Image(label="Segmentation Overlay", type="pil")
-            out_heatmap = gr.Image(label="Probability Heatmap", type="pil")
+            with gr.Accordion("⚙️ Advanced Options", open=False):
+                threshold_slider = gr.Slider(
+                    0.0, 1.0, value=0.5, step=0.01, 
+                    label="Mask Confidence Threshold",
+                    info="Adjust after running to refine the mask edges."
+                )
+
+            run_btn = gr.Button("🚀 Run Segmentation", variant="primary", size="lg")
+
+        # --- Right Column: Outputs ---
+        with gr.Column(scale=1):
+            out_overlay = gr.Image(label="Segmentation Result", type="pil")
+            out_heatmap = gr.Image(label="Confidence Heatmap", type="pil")
 
-    # Full Pipeline
+    # --- Examples Section ---
+    gr.Markdown("### 📝 Try Examples")
+    gr.Examples(
+        examples=[
+            ["./examples/bike.jpg", "Segment the mechanism requiring physical pedaling force"],
+            ["./examples/luggage.jpg", "Segment the suitcase that is easiest to remove without disturbing others"],
+            ["./examples/kitchen.jpg", "Segment the surface suitable for hot cookware"],
+        ],
+        inputs=[input_image, text_prompt],
+        # cache_examples=True # Uncomment if you want to pre-compute these on startup
+    )
+
+    # --- Event Handling ---
+    
+    # 1. Run Inference (GPU)
     run_btn.click(
         fn=run_prediction,
-        inputs=[input_image, text_prompt, threshold_slider, stage_select],
+        inputs=[input_image, text_prompt, threshold_slider],
         outputs=[out_overlay, out_heatmap, prob_state]
     )
 
-    # Lightweight update on slider move
+    # 2. Update Threshold (CPU - Instant)
     threshold_slider.change(
         fn=update_threshold_ui,
         inputs=[input_image, text_prompt, threshold_slider, prob_state],
@@ -248,4 +278,4 @@ with gr.Blocks(title="SAM2 + PLM Segmentation") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue().launch()