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