Update app.py

app.py

@@ -18,13 +18,18 @@ import numpy as np
 import torch
 import cv2
 from PIL import Image, ExifTags
-from scipy import ndimage
 import gradio as gr
 
 from transformers import Sam3Processor, Sam3Model
 from depth_anything_3.api import DepthAnything3
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# Half precision on CUDA gives a large SAM3 speed/memory win with negligible
+# accuracy impact for this pipeline (mask thresholds are coarse-grained);
+# CPU stays fp32 since there's no benefit there. Depth Anything 3 already
+# does its own internal mixed-precision autocast (see DepthAnything3.forward
+# in depth_anything_3/api.py), so it doesn't need this treatment here.
+SAM3_DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 HF_TOKEN = os.environ.get("HF_TOKEN")  # set as a Space secret if sam3 is gated for you
 ARUCO_DICT = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
 
@@ -40,7 +45,9 @@ _depther = None
 def get_segmenter():
     global _segmenter
     if _segmenter is None:
-        model = Sam3Model.from_pretrained("facebook/sam3", token=HF_TOKEN).to(DEVICE)
+        model = Sam3Model.from_pretrained(
+            "facebook/sam3", token=HF_TOKEN, torch_dtype=SAM3_DTYPE
+        ).to(DEVICE)
         processor = Sam3Processor.from_pretrained("facebook/sam3", token=HF_TOKEN)
         _segmenter = (model, processor)
     return _segmenter
@@ -58,10 +65,23 @@ def get_depther(model_id: str):
 # Pipeline stages (same logic as the standalone script)
 # --------------------------------------------------------------------------
 
-def segment(image: Image.Image, text_prompt: str, score_threshold: float = 0.5) -> np.ndarray:
+def segment_batch(image: Image.Image, text_prompts: list, score_threshold: float = 0.5) -> list:
+    """Segments every text prompt against the same primary image in a single
+    batched SAM3 forward pass, instead of one full forward pass (image
+    encoder included) per object as before. The image is simply repeated
+    across the batch dimension so `Sam3Processor`/`Sam3Model` treat it as
+    `len(text_prompts)` independent (image, text) pairs, one call instead of
+    N sequential Python-level calls.
+
+    Returns a list of (mask_or_None, error_message_or_None) tuples, one per
+    entry in `text_prompts`, in the same order.
+    """
     model, processor = get_segmenter()
-    inputs = processor(images=image, text=text_prompt, return_tensors="pt").to(DEVICE)
-    with torch.no_grad():
+    images = [image] * len(text_prompts)
+    # dtype must match the (possibly fp16) model weights, or the forward
+    # pass will fail with a dtype-mismatch error on `pixel_values`.
+    inputs = processor(images=images, text=text_prompts, return_tensors="pt").to(DEVICE, dtype=SAM3_DTYPE)
+    with torch.inference_mode():
         outputs = model(**inputs)
 
     results = processor.post_process_instance_segmentation(
@@ -69,26 +89,35 @@ def segment(image: Image.Image, text_prompt: str, score_threshold: float = 0.5)
         threshold=score_threshold,
         mask_threshold=0.5,
         target_sizes=inputs.get("original_sizes").tolist(),
-    )[0]
+    )
 
-    masks = results["masks"]
-    scores = results["scores"]
+    per_prompt = []
+    for text_prompt, result in zip(text_prompts, results):
+        masks = result["masks"]
+        scores = result["scores"]
 
-    if len(masks) == 0:
-        raise gr.Error(f"No object found matching '{text_prompt}' above the confidence "
-                       f"threshold ({score_threshold}). Try a more specific or different phrase.")
+        if len(masks) == 0:
+            per_prompt.append((None,
+                f"No object found matching '{text_prompt}' above the confidence "
+                f"threshold ({score_threshold}). Try a more specific or different phrase."))
+            continue
 
-    best_idx = int(torch.argmax(scores))
-    mask = masks[best_idx]
-    if hasattr(mask, "cpu"):
-        mask = mask.cpu().numpy()
-    return np.asarray(mask).astype(bool)
+        best_idx = int(torch.argmax(scores))
+        mask = masks[best_idx]
+        if hasattr(mask, "cpu"):
+            mask = mask.cpu().numpy()
+        per_prompt.append((np.asarray(mask).astype(bool), None))
+    return per_prompt
 
 
 def erode_mask(mask: np.ndarray, pixels: int = 3) -> np.ndarray:
     if pixels <= 0:
         return mask
-    eroded = ndimage.binary_erosion(mask, iterations=pixels)
+    # cv2.erode is substantially faster than scipy.ndimage.binary_erosion for
+    # simple binary structuring-element erosion on 2D masks. A 3x3 cross
+    # kernel matches scipy's default 4-connected structuring element.
+    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
+    eroded = cv2.erode(mask.astype(np.uint8), kernel, iterations=pixels).astype(bool)
     return eroded if eroded.sum() > 20 else mask
 
 
@@ -136,8 +165,9 @@ def mask_overlay(image: Image.Image, masks: list) -> Image.Image:
     """Tints each object's mask a distinct color (cycling through
     OBJECT_COLORS if there are more objects than colors) for a quick
     visual sanity-check of what got segmented."""
