app.py

import gradio as gr
import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont
import json
import random
from transformers import (
    AutoModelForZeroShotObjectDetection,
    AutoProcessor,
    Sam2Model,
    Sam2Processor,
    DepthAnythingForDepthEstimation,
)

# -----------------------------
# Device setup
# -----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

print("⏳ Loading models into memory...")


# -----------------------------
# Model loading
# -----------------------------
# 1. Grounding DINO (open-vocabulary object detection -> provides labels + boxes).
#    Not gated, so no HF token/license acceptance is required.
GROUNDING_ID = "IDEA-Research/grounding-dino-tiny"
g_processor = AutoProcessor.from_pretrained(GROUNDING_ID)
g_model = AutoModelForZeroShotObjectDetection.from_pretrained(GROUNDING_ID).to(device).eval()

# 2. SAM2 (segmentation using boxes from Grounding DINO as prompts). Not gated.
SAM2_ID = "facebook/sam2-hiera-base-plus"
sam2_processor = Sam2Processor.from_pretrained(SAM2_ID)
sam2_model = Sam2Model.from_pretrained(SAM2_ID).to(device).eval()

# 3. Depth Anything V2 Large (monocular depth estimation)
DEPTH_ID = "depth-anything/Depth-Anything-V2-Large-hf"
depth_processor = AutoProcessor.from_pretrained(DEPTH_ID)
depth_model = DepthAnythingForDepthEstimation.from_pretrained(DEPTH_ID).to(device).eval()


# -----------------------------
# Pipeline steps
# -----------------------------
def detect_objects(image: Image.Image, text_prompt: str, threshold: float = 0.3):
    """Step 1: Open-vocabulary detection -> boxes + labels.

    Grounding DINO expects each class as a separate phrase. Passing a nested
    list (`[[cls1, cls2, ...]]`) lets the processor join them correctly
    (period-separated) instead of relying on a raw comma-joined string, which
    the model tends to treat as one long phrase and fail to ground.
    """
    class_names = [c.strip() for c in text_prompt.split(",") if c.strip()]
    inputs = g_processor(images=image, text=[class_names], return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = g_model(**inputs)
    results = g_processor.post_process_grounded_object_detection(
        outputs,
        inputs.input_ids,
        threshold=threshold,
        text_threshold=threshold,
        target_sizes=[image.size[::-1]],
    )
    return results[0]


def segment_with_sam2(image: Image.Image, boxes):
    """Step 2: Use detected boxes as SAM2 prompts -> binary masks."""
    if len(boxes) == 0:
        return []
    input_boxes = [[[float(b[0]), float(b[1]), float(b[2]), float(b[3])] for b in boxes]]
    inputs = sam2_processor(images=image, input_boxes=input_boxes, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = sam2_model(**inputs, multimask_output=False)
    masks = sam2_processor.post_process_masks(
        outputs.pred_masks.cpu(),
        inputs["original_sizes"],
    )[0]  # [N, 1, H, W]
    return [m[0].numpy().astype(bool) for m in masks]  # -> list of H x W bool arrays


def estimate_depth(image: Image.Image):
    """Step 3: Monocular depth -> depth map at original resolution."""
    inputs = depth_processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = depth_model(**inputs)
    predicted_depth = outputs.predicted_depth
    interpolation = torch.nn.functional.interpolate(
        predicted_depth.unsqueeze(1),
        size=image.size[::-1],
        mode="bicubic",
        align_corners=False,
    )
    depth_map = interpolation.squeeze().cpu().numpy()
    # Normalize to meters-like scale (relative depth; calibrate if you have absolute metric)
    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min() + 1e-8)
    return depth_map


def average_depth_per_mask(depth_map, masks):
    """Step 4: Compute mean depth within each binary mask."""
    out = []
    for m in masks:  # each mask is already an H x W bool array
        if m.sum() == 0:
            out.append(None)
        else:
            out.append(float(depth_map[m].mean()))
    return out


def render_visualization(image, boxes, labels, masks, avg_depths):
    """Overlay masks, boxes, labels and average depth on the image."""
    overlay = image.copy().convert("RGBA")
    draw_img = image.copy().convert("RGB")
    draw = ImageDraw.Draw(draw_img)

    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255),
              (255, 255, 0), (255, 0, 255), (0, 255, 255)]

    mask_layer = Image.new("RGBA", image.size, (0, 0, 0, 0))

    for i, (box, label, depth) in enumerate(zip(boxes, labels, avg_depths)):
        if depth is None:
            continue
        color = colors[i % len(colors)]
        # Fill mask
        m = masks[i]  # already an H x W bool array
        rgba = np.zeros((m.shape[0], m.shape[1], 4), dtype=np.uint8)
        rgba[m] = color + (120,)  # alpha
        m_img = Image.fromarray(rgba, mode="RGBA")
        mask_layer.paste(m_img, (0, 0), m_img)

        x1, y1, x2, y2 = [int(v) for v in box]
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
        text = f"{label}: {depth:.2f}"
        draw.text((x1, max(0, y1 - 15)), text, fill=color)

    final = Image.alpha_composite(overlay, mask_layer).convert("RGB")
    return final


