""" vision-space — room analysis: depth + segmentation + CLIP. ARCHITECTURE (updated — see PLAN.md "Segmentation Technology Decision"): Depth Anything V2 Large → depth map Grounding DINO Base → bounding boxes with text-prompted labels SAM 2 Large → pixel masks from those boxes CLIP ViT-B/32 → 512-dim room embedding The old pipeline (SAM mask-generation → CLIP crop labeling) is replaced. Root cause of old failures: SAM produced class-agnostic blobs, then CLIP guessed labels by comparing cropped patches to word embeddings — a curtain and a floor lamp are close in CLIP space, so they were routinely swapped. New flow: detect first (Grounding DINO uses text prompts to find specific objects), then segment (SAM 2 draws a precise mask inside each detected box). The label is established before the mask, not inferred after. """ import base64 import io import gradio as gr import numpy as np import spaces import torch import transformers print(f"transformers version: {transformers.__version__}") from PIL import Image from sentence_transformers import SentenceTransformer from transformers import ( AutoImageProcessor, AutoModelForDepthEstimation, AutoProcessor, AutoModelForZeroShotObjectDetection, Sam2Model, ) # --------------------------------------------------------------------------- # Detection vocabulary — passed verbatim to Grounding DINO as text prompt. # The model returns only objects it finds matching these labels. # Add / remove freely; no retraining needed (zero-shot). # --------------------------------------------------------------------------- DETECT_CATEGORIES = [ "sofa", "couch", "sectional sofa", "loveseat", "chair", "armchair", "accent chair", "dining table", "coffee table", "side table", "console table", "floor lamp", "table lamp", "pendant lamp", "chandelier", "rug", "carpet", "curtain", "drape", "blind", "window", "door", "floor", "wall", "ceiling", "shelf", "bookshelf", "cabinet", "wardrobe", "dresser", "bed", "headboard", "plant", "potted plant", "television", "fireplace", "mirror", "artwork", "painting", "cushion", "pillow", ] # Grounding DINO format: "label . label . label ." GROUNDING_PROMPT = " . ".join(DETECT_CATEGORIES) + " ." MAX_IMAGE_SIZE = 1024 MAX_SEGMENTS = 25 BOX_THRESHOLD = 0.25 # min Grounding DINO box confidence TEXT_THRESHOLD = 0.20 # min Grounding DINO text-match confidence # --------------------------------------------------------------------------- # Load all models on CPU at startup — GPU only inside @spaces.GPU # --------------------------------------------------------------------------- print("Loading Depth Anything V2 Large...") depth_processor = AutoImageProcessor.from_pretrained( "depth-anything/Depth-Anything-V2-Large-hf" ) depth_model = AutoModelForDepthEstimation.from_pretrained( "depth-anything/Depth-Anything-V2-Large-hf" ) depth_model.eval() print("Loading Grounding DINO Base...") gdino_processor = AutoProcessor.from_pretrained("IDEA-Research/grounding-dino-base") gdino_model = AutoModelForZeroShotObjectDetection.from_pretrained( "IDEA-Research/grounding-dino-base" ) gdino_model.eval() print("Loading SAM 2 Large...") sam2_processor = AutoProcessor.from_pretrained("facebook/sam2-hiera-large") sam2_model = Sam2Model.from_pretrained("facebook/sam2-hiera-large") sam2_model.eval() print("Loading CLIP ViT-B/32...") clip_model = SentenceTransformer("clip-ViT-B-32") print("All models loaded on CPU.") # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _resize(image: Image.Image, max_size: int) -> Image.Image: """Resize keeping aspect ratio so the longest side = max_size.""" w, h = image.size if max(w, h) <= max_size: return image scale = max_size / max(w, h) return image.resize((int(w * scale), int(h * scale)), Image.LANCZOS) def _mask_to_b64png(mask: np.ndarray) -> str: img = Image.fromarray(mask.astype(np.uint8) * 255, mode="L") buf = io.BytesIO() img.save(buf, format="PNG", optimize=True) return base64.b64encode(buf.getvalue()).decode() # --------------------------------------------------------------------------- # Main analysis function — runs on GPU # --------------------------------------------------------------------------- @spaces.GPU(duration=90) def analyze(image_base64: str) -> dict: import traceback try: device = "cuda" if torch.cuda.is_available() else "cpu" print(f"analyze: device={device}") image = Image.open(io.BytesIO(base64.b64decode(image_base64))).convert("RGB") image = _resize(image, MAX_IMAGE_SIZE) w, h = image.size print(f"analyze: image {w}x{h}") # ------------------------------------------------------------------ # Step 1 — Depth Anything V2 # ------------------------------------------------------------------ depth_model.to(device) depth_inputs = depth_processor(images=image, return_tensors="pt").to(device) with torch.no_grad(): depth_np = depth_model(**depth_inputs).predicted_depth[0].cpu().float().numpy() dmin, dmax = depth_np.min(), depth_np.max() depth_np = (depth_np - dmin) / (dmax - dmin) if dmax > dmin else np.zeros_like(depth_np) depth_model.to("cpu") torch.cuda.empty_cache() print(f"analyze: depth done {depth_np.shape}") # ------------------------------------------------------------------ # Step 2 — Grounding DINO: detect objects by text prompt # ------------------------------------------------------------------ gdino_model.to(device) gdino_inputs = gdino_processor( images=image, text=GROUNDING_PROMPT, return_tensors="pt", ).to(device) with torch.no_grad(): gdino_outputs = gdino_model(**gdino_inputs) # transformers ≥ 4.47 renamed some parameters — try both signatures. try: post_proc = gdino_processor.post_process_grounded_object_detection( gdino_outputs, gdino_inputs.input_ids, box_threshold=BOX_THRESHOLD, text_threshold=TEXT_THRESHOLD, target_sizes=[(h, w)], ) except TypeError: # Older / newer API: call without threshold kwargs, filter manually. post_proc = gdino_processor.post_process_grounded_object_detection( gdino_outputs, gdino_inputs.input_ids, target_sizes=[(h, w)], ) detections = post_proc[0] raw_boxes = detections["boxes"] # transformers ≥ 4.51 returns integer ids in "labels"; prefer "text_labels". raw_labels = detections.get("text_labels", detections["labels"]) raw_scores = detections["scores"] # Ensure scores are a plain float list for filtering. scores_list = raw_scores.cpu().tolist() if hasattr(raw_scores, "cpu") else list(raw_scores) keep = [i for i, s in enumerate(scores_list) if s >= BOX_THRESHOLD] boxes = raw_boxes.cpu().tolist() if hasattr(raw_boxes, "cpu") else list(raw_boxes) boxes = [boxes[i] for i in keep] labels = [raw_labels[i] for i in keep] scores = [scores_list[i] for i in keep] gdino_model.to("cpu") torch.cuda.empty_cache() print(f"analyze: Grounding DINO → {len(boxes)} detections") for lbl, sc in zip(labels, scores): print(f" {lbl} ({sc:.2f})") if not boxes: room_emb = clip_model.encode(image, normalize_embeddings=True).tolist() return {"depth_array": depth_np.tolist(), "width": w, "height": h, "segments": [], "room_clip_embedding": room_emb} # ------------------------------------------------------------------ # Step 3 — SAM 2: generate pixel mask for each detected box # One SAM 2 call per box — reliable and keeps VRAM low. # ------------------------------------------------------------------ sam2_model.to(device) segments = [] # Sort by score descending so highest-confidence detections survive MAX_SEGMENTS cap. sorted_detections = sorted(zip(scores, boxes, labels), reverse=True) scores, boxes, labels = zip(*sorted_detections) if sorted_detections else ([], [], []) for box, label, score in zip(boxes, labels, scores): if len(segments) >= MAX_SEGMENTS: break try: sam2_inputs = sam2_processor( images=image, input_boxes=[[box]], # [image_level][box_level][4 coords] return_tensors="pt", ).to(device) with torch.no_grad(): sam2_out = sam2_model(**sam2_inputs, multimask_output=False) masks = sam2_processor.post_process_masks( sam2_out.pred_masks.cpu(), sam2_inputs["original_sizes"].cpu(), ) # masks[0]: [1 box, 1 mask, H, W] mask_np = masks[0][0, 0].numpy().astype(bool) if int(mask_np.sum()) < 500: continue x1, y1, x2, y2 = [int(v) for v in box] segments.append({ "label": label, "mask_rle": _mask_to_b64png(mask_np), "confidence": round(float(score), 4), "bbox_xyxy": [x1, y1, x2, y2], }) except Exception as seg_err: print(f" SAM2 error for '{label}': {seg_err}") continue sam2_model.to("cpu") torch.cuda.empty_cache() print(f"analyze: {len(segments)} segments with masks") # ------------------------------------------------------------------ # Step 4 — CLIP room embedding (unchanged) # ------------------------------------------------------------------ room_clip_embedding = clip_model.encode(image, normalize_embeddings=True).tolist() return { "depth_array": depth_np.tolist(), "width": w, "height": h, "segments": segments, "room_clip_embedding": room_clip_embedding, } except Exception as e: traceback.print_exc() raise RuntimeError(f"analyze failed: {type(e).__name__}: {e}") from e # --------------------------------------------------------------------------- # Gradio interface # --------------------------------------------------------------------------- with gr.Blocks(title="Vision Space") as demo: gr.Markdown("## Room Analysis — Depth Anything V2 + Grounded SAM 2 + CLIP") with gr.Row(): b64_in = gr.Textbox(label="image_base64", lines=4) json_out = gr.JSON(label="Result") gr.Button("Analyse").click(analyze, inputs=b64_in, outputs=json_out, api_name="analyze") demo.launch(show_error=True)