-    arr = np.array(image).astype(np.float32)
-    overlay = arr.copy()
+    # np.array(image).astype(np.float32) already allocates a fresh array, so
+    # no extra .copy() is needed before mutating it in place.
+    overlay = np.array(image).astype(np.float32)
     for i, mask in enumerate(masks):
         color = np.array(OBJECT_COLORS[i % len(OBJECT_COLORS)])
         overlay[mask] = overlay[mask] * 0.4 + color * 0.6
@@ -187,6 +217,23 @@ def intrinsics_from_exif(image: Image.Image):
         return None
 
 
+def default_intrinsics(image: Image.Image, assumed_focal_35mm: float = 26.0) -> np.ndarray:
+    """Last-resort fallback intrinsics for when neither DA3 nor the image's
+    EXIF data provide any — e.g. screenshots, re-compressed/re-saved
+    photos, or images from sources that strip metadata. Assumes a ~26mm
+    35mm-equivalent focal length (typical of smartphone main cameras) and a
+    36mm-wide full-frame-equivalent sensor, same approximation as
+    `intrinsics_from_exif`. This is a coarse guess, not a calibration —
+    resulting distances should be treated as rough estimates rather than
+    precise measurements, but it lets the pipeline still produce a result
+    instead of hard-failing."""
+    w, h = image.size
+    fx = (assumed_focal_35mm / 36.0) * w
+    fy = fx  # assume square pixels
+    cx, cy = w / 2.0, h / 2.0
+    return np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64)
+
+
 # --------------------------------------------------------------------------
 # Addition 3: automatic scale calibration via an ArUco marker of known
 # physical size, instead of requiring the user to type in a measured
@@ -271,17 +318,17 @@ def run_pipeline(files, objects_text, depth_model_id, erosion_px,
 
     primary_image = images[0]
 
-    # Segment each object independently. A bad prompt for one object
-    # shouldn't discard valid results for the others, so failures are
-    # collected and reported rather than raised immediately.
+    # Segment every object in a single batched SAM3 forward pass (one image,
+    # N text prompts) instead of one full forward pass per object. A bad
+    # prompt for one object shouldn't discard valid results for the others,
+    # so failures are collected and reported rather than raised immediately.
     masks, valid_names, seg_warnings = [], [], []
-    for name in object_names:
-        try:
-            mask = erode_mask(segment(primary_image, name), erosion_px)
-            masks.append(mask)
-            valid_names.append(name)
-        except gr.Error as e:
-            seg_warnings.append(f"'{name}': {e}")
+    for name, (mask, err) in zip(object_names, segment_batch(primary_image, object_names)):
+        if err is not None:
+            seg_warnings.append(f"'{name}': {err}")
+            continue
+        masks.append(erode_mask(mask, erosion_px))
+        valid_names.append(name)
 
     for w in seg_warnings:
         gr.Warning(f"Skipped {w}")
@@ -340,8 +387,10 @@ def run_pipeline(files, objects_text, depth_model_id, erosion_px,
 
     # Addition 2: prefer EXIF-derived intrinsics over DA3's estimated ones
     # when available and requested — a known camera beats a network guess.
-    # Intrinsics are NOT optional (3D back-projection needs them), so if DA3
-    # didn't return them, EXIF becomes required rather than just preferred.
+    # Intrinsics are required for 3D back-projection; if DA3 didn't return
+    # them, EXIF is tried next, and if that's also unavailable we fall back
+    # to a rough default assumption (see default_intrinsics()) rather than
+    # failing the whole request.
     intrinsics = prediction.intrinsics[0] if prediction.intrinsics is not None else None
     intrinsics_source = "DA3 (estimated)"
     if use_exif_intrinsics or intrinsics is None:
@@ -350,11 +399,20 @@ def run_pipeline(files, objects_text, depth_model_id, erosion_px,
             intrinsics = exif_intrinsics
             intrinsics_source = "EXIF (35mm-equivalent focal length)"
         elif intrinsics is None:
-            raise gr.Error(
-                "Camera intrinsics were unavailable from both DA3 and the image's EXIF data, "
-                "so 3D positions can't be computed. Try a photo with intact EXIF metadata "
-                "(avoid re-saving/re-compressing it, which often strips EXIF), or check the "
-                "Space logs for [diag] lines to see what DA3 actually returned."
+            # Neither DA3 nor EXIF gave us intrinsics (e.g. a screenshot or a
+            # re-compressed image with stripped metadata). Rather than hard-
+            # failing the whole pipeline, degrade gracefully to a rough
+            # default focal-length assumption so the user still gets a
+            # result, just flagged as approximate.
+            intrinsics = default_intrinsics(primary_image)
+            intrinsics_source = "default estimate (~26mm-equivalent focal length assumption)"
+            gr.Warning(
+                "Camera intrinsics were unavailable from both DA3 and the image's EXIF data "
+                "(e.g. no EXIF on this image). Falling back to a default focal-length "
+                "assumption (~26mm-equivalent, typical smartphone camera) — treat the "
+                "resulting distances as rough estimates rather than precise measurements. "
+                "For better accuracy, provide a photo with intact EXIF metadata (avoid "
+                "re-saving/re-compressing it, which often strips EXIF)."
             )
         else:
             gr.Warning("No usable focal-length EXIF tag found on the primary image — "