update app

app.py

@@ -4,6 +4,7 @@ import hashlib
 import sys
 import traceback
 import copy
+import tempfile
 
 import cv2
 import numpy as np
@@ -41,6 +42,7 @@ MODEL_CACHE = {
 # ----------------- 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}")
@@ -48,39 +50,63 @@ def download_if_needed(repo_id, filename):
 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(h.lstrip("#")[i:i+2], 16) for i in (0, 2, 4)) for h in EDGE_COLORS_HEX]
+    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:
-    base = Image.fromarray(rgb.astype(np.uint8)).convert("RGB").convert("RGBA")
+    # 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) * 178), "L")
+    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 = Image.new("RGBA", base.size, color + (0,))
-    stroke.putalpha(edges)
+    stroke_layer = Image.new("RGBA", base.size, color + (255,))
+    stroke_layer.putalpha(edges)
 
-    return Image.alpha_composite(base, fill_layer).alpha_composite(stroke).convert("RGB")
+    # Composite safely (Module-level returns new images, no in-place None issues)
+    out = Image.alpha_composite(base, fill_layer)
+    out = Image.alpha_composite(out, stroke_layer)
+    
+    return out.convert("RGB")
 
 def ensure_models_loaded(stage):
     global MODEL_CACHE
-    if MODEL_CACHE[stage]["sam"] is not None: return
+    if MODEL_CACHE[stage]["sam"] is not None: 
+        return
+        
     repo_id = REPO_MAP[stage]
+    logging.info(f"Loading {stage} models from {repo_id} into CPU RAM...")
+    
+    # SAM2
     base_path = download_if_needed(repo_id, BASE_CKPT_NAME)
     model = build_sam2(SAM2_CONFIG, base_path, device="cpu")
-    sd = torch.load(download_if_needed(repo_id, FINAL_CKPT_NAME), map_location="cpu")
+    final_path = download_if_needed(repo_id, FINAL_CKPT_NAME)
+    sd = torch.load(final_path, map_location="cpu")
     model.load_state_dict(sd.get("model", sd), strict=True)
-    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.load_state_dict(torch.load(download_if_needed(repo_id, PLM_CKPT_NAME), map_location="cpu")["plm"], strict=True)
+    
+    # PLM
+    plm_path = download_if_needed(repo_id, PLM_CKPT_NAME)
+    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]["sam"], MODEL_CACHE[stage]["plm"] = model, plm
 
-# ----------------- Core Logic -----------------
+# ----------------- GPU Inference -----------------
 
 @spaces.GPU(duration=120) 
 def run_prediction(image_pil, text_prompt, threshold, stage_choice):
@@ -88,52 +114,77 @@ def run_prediction(image_pil, text_prompt, threshold, stage_choice):
         return None, None, None
 
     ensure_models_loaded(stage_choice)
-    sam_model, plm_model = MODEL_CACHE[stage_choice]["sam"], MODEL_CACHE[stage_choice]["plm"]
-    sam_model.to("cuda"), plm_model.to("cuda")
+    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
 
-            # Preprocess & Encode
-            rgb_sq = cv2.copyMakeBorder(cv2.resize(rgb_orig, (nw, nh)), top, SQUARE_DIM-nh-top, left, SQUARE_DIM-nw-left, cv2.BORDER_CONSTANT, value=0)
+            # 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 & SAM2 Decoder
-            temp_path = "temp.jpg"
-            image_pil.save(temp_path)
-            sp, dp = plm_model([text_prompt], image_emb.shape[2], image_emb.shape[3], [temp_path])
-            low, scores, _, _ = sam_model.sam_mask_decoder(
-                image_embeddings=image_emb.to("cuda"), image_pe=sam_model.sam_prompt_encoder.get_dense_pe().to("cuda"),
-                sparse_prompt_embeddings=sp.to("cuda"), dense_prompt_embeddings=dp.to("cuda"),
-                multimask_output=True, repeat_image=False, high_res_features=[h.to("cuda") for h in hi]
+            # 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 full size
+            # Postprocess to original dimensions
             logits = predictor._transforms.postprocess_masks(low, (SQUARE_DIM, SQUARE_DIM))
-            logit_crop = logits[0, scores.argmax().item(), top:top+nh, left:left+nw].unsqueeze(0).unsqueeze(0)
+            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().detach().cpu().numpy()
+            
+            prob = torch.sigmoid(logit_full).float().cpu().numpy()
 
-        # Initial visualization
-        heatmap = cv2.applyColorMap((prob * 255).astype(np.uint8), cv2.COLORMAP_JET)
-        overlay = make_overlay(rgb_orig, (prob > threshold).astype(np.uint8) * 255, key=text_prompt)
+        # Generate Heatmap
+        heatmap_cv = cv2.applyColorMap((prob * 255).astype(np.uint8), cv2.COLORMAP_JET)
+        heatmap_rgb = cv2.cvtColor(heatmap_cv, cv2.COLOR_BGR2RGB)
         
-        return overlay, Image.fromarray(cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)), prob
+        # 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")
+        sam_model.to("cpu")
+        plm_model.to("cpu")
         torch.cuda.empty_cache()
 
 def update_threshold_ui(image_pil, text_prompt, threshold, cached_prob):
-    """Updates the overlay instantly without rerunning the GPU model."""
+    """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"))
@@ -142,32 +193,35 @@ def update_threshold_ui(image_pil, text_prompt, threshold, cached_prob):
 
 # ----------------- Gradio UI -----------------
 
-with gr.Blocks(title="SAM2 + PLM Interactive") as demo:
-    prob_state = gr.State() # Caches the probability map
+with gr.Blocks(title="SAM2 + PLM Segmentation") as demo:
+    prob_state = gr.State()
     
-    gr.Markdown("# SAM2 + PLM Segmentation\n*Change the model/prompt and click **Run Inference**. Then, adjust the **Threshold** slider for instant mask updates.*")
+    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 blue scissors'")
+            text_prompt = gr.Textbox(label="Text Prompt", placeholder="e.g., 'the surgical forceps'")
+            
             with gr.Row():
-                stage_select = gr.Radio(choices=["Stage 1", "Stage 2"], value="Stage 1", label="Model")
+                stage_select = gr.Radio(choices=["Stage 1", "Stage 2"], value="Stage 1", 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")
 
-    # 1. Clicking the button runs the heavy inference
+    # Full Pipeline
     run_btn.click(
         fn=run_prediction,
         inputs=[input_image, text_prompt, threshold_slider, stage_select],
         outputs=[out_overlay, out_heatmap, prob_state]
     )
 
-    # 2. Moving the slider triggers only the lightweight update
+    # Lightweight update on slider move
     threshold_slider.change(
         fn=update_threshold_ui,
         inputs=[input_image, text_prompt, threshold_slider, prob_state],