Add PXDepth demo

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,17 @@
 ---
-title: PXDepth Demo
-emoji: 👀
-colorFrom: red
-colorTo: gray
+title: PXDepth
+emoji: 🌐
+colorFrom: indigo
+colorTo: blue
 sdk: gradio
-sdk_version: 6.25.0
-python_version: '3.12'
+sdk_version: 6.22.0
+python_version: 3.10.13
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PXDepth Demo
+
+Interactive monocular depth and point-cloud demo for [PXDepth](https://github.com/yuanzhy29/PXDepth).
+
+The Space downloads the released PXDepth and MoGe-2 checkpoints automatically on first startup.

app.py

@@ -0,0 +1,323 @@
+"""Hugging Face Gradio Space for PXDepth."""
+
+from __future__ import annotations
+
+import shutil
+import tempfile
+import time
+from pathlib import Path
+from typing import Optional
+
+# ZeroGPU patches torch during import, so spaces must be imported first.
+try:
+    import spaces
+
+    gpu = spaces.GPU(duration=90)
+except ImportError:
+    gpu = lambda fn: fn
+
+import gradio as gr
+import numpy as np
+import torch
+import torch.nn.functional as F
+import utils3d
+from PIL import Image
+
+from pxdepth.inference import area_size_from_area, resize_image, resize_map
+from pxdepth.model import PXDepth
+from pxdepth.utils.ply import write_point_cloud_ply
+from pxdepth.utils.vis import colorize_depth
+
+
+PXDEPTH_REPO = "yuanzhy29/PXDepth"
+MOGE2_REPO = "Ruicheng/moge-2-vitl-normal"
+PXDEPTH_SIZE = (1022, 770)
+MOGE2_TOKEN_AREA = 1200
+MOGE2_PATCH_SIZE = 14
+MAX_INPUT_PIXELS = 12_000_000
+OUTPUT_MAX_AGE = 60 * 60
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+CSS = """
+#pxdepth-demo { max-width: 1280px; margin: 0 auto; }
+#img-display-input, #img-display-output { max-height: 72vh; }
+#img-display-output img { object-fit: contain !important; }
+#model-3d { min-height: 60vh; }
+"""
+
+
+def load_model() -> PXDepth:
+    """Load PXDepth and its MoGe-2 metric-scale reference once at startup."""
+    print("Loading PXDepth...")
+    model = PXDepth.from_pretrained(PXDEPTH_REPO, strict=True).eval()
+
+    try:
+        from moge.model.v2 import MoGeModel
+    except ImportError as exc:
+        raise RuntimeError(
+            "MoGe-2 is required by this demo. Check the Space requirements."
+        ) from exc
+
+    print("Loading MoGe-2...")
+    model._reference_model = MoGeModel.from_pretrained(MOGE2_REPO).eval()
+    model = model.to(DEVICE).eval()
+    print(f"Models loaded on {DEVICE}.")
+    return model
+
+
+MODEL = load_model()
+
+
+def resize_for_tokens(image: torch.Tensor, tokens: int, patch: int) -> torch.Tensor:
+    """Preserve aspect ratio and resize an RGB tensor to a patch-token area."""
+    height, width = area_size_from_area(
+        image.shape[-2],
+        image.shape[-1],
+        tokens * patch * patch,
+        patch,
+    )
+    if (height, width) == tuple(image.shape[-2:]):
+        return image
+    return F.interpolate(
+        image.unsqueeze(0),
+        (height, width),
+        mode="bilinear",
+        align_corners=False,
+    )[0]
+
+
+def cleanup_outputs(root: Path) -> None:
+    """Remove stale per-session files from the Space's ephemeral storage."""
+    if not root.exists():
+        return
+    cutoff = time.time() - OUTPUT_MAX_AGE
+    for path in root.iterdir():
+        try:
+            if path.is_dir() and path.stat().st_mtime < cutoff:
+                shutil.rmtree(path, ignore_errors=True)
+        except OSError:
+            continue
+
+
+def session_dir(request: Optional[gr.Request]) -> Path:
+    """Create a clean output directory for the current browser session."""
+    session = getattr(request, "session_hash", None) or "local"
+    session = "".join(char for char in session if char.isalnum() or char in "-_")
+    root = Path(tempfile.gettempdir()) / "pxdepth-demo"
+    root.mkdir(parents=True, exist_ok=True)
+    cleanup_outputs(root)
+
+    output = root / (session or "local")
+    shutil.rmtree(output, ignore_errors=True)
+    output.mkdir(parents=True, exist_ok=True)
+    return output
+
+
+def sample_points(
+    points: np.ndarray,
+    colors: np.ndarray,
+    max_points: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Deterministically subsample a point cloud for browser rendering."""
+    if points.shape[0] <= max_points:
+        return points, colors
+    indices = np.linspace(0, points.shape[0] - 1, max_points, dtype=np.int64)
+    return points[indices], colors[indices]
+
+
+@gpu
+@torch.inference_mode()
+def predict_gpu(image: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Run only model inference while holding the ZeroGPU allocation."""
+    tensor = (
+        torch.from_numpy(image.copy())
+        .to(device=DEVICE, dtype=torch.float32)
+        .permute(2, 0, 1)
+        / 255.0
+    )
+    model_image, _ = resize_image(
+        tensor,
+        PXDEPTH_SIZE,
+        True,
+        MODEL.patch_size,
+    )
+    reference_image = resize_for_tokens(
+        tensor,
+        MOGE2_TOKEN_AREA,
+        MOGE2_PATCH_SIZE,
+    )
+
+    result = MODEL.infer(
+        model_image,
+        ref_image=reference_image,
+        apply_mask=False,
+        use_fp16=DEVICE.type == "cuda",
+        use_fp32=DEVICE.type != "cuda",
+    )
+    return (
+        result["depth"].float().cpu().numpy(),
+        result["mask"].cpu().numpy(),
+        result["intrinsics"].float().cpu().numpy(),
+    )
+
+
+def on_submit(
+    image: Optional[np.ndarray],
+    max_points: int,
+    apply_mask: bool,
+    request: gr.Request,
+):
+    """Run inference, build visualizations, and export downloadable files."""
+    if image is None:
+        raise gr.Error("Please upload an image first.")
+    if image.ndim != 3 or image.shape[-1] < 3:
+        raise gr.Error("The input must be an RGB image.")
+    if image.shape[0] * image.shape[1] > MAX_INPUT_PIXELS:
+        raise gr.Error(
+            "The uploaded image is too large. Please use an image below 12 megapixels."
+        )
+
+    image = np.ascontiguousarray(image[..., :3].astype(np.uint8))
+    original_size = image.shape[:2]
+    depth_raw, mask_raw, intrinsics_np = predict_gpu(image)
+
+    # Restore outputs and reconstruct the point map on CPU so ZeroGPU is held
+    # only for neural-network inference.
+    depth = resize_map(torch.from_numpy(depth_raw), original_size).float()
+    mask = resize_map(torch.from_numpy(mask_raw), original_size, is_mask=True)
+    intrinsics = torch.from_numpy(intrinsics_np).float()
+    finite = torch.isfinite(depth) & (depth > 0)
+    valid = finite & mask if apply_mask else finite
+    points = utils3d.pt.depth_map_to_point_map(
+        torch.where(finite, depth, torch.zeros_like(depth)),
+        intrinsics=intrinsics,
+    )
+
+    depth_np = depth.numpy().astype(np.float32)
+    mask_np = mask.numpy().astype(bool)
+    valid_np = valid.numpy().astype(bool)
+    depth_vis = colorize_depth(np.where(mask_np, depth_np, np.inf), mask=None)
+
+    output = session_dir(request)
+    depth_npy = output / "metric_depth.npy"
+    depth_png = output / "depth.png"
+    mask_png = output / "mask.png"
+    ply_path = output / "pointcloud.ply"
+    glb_path = output / "pointcloud_viewer.glb"
+
+    np.save(depth_npy, depth_np)
+    Image.fromarray(depth_vis).save(depth_png)
+    Image.fromarray(mask_np.astype(np.uint8) * 255, mode="L").save(mask_png)
+
+    points_np = points.numpy().reshape(-1, 3)
+    colors_np = image.reshape(-1, 3).astype(np.float32) / 255.0
+    keep = valid_np.reshape(-1) & np.isfinite(points_np).all(axis=1)
+    points_full, colors_full = points_np[keep], colors_np[keep]
+    if points_full.shape[0] == 0:
+        raise gr.Error("No valid 3D points were produced for this image.")
+    write_point_cloud_ply(ply_path, points_full, colors_full)
+
+    import trimesh
+
+    viewer_points, viewer_colors = sample_points(
+        points_full,
+        colors_full,
+        int(max_points),
+    )
+    viewer_points = viewer_points * np.array([1.0, -1.0, -1.0], np.float32)
+    viewer_colors = np.clip(viewer_colors * 255.0, 0, 255).astype(np.uint8)
+    trimesh.PointCloud(viewer_points, colors=viewer_colors).export(glb_path)
+
+    files = [str(depth_png), str(depth_npy), str(mask_png), str(ply_path)]
+    return (image, depth_vis), str(glb_path), files
+
+
+def build_demo() -> gr.Blocks:
+    """Construct the public Gradio interface."""
+    description = """
+Official demo for **PXDepth: Pixel-Space Modeling for Structure Preserving Monocular Depth Estimation**.  
+See the [paper](https://arxiv.org/abs/2608.16984),
+[project page](https://yuanzhy29.github.io/PXDepth-Page/), and
+[GitHub repository](https://github.com/yuanzhy29/PXDepth).
+"""
+    with gr.Blocks(theme=gr.themes.Soft(), css=CSS) as demo:
+        with gr.Column(elem_id="pxdepth-demo"):
+            gr.Markdown("# PXDepth")
+            gr.Markdown(description)
+            gr.Markdown("### Point Cloud & Depth Prediction Demo")
+
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(
+                        label="Input Image",
+                        image_mode="RGB",
+                        type="numpy",
+                        elem_id="img-display-input",
+                    )
+                    with gr.Accordion(label="Settings", open=False):
+                        max_points = gr.Slider(
+                            50_000,
+                            500_000,
+                            value=200_000,
+                            step=50_000,
+                            label="3D Viewer Max Points",
+                            info="The downloaded PLY retains all valid points.",
+                        )
+                        apply_mask = gr.Checkbox(
+                            label="Apply valid-depth mask to point cloud",
+                            value=True,
+                        )
+                    submit = gr.Button("Predict", variant="primary")
+
+                with gr.Column():
+                    with gr.Tabs():
+                        with gr.Tab("3D View"):
+                            model_3d = gr.Model3D(
+                                label="3D Point Map",
+                                clear_color=(1.0, 1.0, 1.0, 1.0),
+                                height="60vh",
+                                elem_id="model-3d",
+                            )
+                        with gr.Tab("Depth"):
+                            depth_map = gr.ImageSlider(
+                                label="RGB / Metric Depth",
+                                image_mode="RGB",
+                                type="numpy",
+                                slider_position=50,
+                                elem_id="img-display-output",
+                            )
+                        with gr.Tab("Download"):
+                            downloads = gr.File(
+                                label="Download Files",
+                                file_count="multiple",
+                                type="filepath",
+                            )
+
+            examples = Path("example_images")
+            example_files = (
+                sorted(
+                    str(path)
+                    for path in examples.iterdir()
+                    if path.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}
+                )
+                if examples.exists()
+                else []
+            )
+            if example_files:
+                gr.Examples(example_files, input_image, cache_examples=False)
+
+            submit.click(
+                on_submit,
+                [input_image, max_points, apply_mask],
+                [depth_map, model_3d, downloads],
+                show_progress="full",
+                concurrency_limit=1,
+            )
+    return demo
+
+
+demo = build_demo()
+
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1).launch()

pxdepth/__init__.py

@@ -0,0 +1,18 @@
+"""Top-level public API for PXDepth monocular depth estimation.
+
+Importing this package exposes the :class:`PXDepth` model without pulling
+evaluation entry points into user code. The model accepts RGB
+tensors and returns normalized depth plus a finite-depth probability map.
+"""
+
+from .build import build_model
+from .model import PXDepth
+from .registry import ENCODERS, MODELS, PREDICTORS
+
+__all__ = [
+    "PXDepth",
+    "build_model",
+    "MODELS",
+    "ENCODERS",
+    "PREDICTORS",
+]

pxdepth/build.py

@@ -0,0 +1,26 @@
+"""Public construction helpers for config-driven PXDepth components."""
+
+from typing import Any, Dict
+
+import torch.nn as nn
+
+from .registry import MODELS
+
+# Import built-in modules once so their registration decorators run. External
+# components are imported by ``load_config`` before these builders are called.
+from . import model as _model  # noqa: F401,E402
+
+
+def build_model(config: Dict[str, Any]) -> nn.Module:
+    """Build a model from its JSON-compatible configuration.
+
+    Args:
+        config: Model dictionary. ``type`` defaults to ``PXDepth`` for old
+            public configs and checkpoints.
+
+    Returns:
+        Constructed ``torch.nn.Module`` registered in :data:`MODELS`.
+    """
+    config = dict(config)
+    config.setdefault("type", "PXDepth")
+    return MODELS.build(config)

pxdepth/config.py