def run_pipeline(image, text_prompt, threshold):
    """End-to-end inference pipeline."""
    if image is None:
        return None, "{}", None

    image_pil = Image.fromarray(image).convert("RGB")

    # Step 1 — Detection
    detection = detect_objects(image_pil, text_prompt, threshold)
    boxes = detection["boxes"].tolist()
    labels = detection["text_labels"]

    if len(boxes) == 0:
        return image, json.dumps({"error": "No objects detected"}), None

    # Step 2 — Segmentation
    masks = segment_with_sam2(image_pil, boxes)

    # Step 3 — Depth estimation
    depth_map = estimate_depth(image_pil)

    # Step 4 — Average depth per mask
    avg_depths = average_depth_per_mask(depth_map, masks)

    # Step 5 — Aggregate by label and compare
    per_label = {}
    for label, depth in zip(labels, avg_depths):
        if depth is None:
            continue
        per_label.setdefault(label, []).append(depth)

    output = {}
    for label, depths in per_label.items():
        output[label] = {"depth": round(sum(depths) / len(depths), 2)}

    if len(output) >= 2:
        # Difference between closest two objects (smallest spread)
        sorted_items = sorted(output.items(), key=lambda kv: kv[1]["depth"])
        d1 = sorted_items[0][1]["depth"]
        d2 = sorted_items[-1][1]["depth"]
        output["difference"] = round(abs(d2 - d1), 2)
        output["closest_pair"] = {
            "a": sorted_items[0][0],
            "b": sorted_items[1][0],
            "gap": round(abs(sorted_items[1][1]["depth"] - sorted_items[0][1]["depth"]), 2),
        }

    # Visualization
    vis = render_visualization(image_pil, boxes, labels, masks, avg_depths)

    # Depth map as numpy for the depth preview
    depth_vis = (depth_map * 255).astype(np.uint8)

    return vis, json.dumps(output, indent=2), depth_vis


# -----------------------------
# Gradio UI
# -----------------------------
demo = gr.Interface(
    fn=run_pipeline,
    inputs=[
        gr.Image(label="Input Image", type="numpy"),
        gr.Textbox(value="person, car, dog, cat, chair, bottle",
                   label="Object classes (comma-separated)"),
        gr.Slider(0.1, 0.9, value=0.3, step=0.05, label="Detection threshold"),
    ],
    outputs=[
        gr.Image(label="Segmentation + Depth Overlay"),
        gr.Textbox(label="JSON Output"),
        gr.Image(label="Depth Map"),
    ],
    title="🧠 Object Depth Comparison Pipeline",
    description=(
        "Pipeline: Image → Grounding DINO (boxes+labels) → SAM2 (masks) → "
        "Depth Anything V2 Large (depth map) → Avg depth per mask → JSON comparison.\n\n"
        "Tip: write the object classes you expect, e.g. `person, car`."
    )
)

if __name__ == "__main__":
    demo.launch()