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()