@@ -0,0 +1,137 @@
+"""Load, compose, expand, and validate PXDepth JSON configurations.
+
+The loader deliberately stays JSON based.  It adds only three conveniences that
+are useful to external users: optional ``_base_`` composition, environment
+variable expansion in strings, and optional module imports for custom registry
+entries.  There is no framework-specific config object or runtime magic.
+"""
+
+import json
+import os
+from copy import deepcopy
+from importlib import import_module
+from pathlib import Path
+from typing import Any, Dict, Iterable
+
+
+def _merge(base: Any, update: Any) -> Any:
+    """Recursively merge one configuration value into another.
+
+    Args:
+        base: Existing value inherited from a base config.
+        update: Value from the child config. Dictionaries merge recursively,
+            while lists and scalar values replace ``base`` completely.
+
+    Returns:
+        A deep-copied merged value. Neither input object is mutated.
+    """
+    if not isinstance(base, dict) or not isinstance(update, dict):
+        return deepcopy(update)
+    result = deepcopy(base)
+    for key, value in update.items():
+        result[key] = _merge(result[key], value) if key in result else deepcopy(value)
+    return result
+
+
+def _expand(value: Any) -> Any:
+    """Expand filesystem shorthand throughout a nested config structure.
+
+    Args:
+        value: Arbitrarily nested dictionaries, lists, strings, and scalar
+            values loaded from JSON.
+
+    Returns:
+        A matching nested structure where every string has environment
+        variables and a leading ``~`` expanded. Non-string values are retained.
+    """
+    if isinstance(value, dict):
+        return {key: _expand(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_expand(item) for item in value]
+    if isinstance(value, str):
+        return os.path.expanduser(os.path.expandvars(value))
+    return value
+
+
+def _read(path: Path, stack: tuple[Path, ...] = ()) -> Dict[str, Any]:
+    """Read one JSON file and recursively compose optional base files.
+
+    Args:
+        path: JSON file to load. Relative ``_base_`` entries resolve beside it.
+        stack: Internal chain of resolved paths used to detect inheritance
+            cycles. Callers normally leave this empty.
+
+    Returns:
+        Merged plain dictionary before string expansion and validation.
+
+    Raises:
+        ValueError: If configs form a circular ``_base_`` dependency.
+    """
+    path = path.resolve()
+    if path in stack:
+        chain = " -> ".join(str(item) for item in (*stack, path))
+        raise ValueError(f"Circular config inheritance: {chain}")
+    config = json.loads(path.read_text())
+    bases = config.pop("_base_", [])
+    if isinstance(bases, str):
+        bases = [bases]
+    merged: Dict[str, Any] = {}
+    for base in bases:
+        base_path = Path(os.path.expanduser(os.path.expandvars(str(base))))
+        if not base_path.is_absolute():
+            base_path = path.parent / base_path
+        merged = _merge(merged, _read(base_path, (*stack, path)))
+    return _merge(merged, config)
+
+
+def import_modules(names: Iterable[str]) -> None:
+    """Import extension modules so their registry decorators execute.
+
+    Args:
+        names: Iterable of importable Python module names, such as
+            ``my_project.components``.
+
+    Returns:
+        ``None``. Imports are performed for their registration side effects.
+    """
+    for name in names:
+        import_module(str(name))
+
+
+def validate_config(config: Dict[str, Any], kind: str | None = None) -> None:
+    """Fail early for missing or structurally invalid public configuration fields.
+
+    Args:
+        config: Fully composed configuration dictionary.
+        kind: Optional ``'eval'`` validation profile.
+
+    Returns:
+        ``None``. A descriptive ``ValueError`` is raised for invalid structure.
+    """
+    if not isinstance(config, dict):
+        raise ValueError("The config root must be a JSON object.")
+    if kind not in {None, "eval"}:
+        raise ValueError(f"Unsupported config kind: {kind!r}")
+    if kind == "eval" and not config:
+        raise ValueError("Evaluation config must contain at least one benchmark.")
+
+
+def load_config(path: str | Path, kind: str | None = None) -> Dict[str, Any]:
+    """Load a resolved config and import optional external extension modules.
+
+    Args:
+        path: JSON config path. Relative ``_base_`` paths resolve beside this file.
+        kind: Optional validation profile passed to :func:`validate_config`.
+
+    Returns:
+        Plain nested dictionaries/lists suitable for JSON serialization. The
+        optional top-level ``imports`` list is retained for external registry
+        extensions.
+    """
+    config = _expand(_read(Path(path)))
+    imports = config.get("imports", [])
+    if isinstance(imports, str):
+        imports = [imports]
+    import_modules(imports)
+    validate_config(config, kind=kind)
+    return config

pxdepth/evaluation/__init__.py

@@ -0,0 +1,12 @@
+"""Public benchmark-evaluation API for PXDepth.
+
+The exported loader produces geometry-aware evaluation samples and the metric
+function aligns raw predictions before computing depth, point-cloud, camera,
+local-structure, and optional boundary measurements. Command-line orchestration
+is kept in ``scripts/eval.py``.
+"""
+
+from .dataloader import EvalDataLoaderPipeline
+from .metrics import compute_metrics
+
+__all__ = ["EvalDataLoaderPipeline", "compute_metrics"]

pxdepth/evaluation/dataloader.py

@@ -0,0 +1,675 @@
+"""Asynchronous RGB-depth benchmark loader with geometry-aware resizing.
+
+Evaluation samples follow the processed benchmark directory contract. This
+module loads RGB, depth, normalized intrinsics, and optional segmentation,
+applies the benchmark-configured view transformation, and returns aligned
+PyTorch tensors plus an organized ground-truth point map.
+"""
+
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+import cv2
+import utils3d
+import pipeline
+
+from ..utils.io import read_depth, read_image, read_json, read_segmentation
+from ..utils.data_augmentation import sample_perspective, warp_perspective, resize_to_cover_center_crop_view
+
+
+def _resolve_image_path(instance_path: Union[str, Path]) -> Path:
+    """Resolve a processed sample's RGB path with PNG precedence.
+
+    Args:
+        instance_path: Directory containing one processed benchmark sample.
+
+    Returns:
+        ``image.png`` when present, otherwise ``image.jpg``.
+    """
+    instance_path = Path(instance_path)
+    image_png = instance_path / 'image.png'
+    if image_png.exists():
+        return image_png
+    return instance_path / 'image.jpg'
+
+
+def _resize_to_cover_center_crop_mask(mask: np.ndarray, raw_width: int, raw_height: int, tgt_width: int, tgt_height: int) -> np.ndarray:
+    """Apply resize-to-cover and center crop to a discrete label mask.
+
+    Args:
+        mask: Integer/boolean source mask ``[H_raw,W_raw]``.
+        raw_width: Source image width.
+        raw_height: Source image height.
+        tgt_width: Output width.
+        tgt_height: Output height.
+
+    Returns:
+        Nearest-resized and center-cropped mask ``[H_tgt,W_tgt]``.
+    """
+    scale = max(tgt_width / raw_width, tgt_height / raw_height)
+    resized_width = max(tgt_width, int(round(raw_width * scale)))
+    resized_height = max(tgt_height, int(round(raw_height * scale)))
+    resized_mask = cv2.resize(
+        mask.astype(np.uint8),
+        (resized_width, resized_height),
+        interpolation=cv2.INTER_NEAREST,
+    )
+    x0 = max(0, (resized_width - tgt_width) // 2)
+    y0 = max(0, (resized_height - tgt_height) // 2)
+    return resized_mask[y0:y0 + tgt_height, x0:x0 + tgt_width].copy()
+
+
+def _intrinsics_normalized_to_pixel(intrinsics: np.ndarray, width: int, height: int) -> np.ndarray:
+    """Convert normalized camera intrinsics to pixel coordinates.
+
+    Args:
+        intrinsics: Floating camera matrix ``[3,3]`` normalized by image size.
+        width: Image width used to scale the first matrix row.
+        height: Image height used to scale the second matrix row.
+
+    Returns:
+        Pixel-space ``float32 [3,3]`` camera matrix.
+    """
+    intrinsics_px = intrinsics.astype(np.float32).copy()
+    intrinsics_px[0, :] *= float(width)
+    intrinsics_px[1, :] *= float(height)
+    intrinsics_px[2, 0] = 0.0
+    intrinsics_px[2, 1] = 0.0
+    intrinsics_px[2, 2] = 1.0
+    return intrinsics_px
+
+
+def _intrinsics_pixel_to_normalized(intrinsics_px: np.ndarray, width: int, height: int) -> np.ndarray:
+    """Convert pixel camera intrinsics to normalized coordinates.
+
+    Args:
+        intrinsics_px: Pixel-space camera matrix ``[3,3]``.
+        width: Image width used to normalize the first matrix row.
+        height: Image height used to normalize the second matrix row.
+
+    Returns:
+        Normalized ``float32 [3,3]`` camera matrix.
+    """
+    intrinsics = intrinsics_px.astype(np.float32).copy()
+    intrinsics[0, :] /= float(width)
+    intrinsics[1, :] /= float(height)
+    intrinsics[2, 0] = 0.0
+    intrinsics[2, 1] = 0.0
+    intrinsics[2, 2] = 1.0
+    return intrinsics
+
+
+def _resize_depth_nearest_preserve_nan(depth: np.ndarray, size: Tuple[int, int]) -> np.ndarray:
+    """Nearest-resize positive finite depth while preserving invalid support.
+
+    Args:
+        depth: Source depth ``float [H,W]`` with NaN/Inf invalid values.
+        size: OpenCV target tuple ``(width,height)``.
+
+    Returns:
+        ``float32 [height,width]`` depth. Pixels whose nearest source was invalid
+        are represented by NaN.
+    """
+    width, height = size
+    valid = np.isfinite(depth) & (depth > 0)
+    resized_depth = cv2.resize(
+        np.where(valid, depth, 0.0).astype(np.float32),
+        (width, height),
+        interpolation=cv2.INTER_NEAREST,
+    )
+    resized_valid = cv2.resize(
+        valid.astype(np.uint8),
+        (width, height),
+        interpolation=cv2.INTER_NEAREST,
+    ).astype(bool)
+    return np.where(resized_valid, resized_depth, np.nan).astype(np.float32)
+
+
+def _mda_boundary_view(
+    image: np.ndarray,
+    depth: np.ndarray,
+    intrinsics: np.ndarray,
+    target_size: Tuple[int, int],
+    segmentation_mask: Optional[np.ndarray] = None,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
+    """Create the principal-point-centered view used by boundary benchmarks.
+
+    The image is first cropped symmetrically around the principal point, then
+    resized to cover the target and center-cropped. RGB uses Lanczos for
+    downsampling and bicubic for upsampling; depth and segmentation use nearest
+    interpolation. Pixel intrinsics are updated after every crop and resize.
+
+    Args:
+        image: RGB uint8 array ``[H,W,3]``.
+        depth: Depth array ``[H,W]`` with NaN invalid values.
+        intrinsics: Normalized camera matrix ``[3,3]``.
+        target_size: Output ``(height,width)``.
+        segmentation_mask: Optional integer labels ``[H,W]``.
+
+    Returns:
+        image: Transformed RGB ``[H_t,W_t,3]``.
+        depth: Transformed depth ``float32 [H_t,W_t]``.
+        intrinsics: Updated normalized matrix ``float32 [3,3]``.
+        segmentation_mask: Transformed labels ``[H_t,W_t]`` or ``None``.
+    """
+    tgt_height, tgt_width = target_size
+    raw_height, raw_width = image.shape[:2]
+    intrinsics_px = _intrinsics_normalized_to_pixel(intrinsics, raw_width, raw_height)
+
+    cx = float(intrinsics_px[0, 2])
+    cy = float(intrinsics_px[1, 2])
+    margin_x = max(1.0, min(cx, raw_width - cx))
+    margin_y = max(1.0, min(cy, raw_height - cy))
+    crop_left = max(0, int(round(cx - margin_x)))
+    crop_right = min(raw_width, int(round(cx + margin_x)))
+    crop_top = max(0, int(round(cy - margin_y)))
+    crop_bottom = min(raw_height, int(round(cy + margin_y)))
+
+    if crop_right - crop_left < 2 or crop_bottom - crop_top < 2:
+        crop_left, crop_top = 0, 0
+        crop_right, crop_bottom = raw_width, raw_height
+
+    image = image[crop_top:crop_bottom, crop_left:crop_right].copy()
+    depth = depth[crop_top:crop_bottom, crop_left:crop_right].copy()
+    if segmentation_mask is not None:
+        segmentation_mask = segmentation_mask[crop_top:crop_bottom, crop_left:crop_right].copy()
+    intrinsics_px[0, 2] -= float(crop_left)
+    intrinsics_px[1, 2] -= float(crop_top)
+
+    crop_height, crop_width = image.shape[:2]
+    scale = max(tgt_width / crop_width, tgt_height / crop_height)
+    resized_width = max(tgt_width, int(np.floor(crop_width * scale)))
+    resized_height = max(tgt_height, int(np.floor(crop_height * scale)))
+    if resized_width < tgt_width or resized_height < tgt_height:
+        resized_width = max(tgt_width, int(np.ceil(crop_width * scale)))
+        resized_height = max(tgt_height, int(np.ceil(crop_height * scale)))
+
+    image_resample = Image.Resampling.LANCZOS if scale < 1.0 else Image.Resampling.BICUBIC
+    resized_image = np.array(Image.fromarray(image).resize((resized_width, resized_height), image_resample))
+    resized_depth = _resize_depth_nearest_preserve_nan(depth, (resized_width, resized_height))
+    resized_segmentation_mask = None
+    if segmentation_mask is not None:
+        resized_segmentation_mask = cv2.resize(
+            segmentation_mask,
+            (resized_width, resized_height),
+            interpolation=cv2.INTER_NEAREST,
+        )
+    intrinsics_px[:2, :] *= float(scale)
+
+    x0 = int(round((resized_width - tgt_width) * 0.5))
+    y0 = int(round((resized_height - tgt_height) * 0.5))
+    x0 = min(max(x0, 0), resized_width - tgt_width)
+    y0 = min(max(y0, 0), resized_height - tgt_height)
+    x1, y1 = x0 + tgt_width, y0 + tgt_height
+
+    tgt_image = resized_image[y0:y1, x0:x1].copy()
+    tgt_depth = resized_depth[y0:y1, x0:x1].copy().astype(np.float32)
+    tgt_segmentation_mask = None
+    if resized_segmentation_mask is not None:
+        tgt_segmentation_mask = resized_segmentation_mask[y0:y1, x0:x1].copy()
+    intrinsics_px[0, 2] -= float(x0)
+    intrinsics_px[1, 2] -= float(y0)
+    tgt_intrinsics = _intrinsics_pixel_to_normalized(intrinsics_px, tgt_width, tgt_height)
+
+    return tgt_image, tgt_depth, tgt_intrinsics, tgt_segmentation_mask
+
+
+class EvalDataLoaderPipeline:
+    """Asynchronously load and geometrically standardize one benchmark dataset.
+
+    The pipeline emits one sample at a time. It supports exact resolutions,
+    center-crop sizes, or aspect-preserving token budgets, and can optionally
+    include segmentation or normal annotations for local and boundary metrics.
+    """
+
+    def __init__(
+        self,
+        path: str,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        center_crop_size: Optional[int] = None,
+        split: int = '.index.txt',
+        drop_max_depth: float = 1000.,
+        num_load_workers: int = 4,
+        num_process_workers: int = 8,
+        include_segmentation: bool = False,
+        include_normal: bool = False,
+        depth_to_normal: bool = False,
+        max_segments: int = 100,
+        min_seg_area: int = 1000,
+        depth_unit: str = None,
+        min_depth: Optional[float] = None,
+        max_depth: Optional[float] = None,
+        has_sharp_boundary = False,
+        subset: int = None,
+        filenames: Optional[List[str]] = None,
+        num_tokens: Optional[int] = None,
+        patch_size: Optional[int] = None,
+        disable_augmentations: bool = True,
+        disable_perspective: bool = True,
+        resize_to_cover_center_crop: bool = False,
+        mda_boundary_transform: bool = False,
+    ):
+        """Configure benchmark indexing, transforms, and worker stages.
+
+        Args:
+            path: Processed benchmark root containing the split index.
+            width: Exact output width when token/crop modes are disabled.
+            height: Exact output height when token/crop modes are disabled.
+            center_crop_size: Optional square output side length.
+            split: Relative index filename under ``path``.
+            drop_max_depth: Relative dynamic-range multiplier used to suppress
+                extreme finite values after transformation.
+            num_load_workers: Number of parallel disk readers.
+            num_process_workers: Number of parallel geometry workers.
+            include_segmentation: Load ``segmentation.png`` and label metadata.
+            include_normal: Derive a normal map from depth.
+            depth_to_normal: Retained benchmark compatibility flag.
+            max_segments: Maximum segmentation labels retained by area.
+            min_seg_area: Minimum number of pixels for a retained label.
+            depth_unit: Optional scalar converting stored depth to metric units.
+            min_depth: Optional lower valid-depth bound in converted units.
+            max_depth: Optional upper valid-depth bound in converted units.
+            has_sharp_boundary: Mark samples for boundary metric computation.
+            subset: Optional number of leading index entries to evaluate.
+            filenames: Optional explicit relative paths replacing the split file.
+            num_tokens: Optional approximate image-token count.
+            patch_size: Patch divisibility used with token sampling.
+            disable_augmentations: Disable random flip/color transforms.
+            disable_perspective: Use identity perspective mapping.
+            resize_to_cover_center_crop: Resize to cover then center crop.
+            mda_boundary_transform: Use the principal-point-centered boundary
+                benchmark transformation.
+
+        Returns:
+            ``None``. Workers start when the context manager is entered.
+        """
+        if filenames is None:
+            filenames = Path(path).joinpath(split).read_text(encoding='utf-8').splitlines()
+        else:
+            filenames = list(filenames)
+        if subset is not None:
+            subset = int(subset)
+            if subset > 0:
+                filenames = filenames[:subset]
+        self.width = int(width) if width is not None else None
+        self.height = int(height) if height is not None else None
+        self.center_crop_size = int(center_crop_size) if center_crop_size is not None else None
+        self.drop_max_depth = drop_max_depth
+        self.path = Path(path)
+        self.filenames = filenames
+        self.include_segmentation = include_segmentation
+        self.include_normal = include_normal
+        self.max_segments = max_segments
+        self.min_seg_area = min_seg_area
+        self.depth_to_normal = depth_to_normal
+        self.depth_unit = depth_unit
+        self.min_depth = float(min_depth) if min_depth is not None else None
+        self.max_depth = float(max_depth) if max_depth is not None else None
+        self.has_sharp_boundary = has_sharp_boundary
+        self.num_tokens = int(num_tokens) if num_tokens is not None else None
+        self.patch_size = int(patch_size) if patch_size is not None else None
+        self.disable_augmentations = bool(disable_augmentations)
+        self.disable_perspective = bool(disable_perspective)
+        self.resize_to_cover_center_crop = bool(resize_to_cover_center_crop)
+        self.mda_boundary_transform = bool(mda_boundary_transform)
+
+        self.rng = np.random.default_rng(seed=0)
+
+        self.pipeline = pipeline.Sequential([
+            self._generator,
+            pipeline.Parallel([self._load_instance] * num_load_workers),
+            pipeline.Parallel([self._process_instance] * num_process_workers),
+            pipeline.Buffer(4)
+        ])
+
+    def __len__(self):
+        """Return the number of configured benchmark samples.
+
+        The value reflects explicit filenames and optional subset truncation.
+
+        Returns:
+            Integer length of the selected filename list.
+        """
+        return len(self.filenames)
+
+    def _resolve_target_size(self, raw_width: int, raw_height: int) -> Tuple[int, int]:
+        """Resolve output dimensions from exact, crop, or token settings.
+
+        Args:
+            raw_width: Source image width.
+            raw_height: Source image height.
+
+        Returns:
+            Integer tuple ``(target_width,target_height)``, optionally rounded to
+            patch multiples.
+        """
+        if self.num_tokens is not None:
+            if self.patch_size is None:
+                raise ValueError("patch_size must be set when using num_tokens.")
+            aspect_ratio = raw_width / raw_height
+            target_area = self.num_tokens * (self.patch_size ** 2)
+            tgt_width = int(round((target_area * aspect_ratio) ** 0.5))
+            tgt_height = int(round(tgt_width / aspect_ratio))
+        elif self.center_crop_size is not None:
+            tgt_width = self.center_crop_size
+            tgt_height = self.center_crop_size
+        else:
+            if self.width is None or self.height is None:
+                raise ValueError("width/height or center_crop_size must be set when num_tokens is not provided.")
+            tgt_width, tgt_height = self.width, self.height
+        if self.patch_size is not None:
+            tgt_width = max(self.patch_size, int(round(tgt_width / self.patch_size)) * self.patch_size)
+            tgt_height = max(self.patch_size, int(round(tgt_height / self.patch_size)) * self.patch_size)
+        return tgt_width, tgt_height
+
+    def _generator(self):
+        """Yield sequential sample indices to the asynchronous pipeline.
+
+        Disk loading and processing are parallelized after this ordered stage.
+
+        Yields:
+            Integer indices from zero through ``len(self)-1``.
+        """
+        for idx in range(len(self)):
+            yield idx
+
+    def _load_instance(self, idx):
+        """Read one indexed RGB-depth sample and optional segmentation.
+
+        Args:
+            idx: Integer index into ``self.filenames``.
+
+        Returns:
+            Dictionary containing RGB ``uint8 [H,W,3]``, depth ``float [H,W]``,
+            normalized intrinsics ``float32 [3,3]``, masks, and optional
+            segmentation; ``None`` for an out-of-range index.
+        """
+        if idx >= len(self.filenames):
+            return None
+
+        path = self.path.joinpath(self.filenames[idx])
+
+        instance = {
+            'filename': self.filenames[idx],
+        }
+        instance['image'] = read_image(_resolve_image_path(path))
+
+        depth = read_depth(Path(path, 'depth.png'))  # ignore depth unit from depth file, use config instead
+        instance.update({
+            'depth': depth,
+            'depth_mask': np.isfinite(depth) & (depth > 0),
+            'depth_mask_inf': np.isinf(depth),
+        })
+
+        if self.include_segmentation:
+            segmentation_mask, segmentation_labels = read_segmentation(Path(path,'segmentation.png'))
+            instance.update({
+                'segmentation_mask': segmentation_mask,
+                'segmentation_labels': segmentation_labels,
+            })
+
+        meta = read_json(Path(path, 'meta.json'))
+        instance['intrinsics'] = np.array(meta['intrinsics'], dtype=np.float32)
+
+        return instance
+
+    def _process_instance(self, instance: dict):
+        """Transform one loaded instance and build its ground-truth point map.
+
+        Args:
+            instance: Raw dictionary returned by :meth:`_load_instance`, or
+                ``None`` propagated from a failed/out-of-range stage.
+
+        Returns:
+            Processed dictionary with image ``float32 [3,H,W]``, depth and masks
+            ``[H,W]``, normalized intrinsics ``[3,3]``, point map ``[H,W,3]``,
+            metadata flags, and optional normals/segmentation tensors. Returns
+            ``None`` when the input is ``None``.
+        """
+        if instance is None:
+            return None
+
+        image = instance['image']
+        depth = instance['depth']
+        intrinsics = instance['intrinsics']
+        segmentation_mask = instance.get('segmentation_mask', None)
+        segmentation_labels = instance.get('segmentation_labels', None)
+
+        raw_height, raw_width = image.shape[:2]
+        tgt_width, tgt_height = self._resolve_target_size(raw_width, raw_height)
+        tgt_aspect = tgt_width / tgt_height
+
+        raw_depth_mask = np.isfinite(depth) & (depth > 0)
+        raw_depth_ratio = raw_depth_mask.mean()
+        if raw_depth_ratio < 0.001:
+            depth = np.ones_like(depth, dtype=np.float32)
+            raw_depth_mask = np.isfinite(depth)
+        else:
+            depth = np.where(raw_depth_mask, depth, np.nan)
+
+        if self.include_normal:
+            raw_normal, raw_normal_mask = utils3d.np.depth_map_to_normal_map(
+                depth, intrinsics=intrinsics, mask=raw_depth_mask, edge_threshold=88
+            )
+            raw_normal = np.where(raw_normal_mask[..., None], raw_normal, np.nan)
+        else:
+            raw_normal = None
+
+        if self.mda_boundary_transform:
+            tgt_image, tgt_depth, tgt_intrinsics, tgt_segmentation_mask = _mda_boundary_view(
+                image,
+                depth,
+                intrinsics,
+                (tgt_height, tgt_width),
+                segmentation_mask=segmentation_mask,
+            )
+            if self.include_normal:
+                tgt_normal, tgt_normal_mask = utils3d.np.depth_map_to_normal_map(
+                    tgt_depth, intrinsics=tgt_intrinsics, mask=np.isfinite(tgt_depth) & (tgt_depth > 0), edge_threshold=88
+                )
+                tgt_normal = np.where(tgt_normal_mask[..., None], tgt_normal, np.nan)
+            else:
+                tgt_normal = None
+        elif self.resize_to_cover_center_crop:
+            tgt_image, tgt_depth, tgt_intrinsics = resize_to_cover_center_crop_view(
+                image,
+                depth,
+                intrinsics,
+                (tgt_height, tgt_width),
+                image_interpolation='lanczos',
+                depth_interpolation='mixed',
+            )
+            if self.include_normal:
+                tgt_normal, tgt_normal_mask = utils3d.np.depth_map_to_normal_map(
+                    tgt_depth, intrinsics=tgt_intrinsics, mask=np.isfinite(tgt_depth) & (tgt_depth > 0), edge_threshold=88
+                )
+                tgt_normal = np.where(tgt_normal_mask[..., None], tgt_normal, np.nan)
+            else:
+                tgt_normal = None
+            tgt_segmentation_mask = None
+            if segmentation_mask is not None:
+                tgt_segmentation_mask = _resize_to_cover_center_crop_mask(
+                    segmentation_mask,
+                    raw_width,
+                    raw_height,
+                    tgt_width,
+                    tgt_height,
+                )
+        elif self.disable_perspective:
+            tgt_intrinsics = intrinsics.copy()
+            R = np.eye(3, dtype=np.float32)
+            transform = np.eye(3, dtype=np.float32)
+        else:
+            tgt_intrinsics, R = sample_perspective(
+                intrinsics,
+                tgt_aspect=tgt_aspect,
+                center_augmentation=0.0,
+                fov_range_absolute=(1, 179),
+                fov_range_relative=(1.0, 1.0),
+                rng=self.rng,
+            )
+            transform = tgt_intrinsics @ R @ np.linalg.inv(intrinsics)
+
+        if not self.resize_to_cover_center_crop and not self.mda_boundary_transform:
+            tgt_image = warp_perspective(image, transform, (tgt_height, tgt_width), interpolation='lanczos')
+
+            depth_edge_mask = utils3d.np.depth_map_edge(depth, mask=raw_depth_mask, kernel_size=5, ltol=0.01)
+            depth_bilinear_mask = raw_depth_mask & ~depth_edge_mask
+            warped_depth_bilinear_mask = warp_perspective(
+                depth_bilinear_mask.astype(np.float32),
+                transform,
+                (tgt_height, tgt_width),
+                interpolation='bilinear',
+            )
+            warped_depth_nearest = warp_perspective(
+                depth,
+                transform,
+                (tgt_height, tgt_width),
+                interpolation='nearest',
+                sparse_mask=~np.isnan(depth),
+            )
+            warped_depth_bilinear = 1 / warp_perspective(
+                1 / depth,
+                transform,
+                (tgt_height, tgt_width),
+                interpolation='bilinear',
+            )
+            warped_depth = np.where(warped_depth_bilinear_mask == 1.0, warped_depth_bilinear, warped_depth_nearest)
+            tgt_uvhomo = np.concatenate(
+                [utils3d.np.uv_map((tgt_height, tgt_width)), np.ones((tgt_height, tgt_width, 1), dtype=np.float32)],
+                axis=-1,
+            )
+            tgt_depth = warped_depth / np.dot(tgt_uvhomo, np.linalg.inv(transform)[2, :])
+
+            if raw_normal is not None:
+                warped_normal = warp_perspective(raw_normal, transform, (tgt_height, tgt_width), interpolation='bilinear')
+                tgt_normal = warped_normal @ R.T
+            else:
+                tgt_normal = None
+
+            if segmentation_mask is not None:
+                tgt_segmentation_mask = warp_perspective(
+                    segmentation_mask, transform, (tgt_height, tgt_width), interpolation='nearest'
+                )
+            else:
+                tgt_segmentation_mask = None
+
+        if not self.disable_augmentations:
+            if self.rng.choice([True, False]):
+                tgt_image = np.flip(tgt_image, axis=1).copy()
+                tgt_depth = np.flip(tgt_depth, axis=1).copy()
+                if tgt_normal is not None:
+                    tgt_normal = np.flip(tgt_normal, axis=1).copy() * [-1, 1, 1]
+
+        if self.depth_unit is not None:
+            tgt_depth *= self.depth_unit
+            is_metric = True
+        else:
+            is_metric = False
+
+        depth_range_mask = np.isfinite(tgt_depth) & (tgt_depth > 0)
+        if self.min_depth is not None:
+            depth_range_mask &= tgt_depth >= self.min_depth
+        if self.max_depth is not None:
+            depth_range_mask &= tgt_depth <= self.max_depth
+        tgt_depth = np.where(depth_range_mask, tgt_depth, np.nan)
+
+        drop_max_depth = np.nanquantile(np.where(np.isfinite(tgt_depth), tgt_depth, np.nan), 0.01) * self.drop_max_depth
+        tgt_depth = np.where(np.isfinite(tgt_depth), np.clip(tgt_depth, 0, drop_max_depth), tgt_depth)
+
+        tgt_depth_mask_inf = np.isinf(tgt_depth)
+        tgt_depth_mask = np.isfinite(tgt_depth) & (tgt_depth > 0)
+        if not np.any(tgt_depth_mask):
+            tgt_depth_mask = np.ones_like(tgt_depth_mask)
+            tgt_depth = np.ones_like(tgt_depth)
+
+        tgt_points = utils3d.np.depth_map_to_point_map(tgt_depth, intrinsics=tgt_intrinsics)
+
+        if self.include_segmentation and tgt_segmentation_mask is not None:
+            for k in ['undefined', 'unannotated', 'background', 'sky']:
+                if k in segmentation_labels:
+                    del segmentation_labels[k]
+            seg_id2count = dict(zip(*np.unique(tgt_segmentation_mask, return_counts=True)))
+            sorted_labels = sorted(segmentation_labels.keys(), key=lambda x: seg_id2count.get(segmentation_labels[x], 0), reverse=True)
+            segmentation_labels = {
+                k: segmentation_labels[k]
+                for k in sorted_labels[:self.max_segments]
+                if seg_id2count.get(segmentation_labels[k], 0) >= self.min_seg_area
+            }
+
+        instance.update({
+            'image': torch.from_numpy(tgt_image.astype(np.float32) / 255.0).permute(2, 0, 1),
+            'depth': torch.from_numpy(tgt_depth).float(),
+            'depth_mask': torch.from_numpy(tgt_depth_mask).bool(),
+            'depth_mask_inf': torch.from_numpy(tgt_depth_mask_inf).bool(),
+            'intrinsics': torch.from_numpy(tgt_intrinsics).float(),
+            'points': torch.from_numpy(tgt_points).float(),
+            'segmentation_mask': torch.from_numpy(tgt_segmentation_mask).long() if tgt_segmentation_mask is not None else None,
+            'segmentation_labels': segmentation_labels,
+            'is_metric': is_metric,
+            'has_sharp_boundary': self.has_sharp_boundary,
+        })
+        if tgt_normal is not None:
+            instance['normal'] = torch.from_numpy(tgt_normal).float()
+
+        instance = {k: v for k, v in instance.items() if v is not None}
+
+        return instance
+
+    def start(self):
+        """Start asynchronous loader workers.
+
+        Call this before :meth:`get` when not using the context manager.
+
+        Returns:
+            ``None``.
+        """
+        self.pipeline.start()
+
+    def stop(self):
+        """Stop asynchronous loader workers and release resources.
+
+        Any prefetched samples are discarded by the pipeline implementation.
+
+        Returns:
+            ``None``.
+        """
+        self.pipeline.stop()
+
+    def __enter__(self):
+        """Start the pipeline and return it as a context-manager value.
+
+        This is equivalent to an explicit :meth:`start` call.
+
+        Returns:
+            This :class:`EvalDataLoaderPipeline` instance.
+        """
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Stop the pipeline when leaving its context.
+
+        Args:
+            exc_type: Exception class raised inside the context, if any.
+            exc_value: Exception instance raised inside the context, if any.
+            traceback: Associated traceback object, if any.
+
+        Returns:
+            ``None``; exceptions are not suppressed.
+        """
+        self.stop()
+
+    def get(self):
+        """Block until the next processed evaluation sample is available.
+
+        Worker-side exceptions are surfaced by the underlying pipeline call.
+
+        Returns:
+            Processed sample dictionary documented by :meth:`_process_instance`.
+        """
+        return self.pipeline.get()

pxdepth/evaluation/metrics.py

@@ -0,0 +1,616 @@
+"""Depth, point-cloud, local-structure, and boundary metrics.
+
+Raw model outputs are aligned with the same low-resolution robust affine
+procedures used by the MoGe evaluation protocol. Depth-space, log-depth-space,
+and disparity-space predictions are converted to positive depth before common
+metrics and point-cloud reconstruction are evaluated.
+"""
+
+from typing import Dict, Literal, Tuple, Union
+from numbers import Number
+
+import cv2
+import torch
+import numpy as np
+import utils3d
+
+from ..utils.alignment import (
+    align_affine_lstsq,
+    align_depth_affine,
+    align_points_scale_xyz_shift,
+)
+from ..utils.tools import key_average
+
+
+ALIGN_MIN_VALID_PIXELS = 16
+
+
+def rel_depth(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6):
+    """Compute mean absolute relative depth error.
+
+    Args:
+        pred: Positive predicted depths at selected pixels, tensor ``[N]``.
+        gt: Positive ground-truth depths at the same pixels, tensor ``[N]``.
+        eps: Denominator stabilizer for near-zero GT values.
+
+    Returns:
+        Python float containing ``mean(abs(pred-gt)/(gt+eps))``.
+    """
+    rel = (torch.abs(pred - gt) / (gt + eps)).mean()
+    return rel.item()
+
+
+def delta1_depth(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6):
+    """Compute the fraction of depth ratios below ``1.25``.
+
+    Args:
+        pred: Positive predicted depth tensor ``[N]``.
+        gt: Positive ground-truth depth tensor ``[N]``.
+        eps: Compatibility argument retained by the public metric API.
+
+    Returns:
+        Python float in ``[0,1]``; larger is better.
+    """
+    delta1 = (torch.maximum(gt / pred, pred / gt) < 1.25).float().mean()
+    return delta1.item()
+
+
+def rel_point(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6):
+    """Compute 3D endpoint error relative to GT camera-space radius.
+
+    Args:
+        pred: Predicted camera-space points ``[N,3]``.
+        gt: Corresponding ground-truth points ``[N,3]``.
+        eps: Stabilizer added to each GT point radius.
+
+    Returns:
+        Python float mean relative Euclidean point error.
+    """
+    dist_gt = torch.norm(gt, dim=-1)
+    dist_err = torch.norm(pred - gt, dim=-1)
+    rel = (dist_err / (dist_gt + eps)).mean()
+    return rel.item()
+
+
+def delta1_point(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6):
+    """Compute the MoGe point accuracy under a 25% radial tolerance.
+
+    Args:
+        pred: Predicted camera-space points ``[N,3]``.
+        gt: Corresponding ground-truth points ``[N,3]``.
+        eps: Compatibility argument retained by the metric API.
+
+    Returns:
+        Python float fraction whose 3D error is below 25% of the smaller
+        predicted/GT camera-space radius.
+    """
+    dist_pred = torch.norm(pred, dim=-1)
+    dist_gt = torch.norm(gt, dim=-1)
+    dist_err = torch.norm(pred - gt, dim=-1)
+
+    delta1 = (dist_err < 0.25 * torch.minimum(dist_gt, dist_pred)).float().mean()
+    return delta1.item()
+
+
+def rel_point_local(pred: torch.Tensor, gt: torch.Tensor, diameter: torch.Tensor):
+    """Normalize local 3D endpoint error by an object's GT diameter.
+
+    Args:
+        pred: Locally aligned predicted points ``[N,3]``.
+        gt: Ground-truth points ``[N,3]`` for the same region.
+        diameter: Scalar tensor containing the largest GT bounding-box extent.
+
+    Returns:
+        Python float mean error divided by ``diameter``.
+    """
+    dist_err = torch.norm(pred - gt, dim=-1)
+    rel = (dist_err / diameter).mean()
+    return rel.item()
+
+
+def delta1_point_local(pred: torch.Tensor, gt: torch.Tensor, diameter: torch.Tensor):
+    """Compute local point accuracy at one quarter of object diameter.
+
+    Args:
+        pred: Locally aligned predicted points ``[N,3]``.
+        gt: Ground-truth points ``[N,3]``.
+        diameter: Scalar GT region diameter.
+
+    Returns:
+        Python float fraction with Euclidean error below ``0.25*diameter``.
+    """
+    dist_err = torch.norm(pred - gt, dim=-1)
+    delta1 = (dist_err < 0.25 * diameter).float().mean()
+    return delta1.item()
+
+
+def _nan_boundary_metrics() -> Dict[str, float]:
+    """Create a complete boundary metric record for invalid edge samples.
+
+    A stable key set keeps aggregation and JSON schemas consistent.
+
+    Returns:
+        Dictionary whose boundary accuracy and Chamfer distance are both NaN.
+    """
+    return {
+        'acc': float('nan'),
+        'cd': float('nan'),
+    }
+
+
+def _mda_boundary_mask(gt_depth: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """Extract the Canny GT depth edge mask used for boundary evaluation.
+
+    Args:
+        gt_depth: Ground-truth depth map ``[H,W]`` in meters.
+        mask: Boolean valid-depth mask ``[H,W]``.
+
+    Returns:
+        Boolean edge tensor ``[H,W]`` on ``gt_depth.device``. Depth is clipped
+        to ``[0.1,65]`` meters and invalid support is dilated by a 2x2 kernel
+        before Canny thresholds 100/200 are applied.
+    """
+    depth_np = gt_depth.detach().float().cpu().numpy()
+    valid_np = mask.detach().cpu().numpy().astype(bool)
+    depth_np = np.nan_to_num(depth_np, nan=0.0, posinf=65.0, neginf=0.0)
+
+    depth_gt_clamp = np.clip(depth_np, 0.1, 65.0)
+    min_val = depth_gt_clamp.min()
+    max_val = depth_gt_clamp.max()
+    norm_depth = (depth_gt_clamp - min_val) / (max_val - min_val + 1e-5)
+    norm_depth = np.clip(norm_depth, 0.0, 1.0)
+    depth_uint8 = (norm_depth * 255).astype(np.uint8)
+
+    edge = cv2.Canny(depth_uint8, 100, 200) > 0.5
+    kernel = np.ones((2, 2), np.uint8)
+    valid_np = cv2.dilate(1 - valid_np.astype(np.uint8), kernel, iterations=1) < 0.5
+    edge = edge & valid_np
+    return torch.from_numpy(edge).to(device=gt_depth.device, dtype=torch.bool)
+
+
+def _as_o3d_point_cloud(points: np.ndarray):
+    """Convert an XYZ NumPy array to an Open3D point cloud.
+
+    Args:
+        points: Finite point array ``float [N,3]``.
+
+    Returns:
+        ``open3d.geometry.PointCloud`` containing the supplied XYZ positions.
+    """
+    import open3d as o3d
+
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points)
+    return pcd
+
+
+def boundary_edge_metrics(
+    pred_depth: torch.Tensor,
+    gt_depth: torch.Tensor,
+    mask: torch.Tensor,
+    intrinsics: torch.Tensor,
+    return_misc: bool = False,
+    edge_mode: Literal['mda'] = 'mda',
+) -> Union[Dict[str, float], Tuple[Dict[str, float], Dict[str, torch.Tensor]]]:
+    """Evaluate edge depth and 3D boundary point-cloud quality.
+
+    GT Canny edges select the boundary point clouds. The predicted cloud is
+    rigidly refined to GT with point-to-point ICP, then bidirectional
+    nearest-neighbor distances produce accuracy and symmetric Chamfer distance
+    in millimeters.
+
+    Args:
+        pred_depth: Globally aligned predicted depth ``[H,W]`` in meters.
+        gt_depth: Ground-truth depth ``[H,W]`` in meters.
+        mask: Boolean GT valid-depth mask ``[H,W]``.
+        intrinsics: Normalized camera matrix ``[3,3]``.
+        return_misc: Also return edge masks, aligned clouds, and ICP transform.
+        edge_mode: Boundary extraction protocol. The release supports ``'mda'``.
+
+    Returns:
+        metrics: Dictionary containing ``acc`` and ``cd`` in millimeters.
+        misc: Returned only when requested. Contains ``edge_mask`` ``[H,W]``,
+            edge point arrays ``[N,3]``, and ``icp_transform`` ``[4,4]``.
+    """
+    from scipy.spatial import cKDTree as KDTree
+    import open3d as o3d
+
+    metrics = _nan_boundary_metrics()
+    misc: Dict[str, torch.Tensor] = {}
+
+    def _finish():
+        """Package the current metric state according to ``return_misc``.
+
+        The closure captures the partially populated dictionaries by reference.
+
+        Returns:
+            Metrics dictionary alone, or ``(metrics,misc)`` when requested.
+        """
+        return (metrics, misc) if return_misc else metrics
+
+    valid = mask & torch.isfinite(gt_depth) & (gt_depth > 0)
+    pred_valid = torch.isfinite(pred_depth) & (pred_depth > 0)
+    if edge_mode == 'mda':
+        edge = _mda_boundary_mask(gt_depth, valid)
+    else:
+        raise ValueError(f"Unknown boundary edge mode: {edge_mode}")
+    gt_edge_mask = edge & valid
+    pred_edge_mask = gt_edge_mask & pred_valid
+    if return_misc:
+        misc['edge_mask'] = edge
+    if gt_edge_mask.sum().item() < 10 or pred_edge_mask.sum().item() < 10:
+        return _finish()
+
+    pred_depth_clean = pred_depth.float().clone()
+    pred_depth_clean[~pred_valid] = 1.0
+    gt_depth_clean = gt_depth.float().clone()
+    gt_depth_clean[~valid] = 1.0
+
+    pred_points_full = utils3d.pt.depth_map_to_point_map(pred_depth_clean, intrinsics=intrinsics)
+    gt_points_full = utils3d.pt.depth_map_to_point_map(gt_depth_clean, intrinsics=intrinsics)
+
+    pred_points = pred_points_full[pred_edge_mask].detach().float().cpu().numpy()
+    gt_points = gt_points_full[gt_edge_mask].detach().float().cpu().numpy()
+    pred_points = pred_points[np.isfinite(pred_points).all(axis=1)]
+    gt_points = gt_points[np.isfinite(gt_points).all(axis=1)]
+    if pred_points.shape[0] < 10 or gt_points.shape[0] < 10:
+        return _finish()
+
+    pcd = _as_o3d_point_cloud(pred_points)
+    pcd_gt = _as_o3d_point_cloud(gt_points)
+    reg_p2p = o3d.pipelines.registration.registration_icp(
+        pcd,
+        pcd_gt,
+        0.1,
+        np.eye(4),
+        o3d.pipelines.registration.TransformationEstimationPointToPoint(),
+    )
+    transform = reg_p2p.transformation
+    pcd.transform(transform)
+    pred_points_aligned = np.asarray(pcd.points)
+
+    gt_tree = KDTree(gt_points)
+    acc_distances, _ = gt_tree.query(pred_points_aligned, workers=-1)
+    pred_tree = KDTree(pred_points_aligned)
+    comp_distances, _ = pred_tree.query(gt_points, workers=-1)
+
+    acc = float(np.mean(acc_distances))
+    comp = float(np.mean(comp_distances))
+    cd = (acc + comp) / 2.0
+    if np.isfinite(acc) and np.isfinite(cd):
+        metrics['acc'] = acc * 1000.0
+        metrics['cd'] = cd * 1000.0
+
+    if return_misc:
+        misc['pred_edge_points'] = torch.from_numpy(pred_points_aligned).to(device=gt_depth.device, dtype=torch.float32)
+        misc['gt_edge_points'] = torch.from_numpy(gt_points).to(device=gt_depth.device, dtype=torch.float32)
+        misc['icp_transform'] = torch.from_numpy(transform.copy()).to(device=gt_depth.device, dtype=torch.float32)
+    return _finish()
+
+
+def _moge_lowres_affine(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    mask: torch.Tensor,
+    weight_depth: torch.Tensor,
+) -> Tuple[torch.Tensor, bool]:
+    """Fit MoGe-style weighted affine alignment on a 64x64 valid subset.
+
+    Args:
+        pred: Raw prediction map ``[H,W]`` in depth or log-depth space.
+        target: GT target map ``[H,W]`` in the same affine space.
+        mask: Boolean candidate-fit mask ``[H,W]``.
+        weight_depth: Positive GT depth ``[H,W]`` used for inverse-depth weights.
+
+    Returns:
+        aligned: Full-resolution floating prediction ``[H,W]``.
+        success: Boolean indicating whether finite affine parameters were found.
+    """
+    valid = (
+        mask
+        & torch.isfinite(pred)
+        & torch.isfinite(target)
+        & torch.isfinite(weight_depth)
+        & (weight_depth > 0)
+    )
+    if valid.sum().item() < ALIGN_MIN_VALID_PIXELS:
+        return pred.float(), False
+
+    pred_clean = torch.where(valid, pred.float(), torch.zeros_like(pred, dtype=torch.float32))
+    target_clean = torch.where(valid, target.float(), torch.zeros_like(target, dtype=torch.float32))
+    weight_depth_clean = torch.where(valid, weight_depth.float(), torch.ones_like(weight_depth, dtype=torch.float32))
+    try:
+        pred_lr, target_lr, weight_depth_lr, mask_lr = utils3d.pt.masked_nearest_resize(
+            pred_clean,
+            target_clean,
+            weight_depth_clean,
+            mask=valid,
+            size=(64, 64),
+        )
+        weight = mask_lr.flatten(-2, -1).float() / weight_depth_lr.flatten(-2, -1).clamp_min(1e-3)
+        if (weight > 0).sum().item() < ALIGN_MIN_VALID_PIXELS:
+            return pred.float(), False
+        scale, shift = align_depth_affine(
+            pred_lr.flatten(-2, -1),
+            target_lr.flatten(-2, -1),
+            weight,
+        )
+        scale = scale.squeeze()
+        shift = shift.squeeze()
+        ok = torch.isfinite(scale) & torch.isfinite(shift)
+        if not bool(ok.item() if ok.ndim == 0 else ok.all().item()):
+            return pred.float(), False
+        return pred.float() * scale + shift, True
+    except Exception:
+        return pred.float(), False
+
+
+def _moge_disparity_affine(
+    pred_disparity: torch.Tensor,
+    gt_disparity: torch.Tensor,
+    mask: torch.Tensor,
+) -> Tuple[torch.Tensor, bool]:
+    """Fit least-squares scale and shift in disparity space.
+
+    Args:
+        pred_disparity: Raw predicted disparity ``[H,W]``.
+        gt_disparity: Ground-truth reciprocal depth ``[H,W]``.
+        mask: Boolean fit mask ``[H,W]``.
+
+    Returns:
+        aligned: Full-resolution disparity ``[H,W]``.
+        success: Boolean indicating a finite affine fit.
+    """
+    valid = mask & torch.isfinite(pred_disparity) & torch.isfinite(gt_disparity) & (gt_disparity > 0)
+    if valid.sum().item() < ALIGN_MIN_VALID_PIXELS:
+        return pred_disparity.float(), False
+    try:
+        scale, shift = align_affine_lstsq(pred_disparity[valid].float(), gt_disparity[valid].float())
+        ok = torch.isfinite(scale) & torch.isfinite(shift)
+        if not bool(ok.item() if ok.ndim == 0 else ok.all().item()):
+            return pred_disparity.float(), False
+        return pred_disparity.float() * scale + shift, True
+    except Exception:
+        return pred_disparity.float(), False
+
+
+def _moge_points_affine(
+    pred_points: torch.Tensor,
+    gt_points: torch.Tensor,
+    mask: torch.Tensor,
+) -> Tuple[torch.Tensor, bool]:
+    """Fit one global scale and XYZ translation to a predicted point map.
+
+    Args:
+        pred_points: Predicted camera-space point map ``[H,W,3]``.
+        gt_points: Ground-truth point map ``[H,W,3]``.
+        mask: Boolean valid correspondence mask ``[H,W]``.
+
+    Returns:
+        aligned: Full-resolution point map ``[H,W,3]``.
+        success: Boolean indicating whether robust alignment succeeded.
+    """
+    valid = mask & torch.isfinite(pred_points).all(dim=-1) & torch.isfinite(gt_points).all(dim=-1)
+    if valid.sum().item() < ALIGN_MIN_VALID_PIXELS:
+        return pred_points.float(), False
+
+    pred_clean = torch.where(valid[..., None], pred_points.float(), torch.zeros_like(pred_points, dtype=torch.float32))
+    gt_clean = torch.where(valid[..., None], gt_points.float(), torch.zeros_like(gt_points, dtype=torch.float32))
+    try:
+        pred_lr, gt_lr, mask_lr = utils3d.pt.masked_nearest_resize(
+            pred_clean,
+            gt_clean,
+            mask=valid,
+            size=(64, 64),
+        )
+        weight = mask_lr.flatten(-2, -1).float() / gt_lr.norm(dim=-1).flatten(-2, -1).clamp_min(1e-6)
+        if (weight > 0).sum().item() < ALIGN_MIN_VALID_PIXELS:
+            return pred_points.float(), False
+        scale, shift = align_points_scale_xyz_shift(
+            pred_lr.flatten(-3, -2),
+            gt_lr.flatten(-3, -2),
+            weight,
+        )
+        scale = scale.squeeze()
+        shift = shift.squeeze()
+        ok = torch.isfinite(scale) & torch.isfinite(shift).all()
+        if not bool(ok.item() if ok.ndim == 0 else ok.all().item()):
+            return pred_points.float(), False
+        return pred_points.float() * scale + shift, True
+    except Exception:
+        return pred_points.float(), False
+
+
+def compute_metrics(
+    pred: Dict[str, torch.Tensor],
+    gt: Dict[str, torch.Tensor],
+    vis: bool = False,
+    compute_boundary: bool = True,
+) -> Tuple[Dict[str, Dict[str, Number]], Dict[str, torch.Tensor]]:
+    """Align one prediction and compute all applicable benchmark metrics.
+
+    Args:
+        pred: Prediction dictionary. It may contain raw ``depth_affine_invariant``
+            ``[H,W]`` plus ``depth_affine_space`` (``'depth'``/``'log'``), raw
+            ``disparity_affine_invariant`` ``[H,W]``, optional point map
+            ``points_affine_invariant`` ``[H,W,3]``, and predicted ``mask``
+            ``[H,W]``.
+        gt: Ground-truth sample containing depth/mask ``[H,W]``, point map
+            ``[H,W,3]``, normalized intrinsics ``[3,3]``, metric/boundary flags,
+            and optional segmentation annotations.
+        vis: Include aligned depth/points and boundary visualization tensors in the
+            auxiliary output.
+        compute_boundary: Evaluate boundary metrics when the dataset is marked
+            ``has_sharp_boundary``.
+
+    Returns:
+        metrics: Nested Python-number dictionary for depth, points, local points,
+            and optional boundary quality.
+        misc: Tensor dictionary containing aligned maps and optional boundary
+            visualization data when ``vis=True``.
+    """
+    metrics = {}
+    misc = {}
+
+    mask = gt['depth_mask']
+    gt_depth = gt['depth']
+    gt_points = gt['points']
+
+    valid_depth = mask & torch.isfinite(gt_depth) & (gt_depth > 0)
+    pred_depth_aligned = None
+    pred_points_aligned = None
+
+    if 'depth_affine_invariant' in pred:
+        raw_depth = pred['depth_affine_invariant'].float()
+        fit_mask = valid_depth & torch.isfinite(raw_depth)
+        affine_space = str(pred.get('depth_affine_space', 'depth')).lower()
+        if affine_space == 'log':
+            target_log = torch.log1p(gt_depth)
+            aligned_log, ok = _moge_lowres_affine(raw_depth, target_log, fit_mask, gt_depth)
+            pred_depth_aligned = torch.expm1(aligned_log if ok else raw_depth)
+        elif affine_space == 'depth':
+            aligned_depth, ok = _moge_lowres_affine(raw_depth, gt_depth, fit_mask, gt_depth)
+            pred_depth_aligned = aligned_depth if ok else raw_depth
+        else:
+            raise ValueError(f"Unsupported depth_affine_space={affine_space!r}")
+
+        metric_mask = fit_mask
+        if metric_mask.any():
+            metrics['depth_affine_invariant'] = {
+                'rel': rel_depth(pred_depth_aligned[metric_mask], gt_depth[metric_mask]),
+                'delta1': delta1_depth(pred_depth_aligned[metric_mask], gt_depth[metric_mask]),
+            }
+
+    elif 'disparity_affine_invariant' in pred:
+        raw_disparity = pred['disparity_affine_invariant'].float()
+        fit_mask = valid_depth & torch.isfinite(raw_disparity)
+        gt_disparity = torch.where(valid_depth, gt_depth.reciprocal(), torch.zeros_like(gt_depth))
+        aligned_disparity, ok = _moge_disparity_affine(raw_disparity, gt_disparity, fit_mask)
+        aligned_disparity = aligned_disparity if ok else raw_disparity
+        if fit_mask.any():
+            max_depth = gt_depth[fit_mask].max()
+            pred_depth_metric = aligned_disparity.clamp_min(max_depth.reciprocal()).reciprocal()
+        else:
+            pred_depth_metric = aligned_disparity.clamp_min(1e-6).reciprocal()
+        pred_depth_aligned = pred_depth_metric
+        metric_mask = fit_mask & torch.isfinite(pred_depth_metric)
+        if metric_mask.any():
+            metrics['depth_affine_invariant'] = {
+                'rel': rel_depth(pred_depth_metric[metric_mask], gt_depth[metric_mask]),
+                'delta1': delta1_depth(pred_depth_metric[metric_mask], gt_depth[metric_mask]),
+            }
+
+    pred_points_affine_invariant = pred.get('points_affine_invariant', None)
+    if pred_points_affine_invariant is None and pred_depth_aligned is not None:
+        point_intrinsics = gt['intrinsics'].to(
+            device=pred_depth_aligned.device,
+            dtype=pred_depth_aligned.dtype,
+        )
+        pred_points_affine_invariant = utils3d.pt.depth_map_to_point_map(
+            pred_depth_aligned,
+            intrinsics=point_intrinsics,
+        )
+
+    if pred_points_affine_invariant is not None:
+        point_mask = (
+            valid_depth
+            & torch.isfinite(pred_points_affine_invariant).all(dim=-1)
+            & torch.isfinite(gt_points).all(dim=-1)
+        )
+        if point_mask.any():
+            aligned_points, ok = _moge_points_affine(pred_points_affine_invariant, gt_points, point_mask)
+            pred_points_aligned = aligned_points if ok else pred_points_affine_invariant
+            metrics['points_affine_invariant'] = {
+                'rel': rel_point(pred_points_aligned[point_mask], gt_points[point_mask]),
+                'delta1': delta1_point(pred_points_aligned[point_mask], gt_points[point_mask]),
+            }
+
+    # Local points
+    if 'segmentation_mask' in gt and 'points' in gt and pred_points_affine_invariant is not None:
+        pred_points = pred_points_affine_invariant
+        gt_points = gt['points']
+        segmentation_mask = gt['segmentation_mask']
+        segmentation_labels = gt['segmentation_labels']
+        local_points_metrics = []
+        for _, seg_id in segmentation_labels.items():
+            valid_mask = (
+                (segmentation_mask == seg_id)
+                & valid_depth
+                & torch.isfinite(pred_points).all(dim=-1)
+                & torch.isfinite(gt_points).all(dim=-1)
+            )
+            if valid_mask.sum().item() < 10:
+                continue
+
+            try:
+                pred_lr, gt_lr, mask_lr = utils3d.pt.masked_nearest_resize(
+                    torch.where(valid_mask[..., None], pred_points.float(), torch.zeros_like(pred_points, dtype=torch.float32)),
+                    torch.where(valid_mask[..., None], gt_points.float(), torch.zeros_like(gt_points, dtype=torch.float32)),
+                    mask=valid_mask,
+                    size=(64, 64),
+                )
+                pred_points_masked = pred_lr[mask_lr]
+                gt_points_masked = gt_lr[mask_lr]
+                if pred_points_masked.shape[0] < 10:
+                    continue
+                diameter = (gt_points_masked.max(dim=0).values - gt_points_masked.min(dim=0).values).max()
+                scale, shift = align_points_scale_xyz_shift(
+                    pred_points_masked.unsqueeze(0),
+                    gt_points_masked.unsqueeze(0),
+                    diameter.clamp_min(1e-6).reciprocal().expand(1, gt_points_masked.shape[0]),
+                )
+                pred_points_masked = pred_points[valid_mask] * scale.squeeze() + shift.squeeze()
+                gt_points_masked = gt_points[valid_mask]
+            except Exception:
+                pred_points_masked = pred_points[valid_mask]
+                gt_points_masked = gt_points[valid_mask]
+                diameter = (gt_points_masked.max(dim=0).values - gt_points_masked.min(dim=0).values).max()
+
+            local_points_metrics.append({
+                'rel': rel_point_local(pred_points_masked, gt_points_masked, diameter),
+                'delta1': delta1_point_local(pred_points_masked, gt_points_masked, diameter),
+            })
+
+        metrics['local_points'] = key_average(local_points_metrics)
+
+    # Boundary Acc/CD with the MDA/Canny edge.
+    boundary_depth = pred_depth_aligned
+    if compute_boundary and boundary_depth is not None and gt['has_sharp_boundary']:
+        if vis:
+            boundary_metrics, boundary_misc = boundary_edge_metrics(
+                boundary_depth,
+                gt_depth,
+                mask,
+                gt['intrinsics'],
+                return_misc=True,
+                edge_mode='mda',
+            )
+        else:
+            boundary_metrics = boundary_edge_metrics(
+                boundary_depth,
+                gt_depth,
+                mask,
+                gt['intrinsics'],
+                edge_mode='mda',
+            )
+            boundary_misc = {}
+        metrics['boundary'] = boundary_metrics
+        if vis:
+            if 'edge_mask' in boundary_misc:
+                misc['boundary_edge_mask'] = boundary_misc['edge_mask']
+            if 'pred_edge_points' in boundary_misc:
+                misc['boundary_pred_edge_points'] = boundary_misc['pred_edge_points']
+            if 'gt_edge_points' in boundary_misc:
+                misc['boundary_gt_edge_points'] = boundary_misc['gt_edge_points']
+            if 'icp_transform' in boundary_misc:
+                misc['boundary_icp_transform'] = boundary_misc['icp_transform']
+
+    if vis:
+        if pred_points_aligned is not None:
+            misc['pred_points'] = pred_points_aligned
+        elif pred_depth_aligned is not None:
+            misc['pred_points'] = utils3d.pt.depth_map_to_point_map(pred_depth_aligned, intrinsics=gt['intrinsics'])
+        if pred_depth_aligned is not None:
+            misc['pred_depth'] = pred_depth_aligned
+
+    return metrics, misc

pxdepth/inference/__init__.py

@@ -0,0 +1,20 @@
+"""Public preprocessing and raw-forward helpers for PXDepth inference.
+
+The package centralizes fixed-size and equal-area image resizing, patch-grid
+alignment, output restoration, and timed model execution. These helpers return
+raw normalized predictions and avoid embedding benchmark-specific alignment in
+the model forward path.
+"""
+
+from .runner import predict_raw
+from .resize import area_size, area_size_from_area, parse_size, patch_size, resize_image, resize_map
+
+__all__ = [
+    "predict_raw",
+    "area_size",
+    "area_size_from_area",
+    "parse_size",
+    "patch_size",
+    "resize_image",
+    "resize_map",
+]

pxdepth/inference/resize.py

@@ -0,0 +1,158 @@
+"""Image and prediction resizing shared by evaluation and inference.
+
+The functions parse user-facing sizes, derive patch-compatible equal-area
+shapes, resize RGB tensors for a model, and restore depth or mask maps to the
+source resolution. Their return values retain the original image dimensions so
+camera-normalized geometry remains consistent after restoration.
+"""
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+
+def parse_size(value: Optional[str]) -> Optional[Tuple[int, int]]:
+    """Parse a CLI image-size string in ``WIDTHxHEIGHT`` notation.
+
+    Args:
+        value: Size string, or ``None``/empty text when no target is requested.
+
+    Returns:
+        Positive integer tuple ``(width, height)``, or ``None``.
+    """
+    if value is None or not str(value).strip():
+        return None
+    parts = str(value).lower().replace(",", "x").split("x")
+    if len(parts) != 2:
+        raise ValueError("Image size must be WIDTHxHEIGHT, for example 1022x770.")
+    width, height = map(int, parts)
+    if width <= 0 or height <= 0:
+        raise ValueError("Image width and height must be positive.")
+    return width, height
+
+
+def area_size(
+    height: int,
+    width: int,
+    target_width: int,
+    target_height: int,
+    patch_size: int,
+) -> Tuple[int, int]:
+    """Preserve aspect ratio while matching a reference width-height area.
+
+    Args:
+        height: Original image height ``H``.
+        width: Original image width ``W``.
+        target_width: Width defining the desired reference area.
+        target_height: Height defining the desired reference area.
+        patch_size: Required divisibility of both output dimensions.
+
+    Returns:
+        Integer ``(new_height, new_width)`` with approximately
+        ``target_width*target_height`` pixels and the original aspect ratio.
+    """
+    area = int(target_width) * int(target_height)
+    return area_size_from_area(height, width, area, patch_size)
+
+
+def area_size_from_area(
+    height: int,
+    width: int,
+    target_area: int,
+    patch_size: int,
+) -> Tuple[int, int]:
+    """Preserve aspect ratio while matching an explicit target pixel area.
+
+    Args:
+        height: Original image height ``H``.
+        width: Original image width ``W``.
+        target_area: Desired number of input pixels before patch rounding.
+        patch_size: Required divisibility of both output dimensions.
+
+    Returns:
+        Integer ``(new_height, new_width)`` rounded to patch multiples.
+    """
+    area = int(target_area)
+    if height <= 0 or width <= 0 or area <= 0:
+        raise ValueError("Image dimensions and target area must be positive.")
+    aspect = width / height
+    new_width = int(round((area * aspect) ** 0.5))
+    new_height = int(round(new_width / aspect))
+    new_width = max(patch_size, int(round(new_width / patch_size)) * patch_size)
+    new_height = max(patch_size, int(round(new_height / patch_size)) * patch_size)
+    return new_height, new_width
+
+
+def patch_size(height: int, width: int, patch: int) -> Tuple[int, int]:
+    """Round spatial dimensions down to valid patch multiples.
+
+    Args:
+        height: Original image height.
+        width: Original image width.
+        patch: Positive encoder patch side length.
+
+    Returns:
+        Integer ``(new_height, new_width)``. Already divisible dimensions are
+        unchanged; smaller results are clamped to one patch.
+    """
+    new_height = height if height % patch == 0 else max(patch, height // patch * patch)
+    new_width = width if width % patch == 0 else max(patch, width // patch * patch)
+    return new_height, new_width
+
+
+def resize_image(
+    image: torch.Tensor,
+    target: Optional[Tuple[int, int]],
+    resize_by_area: bool,
+    patch: int,
+) -> Tuple[torch.Tensor, Tuple[int, int]]:
+    """Resize a CHW/BCHW image to the model input resolution.
+
+    Args:
+        image: RGB tensor ``[3,H,W]`` or ``[B,3,H,W]``.
+        target: Optional ``(width,height)`` reference size.
+        resize_by_area: Preserve aspect ratio and use only ``target`` area when
+            ``True``; otherwise force the exact target dimensions.
+        patch: Required encoder patch divisibility.
+
+    Returns:
+        resized: Bilinearly resized tensor preserving the input batch layout.
+        original_size: Original integer tuple ``(H,W)``.
+    """
+    original = tuple(image.shape[-2:])
+    if target is None:
+        height, width = patch_size(*original, patch)
+    elif resize_by_area:
+        height, width = area_size(*original, target[0], target[1], patch)
+    else:
+        width, height = target
+        if height % patch or width % patch:
+            raise ValueError(f"Fixed input size {width}x{height} must be divisible by patch size {patch}.")
+    if (height, width) == original:
+        return image, original
+    batched = image.ndim == 4
+    source = image if batched else image.unsqueeze(0)
+    resized = F.interpolate(source, (height, width), mode="bilinear", align_corners=False)
+    return (resized if batched else resized[0]), original
+
+
+def resize_map(value: torch.Tensor, size: Tuple[int, int], is_mask: bool = False) -> torch.Tensor:
+    """Nearest-resize a depth/probability map while preserving batch layout.
+
+    Args:
+        value: Map tensor ``[H,W]`` or batch ``[B,H,W]``.
+        size: Target ``(height,width)``.
+        is_mask: Threshold resized values at ``0.5`` and return boolean output.
+
+    Returns:
+        Tensor ``[H_t,W_t]`` or ``[B,H_t,W_t]``. Non-mask output is floating;
+        mask output is boolean.
+    """
+    if tuple(value.shape[-2:]) == tuple(size):
+        return value
+    batched = value.ndim == 3
+    source = value.float().unsqueeze(1) if batched else value.float()[None, None]
+    output = F.interpolate(source, size=size, mode="nearest")
+    output = output[:, 0] if batched else output[0, 0]
+    return output > 0.5 if is_mask else output

pxdepth/inference/runner.py

@@ -0,0 +1,84 @@
+"""Timed raw PXDepth forward path used by evaluation tools.
+
+This module applies the selected fixed-size or equal-area preprocessing,
+measures only model-forward latency, and restores predictions to input
+resolution. It intentionally returns only genuine network outputs rather than
+performing metric alignment or estimating camera intrinsics.
+"""
+
+import time
+from typing import Dict, Optional, Tuple
+
+import torch
+
+from ..model import PXDepth
+from .resize import resize_image, resize_map
+
+
+def synchronize(device: torch.device) -> None:
+    """Synchronize pending CUDA work before or after timing model forward.
+
+    Args:
+        device: Model device. CPU and other devices require no action.
+
+    Returns:
+        ``None``.
+    """
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+@torch.inference_mode()
+def predict_raw(
+    model: PXDepth,
+    image: torch.Tensor,
+    input_size: Optional[Tuple[int, int]] = (1022, 770),
+    resize_by_area: bool = True,
+    use_fp16: bool = False,
+    use_fp32: bool = False,
+) -> Dict[str, torch.Tensor]:
+    """Run raw model forward and return outputs at the original image size.
+
+    Only ``model.forward`` is included in ``inference_time``. Resizing,
+    synchronization overhead, and output packaging are excluded. No GT depth
+    alignment or camera-intrinsics prediction is performed.
+
+    Args:
+        model: Evaluation-mode :class:`PXDepth` model.
+        image: RGB tensor ``[3,H,W]`` or ``[B,3,H,W]`` in ``[0,1]``.
+        input_size: Exact/reference tuple ``(width,height)``. Defaults to
+            ``(1022,770)``.
+        resize_by_area: Preserve aspect ratio at ``input_size`` area. Enabled
+            by default.
+        use_fp16: Use FP16 for attention-heavy model regions.
+        use_fp32: Force full-precision model execution.
+
+    Returns:
+        Dictionary with raw normalized log-depth ``depth_affine_invariant``
+        ``[B,H,W]``, ``depth_affine_space='log'``, boolean ``mask`` ``[B,H,W]``,
+        and scalar forward time. The leading batch dimension is removed for
+        unbatched input.
+    """
+    image, original_size = resize_image(image, input_size, resize_by_area, model.patch_size)
+    batched = image.ndim == 4
+    model_input = image if batched else image.unsqueeze(0)
+    model_input = model_input.to(device=model.device, dtype=torch.float32)
+
+    synchronize(model.device)
+    start = time.perf_counter()
+    output = model.forward(model_input, use_fp16=use_fp16, use_fp32=use_fp32)
+    synchronize(model.device)
+    elapsed = time.perf_counter() - start
+
+    depth = resize_map(output["depth"], original_size)
+    mask = resize_map(output["mask"], original_size, is_mask=True)
+
+    pred = {
+        "depth_affine_invariant": depth,
+        "depth_affine_space": "log",
+        "mask": mask,
+        "inference_time": elapsed,
+    }
+    if not batched:
+        pred = {key: value[0] if isinstance(value, torch.Tensor) and value.ndim > 0 else value for key, value in pred.items()}
+    return pred

pxdepth/model/CM_PiT.py

@@ -0,0 +1,260 @@
+"""Context-Modulated Pixel Transformer (CM-PiT) building blocks.
+
+CM-PiT compresses local dense pixel features into attention tokens, processes
+them with gated self-attention and SwiGLU, and expands them back without losing
+the original pixel lattice. Global encoder tokens generate adaptive shift,
+scale, and residual gates that condition both transformer sublayers.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from .Gated_Attention import GatedAttention
+from .RoPE import RotaryPositionEmbedding2D
+from .precision import full_precision
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """Apply affine context modulation without changing the tensor layout.
+
+    Args:
+        x: Normalized pixel tokens with shape ``[B, N, P, C]``.
+        shift: Context-predicted additive offsets with shape ``[B, N, P, C]``.
+        scale: Context-predicted residual scales with shape ``[B, N, P, C]``.
+
+    Returns:
+        Modulated tokens ``x * (1 + scale) + shift`` with shape
+        ``[B, N, P, C]``.
+    """
+    return x * (1.0 + scale) + shift
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU feed-forward layer operating independently on every pixel token.
+
+    The first projection creates value and gate branches, SiLU activates the
+    gate, and the second projection returns to the pixel-channel dimension.
+    Spatial and patch axes are preserved throughout the module.
+    """
+
+    def __init__(self, dim: int, hidden_dim: int) -> None:
+        """Construct the gated feed-forward projections.
+
+        Args:
+            dim: Input and output channel count ``C``.
+            hidden_dim: Width of each hidden value/gate branch.
+
+        Returns:
+            ``None``. Learnable linear layers are registered on the module.
+        """
+        super().__init__()
+        self.fc1 = nn.Linear(dim, hidden_dim * 2)
+        self.fc2 = nn.Linear(hidden_dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Transform pixel tokens with a SiLU-gated hidden representation.
+
+        Args:
+            x: Floating tensor with arbitrary leading dimensions and final
+                channel dimension ``C=dim``. CM-PiT supplies ``[B,N,P,C]``.
+
+        Returns:
+            Tensor with the same shape and dtype as ``x``.
+        """
+        value, gate = self.fc1(x).chunk(2, dim=-1)
+        return self.fc2(value * F.silu(gate))
+
+
+class ContextAdaNorm(nn.Module):
+    """Predict Context-Guided Adaptive Normalization parameters.
+
+    Each global context token produces shift, scale, and residual-gate values
+    for both the attention and MLP sublayers over every pixel represented by
+    that encoder token. The six parameter groups are unpacked by
+    :class:`CMPiTBlock`.
+    """
+
+    def __init__(self, dim_ctx: int, patch_size: int, dim_pix: int) -> None:
+        """Create the context-to-modulation projection.
+
+        Args:
+            dim_ctx: Channel count of each Global Context Encoder token.
+            patch_size: Encoder patch side length ``P_ctx`` in image pixels.
+            dim_pix: Pixel-feature channel count ``C_pix``.
+
+        Returns:
+            ``None``. The projection outputs ``6 * P_ctx^2 * C_pix`` values per
+            context token.
+        """
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim_ctx, 6 * patch_size * patch_size * dim_pix),
+        )
+
+    def forward(self, ctx: torch.Tensor) -> torch.Tensor:
+        """Project context tokens into six dense pixel-wise parameter fields.
+
+        Args:
+            ctx: Context token tensor ``[B, N_ctx, C_ctx]``.
+
+        Returns:
+            Modulation tensor ``[B, N_ctx, 6 * P_ctx^2 * C_pix]``.
+        """
+        return self.proj(ctx)
+
+
+class CMPiTBlock(nn.Module):
+    """Context-Modulated Pixel Transformer block.
+
+    The block groups a dense pixel feature map into local patches, linearly
+    compresses every patch to an attention token, applies gated global
+    self-attention, expands the token back to pixel features, and follows it
+    with a per-pixel SwiGLU MLP. Both residual branches use Context-Guided
+    Adaptive Normalization generated from DINO context tokens.
+    """
+
+    def __init__(
+        self,
+        dim_ctx: int,
+        ctx_patch_size: int,
+        dim_pix: int,
+        patch_size: int,
+        attn_dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qk_norm: bool = True,
+        rope: Optional[RotaryPositionEmbedding2D] = None,
+        eps: float = 1e-6,
+    ) -> None:
+        """Configure one CM-PiT block.
+
+        Args:
+            dim_ctx: Context-token channel count ``C_ctx``.
+            ctx_patch_size: Image patch size ``P_ctx`` represented by one
+                context token.
+            dim_pix: Dense pixel-feature channel count ``C_pix``.
+            patch_size: Side length ``P`` grouped into one attention token.
+                It must divide ``ctx_patch_size``.
+            attn_dim: Compressed attention-token channel count ``D``.
+            num_heads: Number of attention heads. ``D`` must be divisible by it.
+            mlp_ratio: Expansion ratio controlling the SwiGLU hidden width.
+            qk_norm: Whether to apply FP32 RMSNorm to each query/key head.
+            rope: Optional 2D rotary position embedding shared by decoder blocks.
+            eps: Numerical epsilon used by RMSNorm layers.
+
+        Returns:
+            ``None``. Attention, modulation, MLP, and projection layers are
+            registered on the block.
+        """
+        super().__init__()
+        if ctx_patch_size % patch_size != 0:
+            raise ValueError(
+                f"ctx_patch_size ({ctx_patch_size}) must be divisible by patch_size ({patch_size})"
+            )
+
+        self.dim_ctx = dim_ctx
+        self.dim_pix = dim_pix
+        self.ctx_patch_size = ctx_patch_size
+        self.patch_size = patch_size
+        patch_dim = patch_size * patch_size * dim_pix
+
+        self.norm1 = nn.RMSNorm(dim_pix, eps=eps)
+        self.linear_compress = nn.Linear(patch_dim, attn_dim)
+        self.attn = GatedAttention(attn_dim, num_heads, qk_norm=qk_norm, rope=rope, eps=eps)
+        self.linear_expand = nn.Linear(attn_dim, patch_dim)
+        self.norm2 = nn.RMSNorm(dim_pix, eps=eps)
+        hidden_dim = max(1, int(round(dim_pix * mlp_ratio * 2.0 / 3.0)))
+        self.mlp = SwiGLU(dim_pix, hidden_dim)
+        self.ada_norm = ContextAdaNorm(dim_ctx, ctx_patch_size, dim_pix)
+
+    @staticmethod
+    def _norm(norm: nn.Module, x: torch.Tensor) -> torch.Tensor:
+        """Evaluate a normalization layer in FP32 and restore input dtype.
+
+        Args:
+            norm: Normalization module acting on the final channel dimension.
+            x: Pixel tokens ``[B, N, P^2, C_pix]`` in the active model dtype.
+
+        Returns:
+            Normalized tensor with the same shape and dtype as ``x``.
+        """
+        dtype = x.dtype
+        with full_precision(x.device):
+            out = norm(x.float())
+        return out.to(dtype)
+
+    def _modulation(self, ctx: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        """Align context modulation fields with the block's pixel patches.
+
+        Args:
+            ctx: Global context tokens ``[B, H_ctx*W_ctx, C_ctx]``.
+            height: Dense pixel-map height ``H``.
+            width: Dense pixel-map width ``W``.
+
+        Returns:
+            Six modulation groups with shape
+            ``[B, (H/P)*(W/P), 6, P^2, C_pix]``. Rearrangement is exact and
+            contains no interpolation.
+        """
+        batch = ctx.shape[0]
+        p_ctx, p = self.ctx_patch_size, self.patch_size
+        ctx_h, ctx_w = height // p_ctx, width // p_ctx
+        if ctx.shape[1] != ctx_h * ctx_w:
+            raise ValueError(
+                f"Context token count ({ctx.shape[1]}) does not match grid ({ctx_h}x{ctx_w})"
+            )
+
+        mod = self.ada_norm(ctx).view(batch, ctx_h, ctx_w, 6, p_ctx, p_ctx, self.dim_pix)
+        if p_ctx == p:
+            return rearrange(mod, "b h w m ph pw c -> b (h w) m (ph pw) c")
+
+        ratio = p_ctx // p
+        return rearrange(
+            mod,
+            "b h w m (rh ph) (rw pw) c -> b (h rh w rw) m (ph pw) c",
+            rh=ratio,
+            rw=ratio,
+            ph=p,
+            pw=p,
+        )
+
+    def forward(self, x: torch.Tensor, ctx: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        """Apply context-modulated attention and MLP residual updates.
+
+        Args:
+            x: Dense pixel features ``[B, C_pix, H, W]``.
+            ctx: Global context tokens ``[B, (H/P_ctx)*(W/P_ctx), C_ctx]``.
+            pos: Integer 2D token positions ``[B, (H/P)*(W/P), 2]`` used by
+                rotary position embedding in self-attention.
+
+        Returns:
+            Updated dense pixel features ``[B, C_pix, H, W]``.
+        """
+        batch, _, height, width = x.shape
+        p = self.patch_size
+        pix = rearrange(x, "b c (h ph) (w pw) -> b (h w) (ph pw) c", ph=p, pw=p)
+        shift_attn, scale_attn, gate_attn, shift_mlp, scale_mlp, gate_mlp = self._modulation(
+            ctx, height, width
+        ).unbind(dim=2)
+
+        out = modulate(self._norm(self.norm1, pix), shift_attn, scale_attn)
+        out = self.linear_compress(out.flatten(2))
+        out = self.attn(out, pos=pos)
+        out = self.linear_expand(out).view(batch, -1, p * p, self.dim_pix)
+        pix = pix + gate_attn * out
+
+        out = modulate(self._norm(self.norm2, pix), shift_mlp, scale_mlp)
+        pix = pix + gate_mlp * self.mlp(out)
+        return rearrange(
+            pix,
+            "b (h w) (ph pw) c -> b c (h ph) (w pw)",
+            h=height // p,
+            w=width // p,
+            ph=p,
+            pw=p,
+        )

pxdepth/model/Gated_Attention.py

@@ -0,0 +1,105 @@
+"""Gated multi-head self-attention used inside CM-PiT blocks.
+
+The implementation performs optional FP32 query/key normalization, applies
+two-dimensional rotary position embeddings, and delegates attention to PyTorch
+SDPA. A learned token-channel sigmoid gate modulates the attended features
+before the output projection.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .RoPE import RotaryPositionEmbedding2D
+from .precision import full_precision
+
+
+class GatedAttention(nn.Module):
+    """Multi-head self-attention followed by a learned token-channel gate.
+
+    The Q/K normalization and 2D RoPE calculations follow the numerical path
+    used for the released model. Q/K normalization and RoPE are evaluated in
+    FP32, while SDPA follows the active decoder autocast dtype.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        qk_norm: bool = True,
+        rope: Optional[RotaryPositionEmbedding2D] = None,
+        eps: float = 1e-6,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        """Construct gated multi-head self-attention.
+
+        Args:
+            dim: Token channel count ``D``.
+            num_heads: Number of attention heads. ``D`` must be divisible by it.
+            qk_norm: Enable per-head RMSNorm for queries and keys.
+            rope: Optional 2D rotary position embedding module.
+            eps: Epsilon used by query/key RMSNorm.
+            attn_drop: Attention-probability dropout used during training.
+            proj_drop: Dropout applied after the output projection.
+
+        Returns:
+            ``None``. QKV, gate, output, and optional normalization layers are
+            registered on the module.
+        """
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})")
+
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.k_norm = nn.RMSNorm(self.head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.rope = rope
+        self.attn_drop = float(attn_drop)
+        self.gate = nn.Linear(dim, dim)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor, pos: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Apply self-attention and token-channel gating.
+
+        Args:
+            x: Compressed patch tokens ``[B, N, D]``.
+            pos: Optional integer grid coordinates ``[B, N, 2]``. They are
+                required when a rotary position embedding is configured.
+
+        Returns:
+            Gated and projected attention output ``[B, N, D]``.
+        """
+        batch, length, dim = x.shape
+        gate = torch.sigmoid(self.gate(x))
+        qkv = self.qkv(x).reshape(batch, length, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)
+
+        if not isinstance(self.q_norm, nn.Identity):
+            dtype = q.dtype
+            with full_precision(q.device):
+                q = self.q_norm(q.float())
+                k = self.k_norm(k.float())
+            q, k = q.to(dtype), k.to(dtype)
+
+        if self.rope is not None:
+            dtype = q.dtype
+            with full_precision(q.device):
+                q = self.rope(q.float(), pos)
+                k = self.rope(k.float(), pos)
+            q, k = q.to(dtype), k.to(dtype)
+
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            dropout_p=self.attn_drop if self.training else 0.0,
+        )
+        out = out.transpose(1, 2).reshape(batch, length, dim)
+        out = out * gate.to(out.dtype)
+        return self.proj_drop(self.proj(out))

pxdepth/model/Global_Context_Encoder.py

@@ -0,0 +1,185 @@
+"""Global Context Encoder used to condition pixel-space depth prediction.
+
+A DINOv2 vision transformer extracts selected intermediate patch-token maps.
+Each map is normalized, reshaped to its image grid, projected to a common
+channel width, and summed into the context feature consumed by CM-PiT adaptive
+normalization layers.
+"""
+
+from typing import List, Sequence, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..registry import ENCODERS
+from .dinov2.hub import backbones
+from .utils import wrap_dinov2_attention_with_sdpa, wrap_module_with_gradient_checkpointing
+
+
+@ENCODERS.register()
+class GlobalContextEncoder(nn.Module):
+    """Global Context Encoder based on intermediate DINOv2 features.
+
+    The encoder extracts several normalized patch-token maps from a ViT,
+    projects each map to a shared channel width with a 1x1 convolution, and
+    sums the projected maps. The resulting grid provides global semantic
+    context for Context-Guided Adaptive Normalization in the pixel predictor.
+    """
+
+    def __init__(
+        self,
+        backbone: str = "dinov2_vitl14",
+        intermediate_layers: Union[int, Sequence[int]] = (5, 11, 17, 23),
+        dim_out: int = 1024,
+    ) -> None:
+        """Construct the DINOv2 backbone and intermediate projections.
+
+        Args:
+            backbone: Name of a constructor exposed by ``dinov2.hub.backbones``.
+            intermediate_layers: Explicit zero-based block indices or an
+                integer requesting the last ``n`` intermediate layers.
+            dim_out: Channel count ``C_ctx`` of every projected context map.
+
+        Returns:
+            ``None``. The backbone, output projections, and ImageNet
+            normalization buffers are registered on the module.
+        """
+        super().__init__()
+        if not hasattr(backbones, backbone):
+            raise ValueError(f"Unsupported DINOv2 backbone: {backbone}")
+
+        self.backbone_name = backbone
+        self.intermediate_layers = list(intermediate_layers) if not isinstance(intermediate_layers, int) else intermediate_layers
+        self.backbone = getattr(backbones, backbone)(pretrained=False)
+        if hasattr(self.backbone, "mask_token"):
+            self.backbone.mask_token.requires_grad_(False)
+
+        patch_size = getattr(self.backbone, "patch_size", 14)
+        if isinstance(patch_size, (tuple, list)):
+            patch_size = patch_size[0]
+        self.patch_size = int(patch_size)
+        self.dim_features = int(getattr(self.backbone, "embed_dim"))
+        self.dim_out = int(dim_out)
+        count = self.intermediate_layers if isinstance(self.intermediate_layers, int) else len(self.intermediate_layers)
+        self.output_projections = nn.ModuleList(
+            nn.Conv2d(self.dim_features, dim_out, kernel_size=1) for _ in range(count)
+        )
+
+        self.register_buffer("image_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+        self.register_buffer("image_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+        self._onnx_compatible_mode = False
+
+    @property
+    def onnx_compatible_mode(self) -> bool:
+        """Report whether ONNX-compatible resize behavior is enabled.
+
+        Returns:
+            Boolean flag controlling antialiasing and the vendored backbone's
+            ONNX compatibility path.
+        """
+        return self._onnx_compatible_mode
+
+    @onnx_compatible_mode.setter
+    def onnx_compatible_mode(self, enabled: bool) -> None:
+        """Enable or disable ONNX-compatible encoder operators.
+
+        Args:
+            enabled: Boolean state propagated to the DINOv2 backbone.
+
+        Returns:
+            ``None``. Runtime flags are updated in place.
+        """
+        self._onnx_compatible_mode = bool(enabled)
+        self.backbone.onnx_compatible_mode = bool(enabled)
+
+    def init_weights(self) -> None:
+        """Load official pretrained weights for the configured DINOv2 backbone.
+
+        Returns:
+            ``None``. Backbone parameters are replaced in place while the
+            PXDepth-specific 1x1 projections keep their initialization.
+        """
+        state = getattr(backbones, self.backbone_name)(pretrained=True).state_dict()
+        self.backbone.load_state_dict(state, strict=True)
+
+    def enable_gradient_checkpointing(self) -> None:
+        """Wrap every DINO transformer block with activation checkpointing.
+
+        Parameter names and numerical block behavior remain unchanged; only
+        activation storage during training is affected.
+
+        Returns:
+            ``None``. Each backbone block is modified in place.
+        """
+        for block in self.backbone.blocks:
+            wrap_module_with_gradient_checkpointing(block)
+
+    def enable_pytorch_native_sdpa(self) -> None:
+        """Replace DINO attention forward methods with SDPA-compatible paths.
+
+        Returns:
+            ``None``. Attention modules are wrapped in place and use
+            Flash-Attention when the installed runtime supports it.
+        """
+        for block in self.backbone.blocks:
+            wrap_dinov2_attention_with_sdpa(block.attn)
+
+    def forward(
+        self,
+        image: torch.Tensor,
+        token_rows: int,
+        token_cols: int,
+        return_feature_maps: bool = False,
+        return_class_token: bool = False,
+    ):
+        """Encode RGB images into a summed global context feature map.
+
+        Args:
+            image: RGB tensor ``[B, 3, H_in, W_in]`` with values in ``[0, 1]``.
+            token_rows: Requested context-grid height ``H_ctx``.
+            token_cols: Requested context-grid width ``W_ctx``.
+            return_feature_maps: Also return the list of individually projected
+                feature maps when ``True``.
+            return_class_token: Also return the final selected DINO class token
+                ``[B, C_vit]`` when ``True``.
+
+        Returns:
+            By default, a context map ``[B, C_ctx, H_ctx, W_ctx]``. Optional
+            outputs are appended as a tuple in the order ``feature_maps`` then
+            ``class_token``. Each feature map has shape
+            ``[B, C_ctx, H_ctx, W_ctx]``.
+        """
+        target_size = (token_rows * self.patch_size, token_cols * self.patch_size)
+        if image.shape[-2:] != target_size:
+            image = F.interpolate(
+                image,
+                size=target_size,
+                mode="bilinear",
+                align_corners=False,
+                antialias=not self.onnx_compatible_mode,
+            )
+        image = (image - self.image_mean) / self.image_std
+        features = self.backbone.get_intermediate_layers(
+            image,
+            n=self.intermediate_layers,
+            return_class_token=True,
+            norm=True,
+        )
+        maps = []
+        context = None
+        for projection, (tokens, _) in zip(self.output_projections, features):
+            feature = tokens.permute(0, 2, 1).unflatten(2, (token_rows, token_cols)).contiguous()
+            projected = projection(feature)
+            context = projected if context is None else context + projected
+            if return_feature_maps:
+                maps.append(projected)
+        if context is None:
+            raise RuntimeError("Global Context Encoder did not receive any intermediate features.")
+
+        outputs: List[object] = [context]
+        if return_feature_maps:
+            outputs.append(maps)
+        if return_class_token:
+            outputs.append(features[-1][1])
+        return outputs[0] if len(outputs) == 1 else tuple(outputs)

pxdepth/model/PXDepth.py

@@ -0,0 +1,276 @@
+"""Core PXDepth architecture and stable public model API.
+
+The module connects the Global Context Encoder to the Pixel-Space Depth
+Predictor and defines raw forward computation. Checkpoint translation and
+metric-scale inference live in focused helper modules, while their familiar
+``from_pretrained`` and ``infer`` entry points remain methods on this class.
+"""
+
+from pathlib import Path
+from typing import Any, Dict, IO, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .Global_Context_Encoder import GlobalContextEncoder
+from .Pixel_Space_Depth_Predictor import PixelSpaceDepthPredictor
+from .checkpoint import load_pretrained
+from .inference import infer as infer_model
+from .precision import full_precision, inference_dtype, reduced_precision
+from ..registry import ENCODERS, MODELS, PREDICTORS
+
+
+@MODELS.register()
+class PXDepth(nn.Module):
+    """Complete PXDepth monocular depth model.
+
+    A Global Context Encoder extracts semantic patch features and a Pixel-Space
+    Depth Predictor estimates full-resolution normalized log-depth together
+    with a finite-depth probability. ``forward`` exposes raw network outputs,
+    while ``infer`` aligns them to a GT or MoGe-2 reference for metric-scale
+    visualization and point-cloud reconstruction.
+    """
+
+    def __init__(
+        self,
+        encoder: Union[nn.Module, Dict[str, Any]],
+        predictor: Union[nn.Module, Dict[str, Any]],
+        remap_output: str = "linear",
+        mask_threshold: float = 0.5,
+    ) -> None:
+        """Construct the encoder and CM-PiT pixel predictor.
+
+        Args:
+            encoder: Encoder module or registry config.
+            predictor: Pixel predictor module or registry config. Its context
+                patch size and channel width default to the encoder contract.
+            remap_output: Output remapping applied to normalized log-depth.
+                The released model uses ``'linear'``.
+            mask_threshold: Probability threshold used by :meth:`infer`.
+
+        Returns:
+            ``None``. Model modules and ImageNet normalization buffers are
+            registered on the instance.
+        """
+        super().__init__()
+        if remap_output not in {"linear", "elu"}:
+            raise ValueError(f"Unsupported remap_output: {remap_output}")
+
+        self.remap_output = remap_output
+        self.mask_threshold = float(mask_threshold)
+        if isinstance(encoder, nn.Module):
+            self.encoder = encoder
+        else:
+            encoder_config = dict(encoder)
+            encoder_config.setdefault("type", "GlobalContextEncoder")
+            self.encoder = ENCODERS.build(encoder_config)
+        if not hasattr(self.encoder, "patch_size"):
+            raise TypeError("The encoder must expose an integer patch_size attribute.")
+        self.patch_size = self.encoder.patch_size
+        self.p_enc = self.patch_size
+        dim_ctx = getattr(self.encoder, "dim_out", None)
+        if isinstance(predictor, nn.Module):
+            self.predictor = predictor
+        else:
+            predictor_config = dict(predictor)
+            predictor_config.setdefault("type", "PixelSpaceDepthPredictor")
+            predictor_config.setdefault("in_channels", 3)
+            predictor_config.setdefault("ctx_patch_size", self.patch_size)
+            if dim_ctx is not None:
+                predictor_config.setdefault("dim_ctx", int(dim_ctx))
+            self.predictor = PREDICTORS.build(predictor_config)
+        self.register_buffer("image_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+        self.register_buffer("image_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+        self._reference_model: Optional[nn.Module] = None
+
+    @property
+    def device(self) -> torch.device:
+        """Return the device hosting PXDepth learnable parameters.
+
+        No inputs are required. The value is inferred from the model's first
+        parameter and is used when moving inference inputs or reference models.
+
+        Returns:
+            ``torch.device`` for the current model placement.
+        """
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Return the storage dtype of PXDepth learnable parameters.
+
+        No inputs are required. This reports parameter storage, which is
+        independent from local autocast contexts used inside attention.
+
+        Returns:
+            ``torch.dtype`` of the model's first parameter.
+        """
+        return next(self.parameters()).dtype
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        path_or_repo: Union[str, Path, IO[bytes]],
+        model_kwargs: Optional[Dict[str, Any]] = None,
+        strict: bool = True,
+        **hf_kwargs: Any,
+    ) -> "PXDepth":
+        """Create a model from a local or Hugging Face ``model.pt`` checkpoint.
+
+        Args:
+            path_or_repo: Local checkpoint path, binary file object, or Hugging
+                Face model repository identifier.
+            model_kwargs: Optional constructor overrides applied after reading
+                ``model_config`` from the checkpoint.
+            strict: Forwarded to ``load_state_dict``. Published checkpoints
+                should use the default exact matching.
+            **hf_kwargs: Additional keyword arguments forwarded to
+                ``huggingface_hub.hf_hub_download`` for remote repositories.
+
+        Returns:
+            Initialized :class:`PXDepth` instance on CPU.
+        """
+        return load_pretrained(
+            cls,
+            path_or_repo,
+            model_kwargs=model_kwargs,
+            strict=strict,
+            **hf_kwargs,
+        )
+
+    def init_weights(self) -> None:
+        """Initialize the Global Context Encoder from official DINOv2 weights.
+
+        Predictor parameters retain the initialization created by their own
+        module constructors.
+
+        Returns:
+            ``None``. Encoder parameters are updated in place.
+        """
+        self.encoder.init_weights()
+
+    def enable_gradient_checkpointing(self) -> None:
+        """Enable activation checkpointing in both encoder and predictor.
+
+        This reduces saved activation memory during backward at the cost of
+        recomputing transformer blocks.
+
+        Returns:
+            ``None``. Child module runtime behavior is updated in place.
+        """
+        self.encoder.enable_gradient_checkpointing()
+        self.predictor.enable_gradient_checkpointing()
+
+    def enable_pytorch_native_sdpa(self) -> None:
+        """Enable the optimized SDPA attention path in the DINOv2 backbone.
+
+        Decoder CM-PiT attention already uses PyTorch SDPA directly and is not
+        modified by this method.
+
+        Returns:
+            ``None``. Encoder attention modules are wrapped in place.
+        """
+        self.encoder.enable_pytorch_native_sdpa()
+
+    def _remap(self, depth: torch.Tensor) -> torch.Tensor:
+        """Apply the configured output activation to raw depth predictions.
+
+        Args:
+            depth: Raw normalized-depth tensor with arbitrary batch/spatial
+                shape, normally ``[B, H, W]``.
+
+        Returns:
+            Tensor with the same shape. The released ``linear`` setting returns
+            the input unchanged.
+        """
+        return F.elu(depth) if self.remap_output == "elu" else depth
+
+    def forward(
+        self,
+        image: torch.Tensor,
+        use_fp16: bool = False,
+        use_fp32: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """Run the network without metric-scale alignment.
+
+        Args:
+            image: RGB tensor ``[B, 3, H, W]`` with values in ``[0, 1]``. ``H``
+                and ``W`` must be divisible by the encoder patch size.
+            use_fp16: Run attention-heavy encoder and predictor regions under
+                FP16 autocast.
+            use_fp32: Disable reduced-precision autocast. It is mutually
+                exclusive with ``use_fp16``.
+
+        Returns:
+            Dictionary with normalized log-depth ``depth`` and finite-depth
+            probability ``mask``, both FP32 tensors ``[B, H, W]``.
+        """
+        height, width = image.shape[-2:]
+        if height % self.patch_size or width % self.patch_size:
+            raise ValueError(
+                f"Input resolution ({height}, {width}) must be divisible by patch size {self.patch_size}"
+            )
+        dtype = inference_dtype(use_fp16=use_fp16, use_fp32=use_fp32)
+
+        with full_precision(image.device):
+            image_norm = (image.float() - self.image_mean.float()) / self.image_std.float()
+        with reduced_precision(image.device, dtype):
+            context = self.encoder(image, height // self.patch_size, width // self.patch_size)
+        context = context.flatten(2).permute(0, 2, 1).contiguous()
+        depth, mask = self.predictor(image_norm, context, autocast_dtype=dtype)
+
+        with full_precision(image.device):
+            depth = self._remap(depth.float().squeeze(1))
+            mask = mask.float().squeeze(1).sigmoid()
+        return {"depth": depth, "mask": mask}
+
+    def infer(
+        self,
+        image: torch.Tensor,
+        gt_depth: Optional[torch.Tensor] = None,
+        intrinsics: Optional[torch.Tensor] = None,
+        fov_x: Optional[Union[float, torch.Tensor]] = None,
+        ref_image: Optional[torch.Tensor] = None,
+        apply_mask: bool = True,
+        use_fp16: bool = True,
+        use_fp32: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        """Recover metric-scale depth, validity, intrinsics, and 3D points.
+
+        Raw normalized log-depth is affine-aligned in log space to ``gt_depth``
+        when supplied, otherwise to a lazily loaded MoGe-2 reference. Alignment
+        parameters are estimated on a 64x64 nearest-resized valid subset. The
+        aligned depth is exponentiated and back-projected with normalized camera
+        intrinsics.
+
+        Args:
+            image: RGB tensor ``[3,H,W]`` or batch ``[B,3,H,W]`` in ``[0,1]``.
+            gt_depth: Optional reference depth ``[H,W]`` or ``[B,H,W]``. Finite
+                positive pixels define log-space alignment.
+            intrinsics: Optional normalized camera matrices ``[3,3]`` or
+                ``[B,3,3]`` corresponding to ``gt_depth``.
+            fov_x: Optional horizontal field of view in degrees, scalar or
+                tensor ``[B]``, used when intrinsics are unavailable.
+            ref_image: Optional original-resolution RGB tensor used only by the
+                reference model; PXDepth still consumes ``image``.
+            apply_mask: Replace invalid predicted depth/points with infinity.
+            use_fp16: Use FP16 for attention-heavy model regions.
+            use_fp32: Force those regions to FP32 and override the BF16 default.
+
+        Returns:
+            Dictionary containing aligned ``depth`` ``[B,H,W]``, boolean
+            ``mask`` ``[B,H,W]``, point map ``points`` ``[B,H,W,3]``, normalized
+            ``intrinsics`` ``[B,3,3]``, and horizontal ``fov_x`` ``[B]``. For an
+            unbatched input, the leading batch dimension is removed.
+        """
+        return infer_model(
+            self,
+            image,
+            gt_depth=gt_depth,
+            intrinsics=intrinsics,
+            fov_x=fov_x,
+            ref_image=ref_image,
+            apply_mask=apply_mask,
+            use_fp16=use_fp16,
+            use_fp32=use_fp32,
+        )

pxdepth/model/Pixel_Space_Depth_Predictor.py

@@ -0,0 +1,274 @@
+"""Pixel-Space Depth Predictor that preserves the dense image lattice.
+
+Normalized RGB is embedded with a 1x1 projection and processed by shared
+CM-PiT trunk blocks before branching into depth and finite-mask predictors.
+Linear patch compression is used only within transformer blocks, after which
+features are expanded back to per-pixel tokens for dense output heads.
+"""
+
+from typing import Iterable, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+from ..registry import PREDICTORS
+from .CM_PiT import CMPiTBlock
+from .RoPE import PositionGetter, RotaryPositionEmbedding2D
+from .precision import full_precision, reduced_precision
+
+
+@PREDICTORS.register()
+class PixelSpaceDepthPredictor(nn.Module):
+    """Pixel-Space Depth Predictor built from cascaded CM-PiT blocks.
+
+    A 1x1 projection first embeds normalized RGB into dense pixel features.
+    Shared trunk blocks refine those features, after which independent depth
+    and validity branches predict normalized log-depth and finite-depth logits.
+    No convolution larger than 1x1 is applied to the pixel representation.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        dim_ctx: int = 1024,
+        attn_dim: int = 1536,
+        ctx_patch_size: int = 14,
+        dim_pix: int = 16,
+        trunk_patch_size: int = 14,
+        depth_patch_size: int = 7,
+        mask_patch_size: int = 14,
+        num_heads: int = 24,
+        trunk_depth: int = 4,
+        depth_depth: int = 4,
+        mask_depth: int = 2,
+        mlp_ratio: float = 4.0,
+        qk_norm: bool = True,
+        rope_frequency: float = 100.0,
+        eps: float = 1e-6,
+        gradient_checkpointing: bool = True,
+    ) -> None:
+        """Construct the shared trunk and two prediction branches.
+
+        Args:
+            in_channels: Number of image channels, equal to three for RGB.
+            dim_ctx: Global context-token channel count ``C_ctx``.
+            attn_dim: Channel count ``D`` after linear patch compression.
+            ctx_patch_size: Encoder patch size ``P_ctx`` in image pixels.
+            dim_pix: Channel count ``C_pix`` of the dense pixel feature map.
+            trunk_patch_size: Attention patch size used by shared trunk blocks.
+            depth_patch_size: Attention patch size used by depth blocks.
+            mask_patch_size: Attention patch size used by validity-mask blocks.
+            num_heads: Number of gated-attention heads.
+            trunk_depth: Number of shared CM-PiT blocks.
+            depth_depth: Number of depth-branch CM-PiT blocks.
+            mask_depth: Number of validity-branch CM-PiT blocks.
+            mlp_ratio: SwiGLU expansion ratio inside every block.
+            qk_norm: Enable FP32 RMSNorm for attention queries and keys.
+            rope_frequency: Base frequency of the shared 2D RoPE module.
+            eps: Numerical epsilon for normalization layers.
+            gradient_checkpointing: Recompute CM-PiT blocks during backward to
+                reduce activation memory.
+
+        Returns:
+            ``None``. The complete pixel predictor is registered on the module.
+        """
+        super().__init__()
+        if attn_dim % num_heads != 0:
+            raise ValueError(f"attn_dim ({attn_dim}) must be divisible by num_heads ({num_heads})")
+        for name, patch_size in {
+            "trunk_patch_size": trunk_patch_size,
+            "depth_patch_size": depth_patch_size,
+            "mask_patch_size": mask_patch_size,
+        }.items():
+            if patch_size <= 0 or ctx_patch_size % patch_size != 0:
+                raise ValueError(
+                    f"{name} ({patch_size}) must be positive and divide ctx_patch_size ({ctx_patch_size})"
+                )
+
+        self.dim_ctx = dim_ctx
+        self.dim_pix = dim_pix
+        self.attn_dim = attn_dim
+        self.ctx_patch_size = ctx_patch_size
+        self.trunk_patch_size = trunk_patch_size
+        self.depth_patch_size = depth_patch_size
+        self.mask_patch_size = mask_patch_size
+        self.gradient_checkpointing = gradient_checkpointing
+
+        self.pos = PositionGetter()
+        self.rope = RotaryPositionEmbedding2D(frequency=rope_frequency)
+        self.input_proj = nn.Conv2d(in_channels, dim_pix, kernel_size=1, bias=True)
+
+        block_args = dict(
+            dim_ctx=dim_ctx,
+            ctx_patch_size=ctx_patch_size,
+            dim_pix=dim_pix,
+            attn_dim=attn_dim,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            qk_norm=qk_norm,
+            rope=self.rope,
+            eps=eps,
+        )
+        self.trunk_blocks = nn.ModuleList(
+            CMPiTBlock(patch_size=trunk_patch_size, **block_args) for _ in range(trunk_depth)
+        )
+        self.depth_blocks = nn.ModuleList(
+            CMPiTBlock(patch_size=depth_patch_size, **block_args) for _ in range(depth_depth)
+        )
+        self.mask_blocks = nn.ModuleList(
+            CMPiTBlock(patch_size=mask_patch_size, **block_args) for _ in range(mask_depth)
+        )
+        self.depth_head = nn.Conv2d(dim_pix, 1, kernel_size=1, bias=True)
+        self.mask_head = nn.Conv2d(dim_pix, 1, kernel_size=1, bias=True)
+        self.reset_parameters()
+
+    def _blocks(self) -> Iterable[CMPiTBlock]:
+        """Iterate over every CM-PiT block in execution-independent order.
+
+        Returns:
+            Iterable containing shared trunk, depth, and mask blocks. The method
+            takes no tensor inputs and is used for parameter initialization.
+        """
+        return (*self.trunk_blocks, *self.depth_blocks, *self.mask_blocks)
+
+    def reset_parameters(self) -> None:
+        """Initialize projections and start adaptive modulation at identity.
+
+        Linear and 1x1 convolution weights use Xavier uniform initialization.
+        Normalization scales start at one. The final adaptive-normalization
+        projections are zeroed so every CM-PiT residual branch initially has
+        zero modulation and zero gate.
+
+        Returns:
+            ``None``. Parameters are modified in place.
+        """
+        def init(module: nn.Module) -> None:
+            """Initialize one child module visited by :meth:`nn.Module.apply`.
+
+            Args:
+                module: Child ``nn.Module`` to initialize in place.
+
+            Returns:
+                ``None``.
+            """
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, (nn.LayerNorm, nn.RMSNorm)):
+                if module.weight is not None:
+                    nn.init.ones_(module.weight)
+                if getattr(module, "bias", None) is not None:
+                    nn.init.zeros_(module.bias)
+
+        self.apply(init)
+        for block in self._blocks():
+            nn.init.zeros_(block.ada_norm.proj[-1].weight)
+            nn.init.zeros_(block.ada_norm.proj[-1].bias)
+
+    def enable_gradient_checkpointing(self) -> None:
+        """Enable activation recomputation for CM-PiT blocks.
+
+        The flag is consulted only while the module is in training mode.
+
+        Returns:
+            ``None``. The runtime flag is changed in place.
+        """
+        self.gradient_checkpointing = True
+
+    def disable_gradient_checkpointing(self) -> None:
+        """Disable activation recomputation for CM-PiT blocks.
+
+        Subsequent training forwards retain block activations for backward.
+
+        Returns:
+            ``None``. The runtime flag is changed in place.
+        """
+        self.gradient_checkpointing = False
+
+    def _position(self, batch: int, height: int, width: int, patch_size: int, device: torch.device):
+        """Create cached 2D coordinates for one decoder patch grid.
+
+        Args:
+            batch: Batch size ``B``.
+            height: Dense image-feature height ``H``.
+            width: Dense image-feature width ``W``.
+            patch_size: Block patch side length ``P``.
+            device: Device on which coordinates are allocated.
+
+        Returns:
+            Integer position tensor ``[B, (H/P)*(W/P), 2]``.
+        """
+        return self.pos(batch, height // patch_size, width // patch_size, device=device).to(device)
+
+    def _run(
+        self,
+        x: torch.Tensor,
+        blocks: nn.ModuleList,
+        ctx: torch.Tensor,
+        pos: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run a sequence of CM-PiT blocks with optional checkpointing.
+
+        Args:
+            x: Dense pixel features ``[B, C_pix, H, W]``.
+            blocks: Ordered CM-PiT block collection for one branch.
+            ctx: Global context tokens ``[B, N_ctx, C_ctx]``.
+            pos: 2D positions ``[B, N, 2]`` matching the blocks' patch grid.
+
+        Returns:
+            Refined dense features ``[B, C_pix, H, W]``.
+        """
+        for block in blocks:
+            if self.training and self.gradient_checkpointing:
+                x = checkpoint(block, x, ctx, pos, use_reentrant=False)
+            else:
+                x = block(x, ctx, pos)
+        return x
+
+    def forward(
+        self,
+        image: torch.Tensor,
+        ctx: torch.Tensor,
+        autocast_dtype: Optional[torch.dtype] = torch.bfloat16,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Predict normalized log-depth and finite-depth logits from RGB.
+
+        Args:
+            image: ImageNet-normalized RGB tensor ``[B, 3, H, W]``.
+            ctx: Global context tokens ``[B, (H/P_ctx)*(W/P_ctx), C_ctx]``.
+            autocast_dtype: Decoder attention dtype. Use ``None`` for FP32,
+                ``torch.float16`` for FP16, or ``torch.bfloat16`` for BF16.
+
+        Returns:
+            depth: Raw normalized log-depth tensor ``[B, 1, H, W]``.
+            mask: Raw finite-depth logit tensor ``[B, 1, H, W]``.
+        """
+        batch, _, height, width = image.shape
+        p_ctx = self.ctx_patch_size
+        if height % p_ctx != 0 or width % p_ctx != 0:
+            raise ValueError(f"Input resolution ({height}, {width}) must be divisible by {p_ctx}")
+        expected = (height // p_ctx) * (width // p_ctx)
+        if tuple(ctx.shape) != (batch, expected, self.dim_ctx):
+            raise ValueError(
+                f"Context shape {tuple(ctx.shape)} does not match ({batch}, {expected}, {self.dim_ctx})"
+            )
+
+        with full_precision(image.device):
+            pix = self.input_proj(image.float())
+
+        with reduced_precision(image.device, autocast_dtype):
+            trunk_pos = self._position(batch, height, width, self.trunk_patch_size, image.device)
+            pix = self._run(pix, self.trunk_blocks, ctx, trunk_pos)
+
+            depth_pos = self._position(batch, height, width, self.depth_patch_size, image.device)
+            depth_feat = self._run(pix, self.depth_blocks, ctx, depth_pos)
+
+            mask_pos = self._position(batch, height, width, self.mask_patch_size, image.device)
+            mask_feat = self._run(pix, self.mask_blocks, ctx, mask_pos)
+
+        with full_precision(image.device):
+            depth = self.depth_head(depth_feat.float())
+            mask = self.mask_head(mask_feat.float())
+        return depth, mask

pxdepth/model/RoPE.py

@@ -0,0 +1,208 @@
+"""Cached two-dimensional rotary position embeddings for image-token grids.
+
+Positions are represented as integer ``(row, column)`` pairs. Half of every
+attention-head feature is rotated by row and the other half by column, allowing
+CM-PiT attention to operate at arbitrary patch-grid aspect ratios.
+"""
+
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+
+# Implementation of 2D Rotary Position Embeddings (RoPE).
+
+# This module provides a clean implementation of 2D Rotary Position Embeddings,
+# which extends the original RoPE concept to handle 2D spatial positions.
+
+# Inspired by:
+#         https://github.com/meta-llama/codellama/blob/main/llama/model.py
+#         https://github.com/naver-ai/rope-vit
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Tuple
+
+
+class PositionGetter:
+    """Generates and caches 2D spatial positions for patches in a grid.
+
+    This class efficiently manages the generation of spatial coordinates for patches
+    in a 2D grid, caching results to avoid redundant computations.
+
+    Attributes:
+        position_cache: Dictionary storing precomputed position tensors for different
+            grid dimensions.
+    """
+
+    def __init__(self):
+        """Initialize the position generator with an empty grid cache.
+
+        The constructor takes no inputs. Coordinate tensors are generated lazily
+        for each unique ``(height,width)`` requested by decoder blocks.
+
+        Returns:
+            ``None``.
+        """
+        self.position_cache: Dict[Tuple[int, int], torch.Tensor] = {}
+
+    def __call__(self, batch_size: int, height: int, width: int, device: torch.device) -> torch.Tensor:
+        """Generates spatial positions for a batch of patches.
+
+        Args:
+            batch_size: Number of samples in the batch.
+            height: Height of the grid in patches.
+            width: Width of the grid in patches.
+            device: Target device for the position tensor.
+
+        Returns:
+            Tensor of shape (batch_size, height*width, 2) containing y,x coordinates
+            for each position in the grid, repeated for each batch item.
+        """
+        if (height, width) not in self.position_cache:
+            y_coords = torch.arange(height, device=device)
+            x_coords = torch.arange(width, device=device)
+            positions = torch.cartesian_prod(y_coords, x_coords)
+            self.position_cache[height, width] = positions
+
+        cached_positions = self.position_cache[height, width]
+        return cached_positions.view(1, height * width, 2).expand(batch_size, -1, -1).clone()
+
+
+class RotaryPositionEmbedding2D(nn.Module):
+    """2D Rotary Position Embedding implementation.
+
+    This module applies rotary position embeddings to input tokens based on their
+    2D spatial positions. It handles the position-dependent rotation of features
+    separately for vertical and horizontal dimensions.
+
+    Args:
+        frequency: Base frequency for the position embeddings. Default: 100.0
+        scaling_factor: Scaling factor for frequency computation. Default: 1.0
+
+    Attributes:
+        base_frequency: Base frequency for computing position embeddings.
+        scaling_factor: Factor to scale the computed frequencies.
+        frequency_cache: Cache for storing precomputed frequency components.
+    """
+
+    def __init__(self, frequency: float = 100.0, scaling_factor: float = 1.0):
+        """Initialize frequency settings and an empty component cache.
+
+        Args:
+            frequency: Base controlling the geometric frequency progression.
+            scaling_factor: Reserved multiplicative frequency scale retained
+                for checkpoint and API compatibility.
+
+        Returns:
+            ``None``. Frequency tables are generated lazily during ``forward``.
+        """
+        super().__init__()
+        self.base_frequency = frequency
+        self.scaling_factor = scaling_factor
+        self.frequency_cache: Dict[Tuple, Tuple[torch.Tensor, torch.Tensor]] = {}
+
+    def _compute_frequency_components(
+        self, dim: int, seq_len: int, device: torch.device, dtype: torch.dtype
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Computes frequency components for rotary embeddings.
+
+        Args:
+            dim: Feature dimension (must be even).
+            seq_len: Maximum sequence length.
+            device: Target device for computations.
+            dtype: Data type for the computed tensors.
+
+        Returns:
+            Tuple of (cosine, sine) tensors for frequency components.
+        """
+        cache_key = (dim, seq_len, device, dtype)
+        if cache_key not in self.frequency_cache:
+            # Compute frequency bands
+            exponents = torch.arange(0, dim, 2, device=device).float() / dim
+            inv_freq = 1.0 / (self.base_frequency**exponents)
+
+            # Generate position-dependent frequencies
+            positions = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
+            angles = torch.einsum("i,j->ij", positions, inv_freq)
+
+            # Compute and cache frequency components
+            angles = angles.to(dtype)
+            angles = torch.cat((angles, angles), dim=-1)
+            cos_components = angles.cos().to(dtype)
+            sin_components = angles.sin().to(dtype)
+            self.frequency_cache[cache_key] = (cos_components, sin_components)
+
+        return self.frequency_cache[cache_key]
+
+    @staticmethod
+    def _rotate_features(x: torch.Tensor) -> torch.Tensor:
+        """Performs feature rotation by splitting and recombining feature dimensions.
+
+        Args:
+            x: Tensor ``[..., D]`` whose final dimension is split in half.
+
+        Returns:
+            Rotated tensor with the same shape and dtype as ``x``.
+        """
+        feature_dim = x.shape[-1]
+        x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def _apply_1d_rope(
+        self, tokens: torch.Tensor, positions: torch.Tensor, cos_comp: torch.Tensor, sin_comp: torch.Tensor
+    ) -> torch.Tensor:
+        """Applies 1D rotary position embeddings along one dimension.
+
+        Args:
+            tokens: One-axis head features ``[B, H, N, D_axis]``.
+            positions: Integer position indices ``[B, N]``.
+            cos_comp: Cosine lookup table ``[L, D_axis]``.
+            sin_comp: Sine lookup table ``[L, D_axis]``.
+
+        Returns:
+            Rotated features ``[B, H, N, D_axis]``.
+        """
+        # Embed positions with frequency components
+        cos = F.embedding(positions, cos_comp)[:, None, :, :]
+        sin = F.embedding(positions, sin_comp)[:, None, :, :]
+
+        # Apply rotation
+        return (tokens * cos) + (self._rotate_features(tokens) * sin)
+
+    def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """Applies 2D rotary position embeddings to input tokens.
+
+        Args:
+            tokens: Input tensor of shape (batch_size, n_heads, n_tokens, dim).
+                   The feature dimension (dim) must be divisible by 4.
+            positions: Position tensor of shape (batch_size, n_tokens, 2) containing
+                      the y and x coordinates for each token.
+
+        Returns:
+            Tensor of same shape as input with applied 2D rotary position embeddings.
+
+        Raises:
+            AssertionError: If input dimensions are invalid or positions are malformed.
+        """
+        # Validate inputs
+        assert tokens.size(-1) % 2 == 0, "Feature dimension must be even"
+        assert positions.ndim == 3 and positions.shape[-1] == 2, "Positions must have shape (batch_size, n_tokens, 2)"
+
+        # Compute feature dimension for each spatial direction
+        feature_dim = tokens.size(-1) // 2
+
+        # Get frequency components
+        max_position = int(positions.max()) + 1
+        cos_comp, sin_comp = self._compute_frequency_components(feature_dim, max_position, tokens.device, tokens.dtype)
+
+        # Split features for vertical and horizontal processing
+        vertical_features, horizontal_features = tokens.chunk(2, dim=-1)
+
+        # Apply RoPE separately for each dimension
+        vertical_features = self._apply_1d_rope(vertical_features, positions[..., 0], cos_comp, sin_comp)
+        horizontal_features = self._apply_1d_rope(horizontal_features, positions[..., 1], cos_comp, sin_comp)
+
+        # Combine processed features
+        return torch.cat((vertical_features, horizontal_features), dim=-1)

pxdepth/model/__init__.py

@@ -0,0 +1,12 @@
+"""Public model namespace for PXDepth.
+
+Only the complete :class:`PXDepth` network is re-exported here. Internal
+encoder, predictor, attention, and positional-encoding modules remain available
+for development without becoming part of the stable package-level API.
+"""
+
+from .PXDepth import PXDepth
+from .Global_Context_Encoder import GlobalContextEncoder
+from .Pixel_Space_Depth_Predictor import PixelSpaceDepthPredictor
+
+__all__ = ["PXDepth", "GlobalContextEncoder", "PixelSpaceDepthPredictor"]

pxdepth/model/checkpoint.py

@@ -0,0 +1,83 @@
+"""Load model-only checkpoints written in the public PXDepth format."""
+
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Dict, IO, Optional, Type, TypeVar, Union
+
+import torch
+import torch.nn as nn
+from huggingface_hub import hf_hub_download
+
+
+ModelT = TypeVar("ModelT", bound=nn.Module)
+
+
+def _merge(base: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any]:
+    """Recursively merge model-constructor overrides into a copied config.
+
+    Args:
+        base: Original nested model configuration. The dictionary is not
+            modified.
+        update: User overrides. Nested dictionaries update individual fields,
+            while lists and scalar values replace their counterparts.
+
+    Returns:
+        A new merged dictionary suitable for constructing the model.
+    """
+    result = deepcopy(base)
+    for key, value in update.items():
+        if isinstance(result.get(key), dict) and isinstance(value, dict):
+            result[key] = _merge(result[key], value)
+        else:
+            result[key] = deepcopy(value)
+    return result
+
+
+def load_pretrained(
+    model_class: Type[ModelT],
+    path_or_repo: Union[str, Path, IO[bytes]],
+    model_kwargs: Optional[Dict[str, Any]] = None,
+    strict: bool = True,
+    **hf_kwargs: Any,
+) -> ModelT:
+    """Construct a model from a local or Hugging Face ``model.pt`` file.
+
+    Args:
+        model_class: Model class whose constructor accepts the canonical public
+            configuration.
+        path_or_repo: Existing local checkpoint path, binary file object, or
+            Hugging Face model repository identifier.
+        model_kwargs: Optional nested constructor overrides. Nested encoder or
+            predictor fields are merged without discarding sibling settings.
+        strict: Forwarded to :meth:`torch.nn.Module.load_state_dict`. Published
+            checkpoints should normally use the default strict loading.
+        **hf_kwargs: Extra arguments forwarded to ``hf_hub_download`` when
+            ``path_or_repo`` is a repository identifier.
+
+    Returns:
+        An initialized model instance on CPU.
+    """
+    path = Path(path_or_repo) if isinstance(path_or_repo, (str, Path)) else None
+    if path is not None and path.exists():
+        checkpoint_path: Union[Path, IO[bytes]] = path
+    elif isinstance(path_or_repo, str):
+        checkpoint_path = Path(
+            hf_hub_download(path_or_repo, repo_type="model", filename="model.pt", **hf_kwargs)
+        )
+    else:
+        checkpoint_path = path_or_repo
+
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+    if "model_config" not in checkpoint or "model" not in checkpoint:
+        raise ValueError("PXDepth checkpoints must contain 'model_config' and 'model'.")
+    model_config = deepcopy(checkpoint["model_config"])
+    model_type = model_config.pop("type", None)
+    if model_type != model_class.__name__:
+        raise ValueError(
+            f"Expected a {model_class.__name__} checkpoint, got model type {model_type!r}."
+        )
+    if model_kwargs:
+        model_config = _merge(model_config, model_kwargs)
+    model = model_class(**model_config)
+    model.load_state_dict(checkpoint["model"], strict=strict)
+    return model

pxdepth/model/dinov2/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+__version__ = "0.0.1"

pxdepth/model/dinov2/hub/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.

pxdepth/model/dinov2/hub/backbones.py

@@ -0,0 +1,156 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+from enum import Enum
+from typing import Union
+
+import torch
+
+from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name
+
+
+class Weights(Enum):
+    LVD142M = "LVD142M"
+
+
+def _make_dinov2_model(
+    *,
+    arch_name: str = "vit_large",
+    img_size: int = 518,
+    patch_size: int = 14,
+    init_values: float = 1.0,
+    ffn_layer: str = "mlp",
+    block_chunks: int = 0,
+    num_register_tokens: int = 0,
+    interpolate_antialias: bool = False,
+    interpolate_offset: float = 0.1,
+    pretrained: bool = True,
+    weights: Union[Weights, str] = Weights.LVD142M,
+    **kwargs,
+):
+    from ..models import vision_transformer as vits
+
+    if isinstance(weights, str):
+        try:
+            weights = Weights[weights]
+        except KeyError:
+            raise AssertionError(f"Unsupported weights: {weights}")
+
+    model_base_name = _make_dinov2_model_name(arch_name, patch_size)
+    vit_kwargs = dict(
+        img_size=img_size,
+        patch_size=patch_size,
+        init_values=init_values,
+        ffn_layer=ffn_layer,
+        block_chunks=block_chunks,
+        num_register_tokens=num_register_tokens,
+        interpolate_antialias=interpolate_antialias,
+        interpolate_offset=interpolate_offset,
+    )
+    vit_kwargs.update(**kwargs)
+    model = vits.__dict__[arch_name](**vit_kwargs)
+
+    if pretrained:
+        model_full_name = _make_dinov2_model_name(arch_name, patch_size, num_register_tokens)
+        url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_pretrain.pth"
+        state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu")
+        model.load_state_dict(state_dict, strict=True)
+
+    return model
+
+
+def dinov2_vits14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+    """
+    DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(arch_name="vit_small", pretrained=pretrained, weights=weights, **kwargs)
+
+
+def dinov2_vitb14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+    """
+    DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(arch_name="vit_base", pretrained=pretrained, weights=weights, **kwargs)
+
+
+def dinov2_vitl14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+    """
+    DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(arch_name="vit_large", pretrained=pretrained, weights=weights, **kwargs)
+
+
+def dinov2_vitg14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+    """
+    DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_giant2",
+        ffn_layer="swiglufused",
+        weights=weights,
+        pretrained=pretrained,
+        **kwargs,
+    )
+
+
+def dinov2_vits14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+    """
+    DINOv2 ViT-S/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_small",
+        pretrained=pretrained,
+        weights=weights,
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )
+
+
+def dinov2_vitb14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+    """
+    DINOv2 ViT-B/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_base",
+        pretrained=pretrained,
+        weights=weights,
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )
+
+
+def dinov2_vitl14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+    """
+    DINOv2 ViT-L/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_large",
+        pretrained=pretrained,
+        weights=weights,
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )
+
+
+def dinov2_vitg14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs):
+    """
+    DINOv2 ViT-g/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_giant2",
+        ffn_layer="swiglufused",
+        weights=weights,
+        pretrained=pretrained,
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )

pxdepth/model/dinov2/hub/utils.py

@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import itertools
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+_DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2"
+
+
+def _make_dinov2_model_name(arch_name: str, patch_size: int, num_register_tokens: int = 0) -> str:
+    compact_arch_name = arch_name.replace("_", "")[:4]
+    registers_suffix = f"_reg{num_register_tokens}" if num_register_tokens else ""
+    return f"dinov2_{compact_arch_name}{patch_size}{registers_suffix}"
+
+
+class CenterPadding(nn.Module):
+    def __init__(self, multiple):
+        super().__init__()
+        self.multiple = multiple
+
+    def _get_pad(self, size):
+        new_size = math.ceil(size / self.multiple) * self.multiple
+        pad_size = new_size - size
+        pad_size_left = pad_size // 2
+        pad_size_right = pad_size - pad_size_left
+        return pad_size_left, pad_size_right
+
+    @torch.inference_mode()
+    def forward(self, x):
+        pads = list(itertools.chain.from_iterable(self._get_pad(m) for m in x.shape[:1:-1]))
+        output = F.pad(x, pads)
+        return output

pxdepth/model/dinov2/layers/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+from .mlp import Mlp
+from .patch_embed import PatchEmbed
+from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
+from .block import NestedTensorBlock
+from .attention import MemEffAttention

pxdepth/model/dinov2/layers/attention.py

@@ -0,0 +1,151 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import logging
+import os
+import warnings
+
+import torch.nn.functional as F
+from torch import Tensor
+from torch import nn
+import torch
+
+
+logger = logging.getLogger("dinov2")
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import memory_efficient_attention, unbind
+
+        XFORMERS_AVAILABLE = True
+        # warnings.warn("xFormers is available (Attention)")
+    else:
+        # warnings.warn("xFormers is disabled (Attention)")
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+    # warnings.warn("xFormers is not available (Attention)")
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_func
+    FLASH_ATTN_AVAILABLE = True
+except Exception:
+    flash_attn_func = None
+    FLASH_ATTN_AVAILABLE = False
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    # # Deprecated implementation, extremely slow
+    # def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+    #     B, N, C = x.shape
+    #     qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+    #     q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+    #     attn = q @ k.transpose(-2, -1)
+    #     attn = attn.softmax(dim=-1)
+    #     attn = self.attn_drop(attn)
+    #     x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+    #     x = self.proj(x)
+    #     x = self.proj_drop(x)
+    #     return x
+
+    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)  # (3, B, H, N, C // H)
+
+        q, k, v = qkv.unbind(0)      # (B, H, N, C // H)
+
+        use_flash_attn = (
+            FLASH_ATTN_AVAILABLE
+            and attn_bias is None
+            and q.is_cuda
+            and q.dtype in (torch.float16, torch.bfloat16)
+        )
+        if use_flash_attn:
+            q_f = q.permute(0, 2, 1, 3).contiguous()
+            k_f = k.permute(0, 2, 1, 3).contiguous()
+            v_f = v.permute(0, 2, 1, 3).contiguous()
+            x = flash_attn_func(
+                q_f,
+                k_f,
+                v_f,
+                dropout_p=0.0,
+                softmax_scale=self.scale,
+                causal=False,
+            )
+            x = x.reshape(B, N, C)
+        else:
+            x = F.scaled_dot_product_attention(q, k, v, attn_bias)
+            x = x.permute(0, 2, 1, 3).reshape(B, N, C)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+class MemEffAttention(Attention):
+    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+        B, N, C = x.shape
+        use_flash_attn = (
+            FLASH_ATTN_AVAILABLE
+            and attn_bias is None
+            and x.is_cuda
+            and x.dtype in (torch.float16, torch.bfloat16)
+        )
+        if use_flash_attn:
+            qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            q, k, v = qkv.unbind(2)
+            x = flash_attn_func(
+                q.contiguous(),
+                k.contiguous(),
+                v.contiguous(),
+                dropout_p=0.0,
+                softmax_scale=self.scale,
+                causal=False,
+            )
+            x = x.reshape([B, N, C])
+
+            x = self.proj(x)
+            x = self.proj_drop(x)
+            return x
+
+        if not XFORMERS_AVAILABLE:
+            if attn_bias is not None:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return super().forward(x, attn_bias=attn_bias)
+
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

pxdepth/model/dinov2/layers/block.py

@@ -0,0 +1,259 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+import logging
+import os
+from typing import Callable, List, Any, Tuple, Dict
+import warnings
+
+import torch
+from torch import nn, Tensor
+
+from .attention import Attention, MemEffAttention
+from .drop_path import DropPath
+from .layer_scale import LayerScale
+from .mlp import Mlp
+
+
+logger = logging.getLogger("dinov2")
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import fmha, scaled_index_add, index_select_cat
+
+        XFORMERS_AVAILABLE = True
+        # warnings.warn("xFormers is available (Block)")
+    else:
+        # warnings.warn("xFormers is disabled (Block)")
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+    # warnings.warn("xFormers is not available (Block)")
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = Attention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+    ) -> None:
+        super().__init__()
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    def forward(self, x: Tensor) -> Tensor:
+        def attn_residual_func(x: Tensor) -> Tensor:
+            return self.ls1(self.attn(self.norm1(x)))
+
+        def ffn_residual_func(x: Tensor) -> Tensor:
+            return self.ls2(self.mlp(self.norm2(x)))
+
+        if self.training and self.sample_drop_ratio > 0.1:
+            # the overhead is compensated only for a drop path rate larger than 0.1
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+        elif self.training and self.sample_drop_ratio > 0.0:
+            x = x + self.drop_path1(attn_residual_func(x))
+            x = x + self.drop_path1(ffn_residual_func(x))
+        else:
+            x = x + attn_residual_func(x)
+            x = x + ffn_residual_func(x)
+        return x
+
+
+def drop_add_residual_stochastic_depth(
+    x: Tensor,
+    residual_func: Callable[[Tensor], Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> Tensor:
+    # 1) extract subset using permutation
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    x_subset = x[brange]
+
+    # 2) apply residual_func to get residual
+    residual = residual_func(x_subset)
+
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+
+    residual_scale_factor = b / sample_subset_size
+
+    # 3) add the residual
+    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    residual_scale_factor = b / sample_subset_size
+    return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+    if scaling_vector is None:
+        x_flat = x.flatten(1)
+        residual = residual.flatten(1)
+        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    else:
+        x_plus_residual = scaled_index_add(
+            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+        )
+    return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+    """
+    this will perform the index select, cat the tensors, and provide the attn_bias from cache
+    """
+    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+    if all_shapes not in attn_bias_cache.keys():
+        seqlens = []
+        for b, x in zip(batch_sizes, x_list):
+            for _ in range(b):
+                seqlens.append(x.shape[1])
+        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+        attn_bias._batch_sizes = batch_sizes
+        attn_bias_cache[all_shapes] = attn_bias
+
+    if branges is not None:
+        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+    else:
+        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+        cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+    return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[Tensor],
+    residual_func: Callable[[Tensor, Any], Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> Tensor:
+    # 1) generate random set of indices for dropping samples in the batch
+    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+    branges = [s[0] for s in branges_scales]
+    residual_scale_factors = [s[1] for s in branges_scales]
+
+    # 2) get attention bias and index+concat the tensors
+    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+    # 3) apply residual_func to get residual, and split the result
+    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+    outputs = []
+    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+    return outputs
+
+
+class NestedTensorBlock(Block):
+    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
+        """
+        x_list contains a list of tensors to nest together and run
+        """
+        assert isinstance(self.attn, MemEffAttention)
+
+        if self.training and self.sample_drop_ratio > 0.0:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.mlp(self.norm2(x))
+
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            return x_list
+        else:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls2(self.mlp(self.norm2(x)))
+
+            attn_bias, x = get_attn_bias_and_cat(x_list)
+            x = x + attn_residual_func(x, attn_bias=attn_bias)
+            x = x + ffn_residual_func(x)
+            return attn_bias.split(x)
+
+    def forward(self, x_or_x_list):
+        if isinstance(x_or_x_list, Tensor):
+            return super().forward(x_or_x_list)
+        elif isinstance(x_or_x_list, list):
+            if not XFORMERS_AVAILABLE:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise AssertionError

pxdepth/model/dinov2/layers/drop_path.py

@@ -0,0 +1,34 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
+
+
+from torch import nn
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)

pxdepth/model/dinov2/layers/layer_scale.py

@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
+
+from typing import Union
+
+import torch
+from torch import Tensor
+from torch import nn
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma

pxdepth/model/dinov2/layers/mlp.py

@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
+
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x