app.py

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