Update app.py

app.py

@@ -6,10 +6,9 @@ from PIL import Image, ImageDraw, ImageFont
 import json
 import random
 from transformers import (
-    AutoModelForZeroShotObjectDetection,
     AutoProcessor,
-    SamModel,
-    SamProcessor,
+    Sam3Model,
+    Sam3Processor,
     DepthAnythingForDepthEstimation,
 )
 
@@ -22,18 +21,17 @@ print(f"Using device: {device}")
 # -----------------------------
 # Model loading
 # -----------------------------
-# 1. Grounding DINO (open-vocabulary object detection -> provides labels + boxes)
-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. SAM (segmentation using boxes from Grounding DINO as prompts)
-SAM_ID = "facebook/sam-vit-base"
-sam_processor = SamProcessor.from_pretrained(SAM_ID)
-sam_model = SamModel.from_pretrained(SAM_ID).to(device).eval()
-
-# 3. Depth Anything V2 (monocular depth estimation)
-DEPTH_ID = "depth-anything/Depth-Anything-V2-Small-hf"
+# 1. SAM3 (promptable concept segmentation -> detects + segments objects from a
+#    text prompt in a single step, replacing the old Grounding DINO + SAM combo).
+#    NOTE: facebook/sam3 is a gated Meta model on the Hub — the environment needs
+#    a HF token that has accepted the model license (e.g. `huggingface-cli login`
+#    or the `HF_TOKEN` env var) or `from_pretrained` will fail to download it.
+SAM3_ID = "facebook/sam3"
+sam3_processor = Sam3Processor.from_pretrained(SAM3_ID)
+sam3_model = Sam3Model.from_pretrained(SAM3_ID).to(device).eval()
+
+# 2. 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()
 
@@ -41,35 +39,38 @@ depth_model = DepthAnythingForDepthEstimation.from_pretrained(DEPTH_ID).to(devic
 # -----------------------------
 # Pipeline steps
 # -----------------------------
-def detect_objects(image: Image.Image, text_prompt: str, threshold: float = 0.3):
-    """Step 1: Open-vocabulary detection -> boxes + labels."""
-    inputs = g_processor(images=image, text=text_prompt, 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 detect_and_segment(image: Image.Image, text_prompt: str, threshold: float = 0.3):
+    """Step 1+2: SAM3 promptable concept segmentation -> boxes + masks + labels.
 
+    Runs one SAM3 forward pass per requested class, reusing the image's vision
+    embeddings so the (expensive) vision backbone only runs once per image.
+    """
+    class_names = [c.strip() for c in text_prompt.split(",") if c.strip()]
+    if not class_names:
+        return [], [], []
 
-def segment_with_sam(image: Image.Image, boxes):
-    """Step 2: Use detected boxes as SAM 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 = sam_processor(image, input_boxes=input_boxes, return_tensors="pt").to(device)
+    img_inputs = sam3_processor(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
-        outputs = sam_model(**inputs)
-    masks = sam_processor.image_processor.post_process_masks(
-        outputs.pred_masks.cpu(),
-        inputs["original_sizes"].cpu(),
-        inputs["reshaped_input_sizes"].cpu(),
-    )
-    return masks[0]  # [N, 1, H, W]
+        vision_embeds = sam3_model.get_vision_features(pixel_values=img_inputs.pixel_values)
+
+    all_boxes, all_labels, all_masks = [], [], []
+    for class_name in class_names:
+        text_inputs = sam3_processor(text=class_name, return_tensors="pt").to(device)
+        with torch.no_grad():
+            outputs = sam3_model(vision_embeds=vision_embeds, **text_inputs)
+        result = sam3_processor.post_process_instance_segmentation(
+            outputs,
+            threshold=threshold,
+            mask_threshold=0.5,
+            target_sizes=img_inputs.get("original_sizes").tolist(),
+        )[0]
+
+        for box, mask in zip(result["boxes"], result["masks"]):
+            all_boxes.append([float(v) for v in box.tolist()])
+            all_labels.append(class_name)
+            all_masks.append(mask.cpu().numpy().astype(bool))
+
+    return all_boxes, all_labels, all_masks
 
 
 def estimate_depth(image: Image.Image):
@@ -93,8 +94,7 @@ def estimate_depth(image: Image.Image):
 def average_depth_per_mask(depth_map, masks):
     """Step 4: Compute mean depth within each binary mask."""
     out = []
-    for mask in masks:
-        m = mask[0].numpy().astype(bool)  # H x W
+    for m in masks:  # each mask is already an H x W bool array
         if m.sum() == 0:
             out.append(None)
         else:
@@ -118,7 +118,7 @@ def render_visualization(image, boxes, labels, masks, avg_depths):
             continue
         color = colors[i % len(colors)]
         # Fill mask
-        m = masks[i][0].numpy().astype(bool)
+        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")
@@ -140,17 +140,12 @@ def run_pipeline(image, text_prompt, threshold):
 
     image_pil = Image.fromarray(image).convert("RGB")
 
-    # Step 1 — Detection
-    detection = detect_objects(image_pil, text_prompt, threshold)
-    boxes = detection["boxes"].tolist()
-    labels = detection["labels"]
+    # Step 1+2 — SAM3 detection + segmentation
+    boxes, labels, masks = detect_and_segment(image_pil, text_prompt, threshold)
 
     if len(boxes) == 0:
         return image, json.dumps({"error": "No objects detected"}), None
 
-    # Step 2 — Segmentation
-    masks = segment_with_sam(image_pil, boxes)
-
     # Step 3 — Depth estimation
     depth_map = estimate_depth(image_pil)
 
@@ -207,8 +202,8 @@ demo = gr.Interface(
     ],
     title="🧠 Object Depth Comparison Pipeline",
     description=(
-        "Pipeline: Image → Grounding DINO (boxes+labels) → SAM (masks) → "
-        "Depth Anything V2 (depth map) → Avg depth per mask → JSON comparison.\n\n"
+        "Pipeline: Image → SAM3 (text-prompted boxes+masks per class) → "
+        "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`."
     )
 )