diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..a4e72aa7d623b7b5f587ecbe3d669775e8b768e8 100644 --- a/.gitattributes +++ b/.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 diff --git a/README.md b/README.md index 4a4cd4df874719879dffe153594e5477f56e2bd1..981c64a09dd7d471e43f9ca70a0a1485c39addff 100644 --- a/README.md +++ b/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. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..041b4d0f856f9212db95bf496203becc7d0febf1 --- /dev/null +++ b/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() diff --git a/example_images/KITTI_01.jpg b/example_images/KITTI_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..39ab1f4dd3c68df91e3ed993cc50e4dfbeda9929 --- /dev/null +++ b/example_images/KITTI_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f6a3bc776a4f4b0e20f9e9d0ec77912db738d3b2b17a461425dc2fb9686e2d73 +size 144310 diff --git a/example_images/KITTI_02.jpg b/example_images/KITTI_02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..052fc375caca00f35bd69c8cb7697f9a83e463f2 --- /dev/null +++ b/example_images/KITTI_02.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:738344a8690fa43b21cb6c98bcdbe9ffdbdfeb65629e344b3d6aa48245b742e8 +size 217965 diff --git a/example_images/KITTI_03.jpg b/example_images/KITTI_03.jpg new file mode 100644 index 0000000000000000000000000000000000000000..0397a582ab8f282d6bbed7f02bedcecf831affb1 --- /dev/null +++ b/example_images/KITTI_03.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5a5eac2ac8f8cc1f90cf9e5426984611358baa066a1c3794854eea0752d81c8f +size 214644 diff --git a/example_images/Libary_01.jpg b/example_images/Libary_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b4df6c288b18d4b3f49db85609fbb844ef28e36a --- /dev/null +++ b/example_images/Libary_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ab19a901bb333ccd85b9f5cd17508d4dd83a9f1f7fabda42053c30f6b2a2acf0 +size 47318 diff --git a/example_images/NYUv2_00175.jpg b/example_images/NYUv2_00175.jpg new file mode 100644 index 0000000000000000000000000000000000000000..06cc0a83c75d7fe833b6315b3d1c01bbd7d3a774 --- /dev/null +++ b/example_images/NYUv2_00175.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:01c654b2864893d77e46c73463f196c9046098577832d0cca382edaf33369387 +size 85387 diff --git a/example_images/bird_01.jpg b/example_images/bird_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..f6cf1763ddbd6793fd01f05717094d60b57ec712 --- /dev/null +++ b/example_images/bird_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c0f44133d753a851044d3ab5d135fab07314858ec8277df76c22b044c6556cdb +size 407514 diff --git a/example_images/chair_01.jpg b/example_images/chair_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ed25e69ed86cd8af04bf9990199ff9cfa8899c9c --- /dev/null +++ b/example_images/chair_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:845be5c970de293f7da8b9399b94eaec9ce30f9eba9af185c77f1c58099647a2 +size 139287 diff --git a/example_images/courtyard_01.jpg b/example_images/courtyard_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..166af1e16bde998e5536843c2bca089fcd83f274 --- /dev/null +++ b/example_images/courtyard_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fd404a448717b76eff4b1f9ffffd3a1c09b9a7115ad40dd01716e2f55743d12f +size 56009 diff --git a/example_images/courtyard_02.jpg b/example_images/courtyard_02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d9b5e392994e696415e0b51a44e0a17c29c8c2b7 --- /dev/null +++ b/example_images/courtyard_02.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d12f12f53dcd68860cd35d04e4c881673256fb05dae5942da0a991d89bc3f9f0 +size 59525 diff --git a/example_images/pipe_01.jpg b/example_images/pipe_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..886e5b610530e2ef5ea17cce35e3aa54c10aba8f --- /dev/null +++ b/example_images/pipe_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:820c7504377a971b2541d373d261d645b98f50b61a7b5b2d4a7eba20296e8603 +size 191925 diff --git a/example_images/room_01.jpg b/example_images/room_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1f579dab88ffa7bf090e668a51e897829e37a95d --- /dev/null +++ b/example_images/room_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:523549a80a417a5543d185471cd8373d67605498e8a72f397e77208638184374 +size 100348 diff --git a/example_images/room_02.jpg b/example_images/room_02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d217b17b5e81d31ca9e5e9f14f24c4cb48aaaa15 --- /dev/null +++ b/example_images/room_02.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f2ac3b7dabd4699c2899c997852cf81676df10aa1256e9d107a54a641d4f3d2d +size 85299 diff --git a/example_images/room_03.jpg b/example_images/room_03.jpg new file mode 100644 index 0000000000000000000000000000000000000000..5972aecfb5314813ebe6d9744fc7753d74fb9356 --- /dev/null +++ b/example_images/room_03.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0b3ea8c8cbdd30edb300af35535a307247f1b1b8a632afb37423cf256c331bae +size 127301 diff --git a/example_images/stair_01.jpg b/example_images/stair_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..bd44e983a15d81d27f3951f0044f65070210c6f7 --- /dev/null +++ b/example_images/stair_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7838f9ab682c71717612462fe256bda572ff556e0f2ab90aa64eb0f46495fa5e +size 15988 diff --git a/example_images/stand_01.jpg b/example_images/stand_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..4ec067ff021697ea65aa15ba2b846a042fa17a6a --- /dev/null +++ b/example_images/stand_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:39921cfa3a7a48a82c4ca442704dd654903b44a782e44e1bc42d6bf05c3d49af +size 227022 diff --git a/example_images/street_01.jpg b/example_images/street_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2e232edf0b23a3c41428f5f427214a445942cc2c --- /dev/null +++ b/example_images/street_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7d70e2bdeeb71218e1dd66a06f919a244c26178b97d09b593c6a5cc7f1c63569 +size 339088 diff --git a/example_images/street_02.png b/example_images/street_02.png new file mode 100644 index 0000000000000000000000000000000000000000..0f28aef89dd9d895a62df92997741b7fc85b43dd --- /dev/null +++ b/example_images/street_02.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6f00094173568cf0b9295735668cc7a7aee6d3cc75f61897e12ce0cff24922e6 +size 2930541 diff --git a/example_images/street_03.jpg b/example_images/street_03.jpg new file mode 100644 index 0000000000000000000000000000000000000000..86b5c74685e75b3534f6587232bda8325e5a2779 --- /dev/null +++ b/example_images/street_03.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e6029ec36beb9779ae99449ed6d0fe327e6d14ceac82c08aa53849e81cd5cee6 +size 54136 diff --git a/example_images/umic_building.jpg b/example_images/umic_building.jpg new file mode 100644 index 0000000000000000000000000000000000000000..4d0ff07538435c9ca4e6a700c9a750ad09f7a00b --- /dev/null +++ b/example_images/umic_building.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e375c42219bd9a4b63d710f52cb5c8dd247f6c05fb186f33bb7c042bf4f4106d +size 685564 diff --git a/example_images/volterra.jpg b/example_images/volterra.jpg new file mode 100644 index 0000000000000000000000000000000000000000..21933ccfada46ba342f9ec2ec6ca9e32b1a0192f --- /dev/null +++ b/example_images/volterra.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4e4aea6dad1c424deaa6c03a61ea58d4787f57cafb3b1178ec59061120d4478d +size 332406 diff --git a/pxdepth/__init__.py b/pxdepth/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..504c104cf86d233c054f0fb0f4bc1eb641ddfe9c --- /dev/null +++ b/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", +] diff --git a/pxdepth/build.py b/pxdepth/build.py new file mode 100644 index 0000000000000000000000000000000000000000..91330eb21863f37b4fa30ee41b01db45aae99fa0 --- /dev/null +++ b/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) diff --git a/pxdepth/config.py b/pxdepth/config.py new file mode 100644 index 0000000000000000000000000000000000000000..b285b514381c9e330704e127aed5c924d6dac94a --- /dev/null +++ b/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 diff --git a/pxdepth/evaluation/__init__.py b/pxdepth/evaluation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fae96c408cba3c6cf73659a6c801fbbd62afa48d --- /dev/null +++ b/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"] diff --git a/pxdepth/evaluation/dataloader.py b/pxdepth/evaluation/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..4b74e53cfce180ddabe10814f1e746f3cd20d037 --- /dev/null +++ b/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() diff --git a/pxdepth/evaluation/metrics.py b/pxdepth/evaluation/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..eaf55bb83168112b562cfb8164bb63b2cc74946a --- /dev/null +++ b/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 diff --git a/pxdepth/inference/__init__.py b/pxdepth/inference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..23c9c081d4a62ff38279e966ba088539ac8adea6 --- /dev/null +++ b/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", +] diff --git a/pxdepth/inference/resize.py b/pxdepth/inference/resize.py new file mode 100644 index 0000000000000000000000000000000000000000..4ce5319c687de36c315d1ea13054761e4e6aa443 --- /dev/null +++ b/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 diff --git a/pxdepth/inference/runner.py b/pxdepth/inference/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..5e6ba997a06f74d78ed7e686e1a94e4a170c73ca --- /dev/null +++ b/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 diff --git a/pxdepth/model/CM_PiT.py b/pxdepth/model/CM_PiT.py new file mode 100644 index 0000000000000000000000000000000000000000..5b4e82a57409ba3a63eef10291a58cecc3f7ec4e --- /dev/null +++ b/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, + ) diff --git a/pxdepth/model/Gated_Attention.py b/pxdepth/model/Gated_Attention.py new file mode 100644 index 0000000000000000000000000000000000000000..95ebefc677efa19467ab27794bae1460feca37c9 --- /dev/null +++ b/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)) diff --git a/pxdepth/model/Global_Context_Encoder.py b/pxdepth/model/Global_Context_Encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..3f2158f2179b641e406b0bb9f5f134039516fe6d --- /dev/null +++ b/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) diff --git a/pxdepth/model/PXDepth.py b/pxdepth/model/PXDepth.py new file mode 100644 index 0000000000000000000000000000000000000000..17b360359263b9cfc3e08754db28524764d72b72 --- /dev/null +++ b/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, + ) diff --git a/pxdepth/model/Pixel_Space_Depth_Predictor.py b/pxdepth/model/Pixel_Space_Depth_Predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..26d1a08aca49b824b98440b4b98b2dfe32add294 --- /dev/null +++ b/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 diff --git a/pxdepth/model/RoPE.py b/pxdepth/model/RoPE.py new file mode 100644 index 0000000000000000000000000000000000000000..47005d85b7ec19d6c5c58d469d52147395214936 --- /dev/null +++ b/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) diff --git a/pxdepth/model/__init__.py b/pxdepth/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..78fc70b35c0de6f5d63e34ae1183ba6731a94af4 --- /dev/null +++ b/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"] diff --git a/pxdepth/model/checkpoint.py b/pxdepth/model/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..d282401c1d955ba412702e83ff0f8820732d1e13 --- /dev/null +++ b/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 diff --git a/pxdepth/model/dinov2/__init__.py b/pxdepth/model/dinov2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ae847e46898077fe3d8701b8a181d7b4e3d41cd9 --- /dev/null +++ b/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" diff --git a/pxdepth/model/dinov2/hub/__init__.py b/pxdepth/model/dinov2/hub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b88da6bf80be92af00b72dfdb0a806fa64a7a2d9 --- /dev/null +++ b/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. diff --git a/pxdepth/model/dinov2/hub/backbones.py b/pxdepth/model/dinov2/hub/backbones.py new file mode 100644 index 0000000000000000000000000000000000000000..53fe83719d5107eb77a8f25ef1814c3d73446002 --- /dev/null +++ b/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, + ) diff --git a/pxdepth/model/dinov2/hub/utils.py b/pxdepth/model/dinov2/hub/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9c6641404093652d5a2f19b4cf283d976ec39e64 --- /dev/null +++ b/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 diff --git a/pxdepth/model/dinov2/layers/__init__.py b/pxdepth/model/dinov2/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..395df58d510e09f74748926631ec9bbdee3e18e6 --- /dev/null +++ b/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 diff --git a/pxdepth/model/dinov2/layers/attention.py b/pxdepth/model/dinov2/layers/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..83d53c67cef0a67038235a1dc19830cab8e174c0 --- /dev/null +++ b/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 diff --git a/pxdepth/model/dinov2/layers/block.py b/pxdepth/model/dinov2/layers/block.py new file mode 100644 index 0000000000000000000000000000000000000000..4d2687fe75fb9b75d923c924e188e63ee8518752 --- /dev/null +++ b/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 diff --git a/pxdepth/model/dinov2/layers/drop_path.py b/pxdepth/model/dinov2/layers/drop_path.py new file mode 100644 index 0000000000000000000000000000000000000000..1d640e0b969b8dcba96260243473700b4e5b24b5 --- /dev/null +++ b/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) diff --git a/pxdepth/model/dinov2/layers/layer_scale.py b/pxdepth/model/dinov2/layers/layer_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..51df0d7ce61f2b41fa9e6369f52391dd7fe7d386 --- /dev/null +++ b/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 diff --git a/pxdepth/model/dinov2/layers/mlp.py b/pxdepth/model/dinov2/layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..bbf9432aae9258612caeae910a7bde17999e328e --- /dev/null +++ b/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 diff --git a/pxdepth/model/dinov2/layers/patch_embed.py b/pxdepth/model/dinov2/layers/patch_embed.py new file mode 100644 index 0000000000000000000000000000000000000000..8b7c0804784a42cf80c0297d110dcc68cc85b339 --- /dev/null +++ b/pxdepth/model/dinov2/layers/patch_embed.py @@ -0,0 +1,88 @@ +# 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 + +from typing import Callable, Optional, Tuple, Union + +from torch import Tensor +import torch.nn as nn + + +def make_2tuple(x): + if isinstance(x, tuple): + assert len(x) == 2 + return x + + assert isinstance(x, int) + return (x, x) + + +class PatchEmbed(nn.Module): + """ + 2D image to patch embedding: (B,C,H,W) -> (B,N,D) + + Args: + img_size: Image size. + patch_size: Patch token size. + in_chans: Number of input image channels. + embed_dim: Number of linear projection output channels. + norm_layer: Normalization layer. + """ + + def __init__( + self, + img_size: Union[int, Tuple[int, int]] = 224, + patch_size: Union[int, Tuple[int, int]] = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Optional[Callable] = None, + flatten_embedding: bool = True, + ) -> None: + super().__init__() + + image_HW = make_2tuple(img_size) + patch_HW = make_2tuple(patch_size) + patch_grid_size = ( + image_HW[0] // patch_HW[0], + image_HW[1] // patch_HW[1], + ) + + self.img_size = image_HW + self.patch_size = patch_HW + self.patches_resolution = patch_grid_size + self.num_patches = patch_grid_size[0] * patch_grid_size[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.flatten_embedding = flatten_embedding + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + _, _, H, W = x.shape + patch_H, patch_W = self.patch_size + + assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}" + assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}" + + x = self.proj(x) # B C H W + H, W = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) # B HW C + x = self.norm(x) + if not self.flatten_embedding: + x = x.reshape(-1, H, W, self.embed_dim) # B H W C + return x + + def flops(self) -> float: + Ho, Wo = self.patches_resolution + flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + if self.norm is not None: + flops += Ho * Wo * self.embed_dim + return flops diff --git a/pxdepth/model/dinov2/layers/swiglu_ffn.py b/pxdepth/model/dinov2/layers/swiglu_ffn.py new file mode 100644 index 0000000000000000000000000000000000000000..5ce211515774d42e04c8b51003bae53b88f14b35 --- /dev/null +++ b/pxdepth/model/dinov2/layers/swiglu_ffn.py @@ -0,0 +1,72 @@ +# 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 os +from typing import Callable, Optional +import warnings + +from torch import Tensor, nn +import torch.nn.functional as F + + +class SwiGLUFFN(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + 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.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.chunk(2, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None +try: + if XFORMERS_ENABLED: + from xformers.ops import SwiGLU + + XFORMERS_AVAILABLE = True + # warnings.warn("xFormers is available (SwiGLU)") + else: + # warnings.warn("xFormers is disabled (SwiGLU)") + raise ImportError +except ImportError: + SwiGLU = SwiGLUFFN + XFORMERS_AVAILABLE = False + + # warnings.warn("xFormers is not available (SwiGLU)") + + +class SwiGLUFFNFused(SwiGLU): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + out_features = out_features or in_features + hidden_features = hidden_features or in_features + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_features=out_features, + bias=bias, + ) diff --git a/pxdepth/model/dinov2/models/__init__.py b/pxdepth/model/dinov2/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..af9e96cee81846e08c09a3a0f55812f2f2764623 --- /dev/null +++ b/pxdepth/model/dinov2/models/__init__.py @@ -0,0 +1,8 @@ +# 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 . import vision_transformer + +__all__ = ["vision_transformer"] diff --git a/pxdepth/model/dinov2/models/vision_transformer.py b/pxdepth/model/dinov2/models/vision_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..f0bed9d0b7cdcff2b5e129121251c58e41c4c61d --- /dev/null +++ b/pxdepth/model/dinov2/models/vision_transformer.py @@ -0,0 +1,407 @@ +# 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/main/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +from functools import partial +import math +import logging +from typing import Sequence, Tuple, Union, Callable, Optional, List + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from torch.nn.init import trunc_normal_ + +from ..layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block + + +logger = logging.getLogger("dinov2") + + +def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +class BlockChunk(nn.ModuleList): + def forward(self, x): + for b in self: + x = b(x) + return x + + +class DinoVisionTransformer(nn.Module): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=True, + ffn_bias=True, + proj_bias=True, + drop_path_rate=0.0, + drop_path_uniform=False, + init_values=None, # for layerscale: None or 0 => no layerscale + embed_layer=PatchEmbed, + act_layer=nn.GELU, + block_fn=Block, + ffn_layer="mlp", + block_chunks=1, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + ): + """ + Args: + img_size (int, tuple): input image size + patch_size (int, tuple): patch size + in_chans (int): number of input channels + embed_dim (int): embedding dimension + depth (int): depth of transformer + num_heads (int): number of attention heads + mlp_ratio (int): ratio of mlp hidden dim to embedding dim + qkv_bias (bool): enable bias for qkv if True + proj_bias (bool): enable bias for proj in attn if True + ffn_bias (bool): enable bias for ffn if True + drop_path_rate (float): stochastic depth rate + drop_path_uniform (bool): apply uniform drop rate across blocks + weight_init (str): weight init scheme + init_values (float): layer-scale init values + embed_layer (nn.Module): patch embedding layer + act_layer (nn.Module): MLP activation layer + block_fn (nn.Module): transformer block class + ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity" + block_chunks: (int) split block sequence into block_chunks units for FSDP wrap + num_register_tokens: (int) number of extra cls tokens (so-called "registers") + interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings + interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings + """ + super().__init__() + norm_layer = partial(nn.LayerNorm, eps=1e-6) + + self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models + self.num_tokens = 1 + self.n_blocks = depth + self.num_heads = num_heads + self.patch_size = patch_size + self.num_register_tokens = num_register_tokens + self.interpolate_antialias = interpolate_antialias + self.interpolate_offset = interpolate_offset + + self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) + num_patches = self.patch_embed.num_patches + + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) + assert num_register_tokens >= 0 + self.register_tokens = ( + nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None + ) + + if drop_path_uniform is True: + dpr = [drop_path_rate] * depth + else: + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule + + if ffn_layer == "mlp": + logger.info("using MLP layer as FFN") + ffn_layer = Mlp + elif ffn_layer == "swiglufused" or ffn_layer == "swiglu": + logger.info("using SwiGLU layer as FFN") + ffn_layer = SwiGLUFFNFused + elif ffn_layer == "identity": + logger.info("using Identity layer as FFN") + + def f(*args, **kwargs): + return nn.Identity() + + ffn_layer = f + else: + raise NotImplementedError + + blocks_list = [ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer, + ffn_layer=ffn_layer, + init_values=init_values, + ) + for i in range(depth) + ] + if block_chunks > 0: + self.chunked_blocks = True + chunked_blocks = [] + chunksize = depth // block_chunks + for i in range(0, depth, chunksize): + # this is to keep the block index consistent if we chunk the block list + chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize]) + self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks]) + else: + self.chunked_blocks = False + self.blocks = nn.ModuleList(blocks_list) + + self.norm = norm_layer(embed_dim) + self.head = nn.Identity() + + self.mask_token = nn.Parameter(torch.zeros(1, embed_dim)) + + self.init_weights() + + @property + def onnx_compatible_mode(self): + return getattr(self, "_onnx_compatible_mode", False) + + @onnx_compatible_mode.setter + def onnx_compatible_mode(self, value: bool): + self._onnx_compatible_mode = value + + def init_weights(self): + trunc_normal_(self.pos_embed, std=0.02) + nn.init.normal_(self.cls_token, std=1e-6) + if self.register_tokens is not None: + nn.init.normal_(self.register_tokens, std=1e-6) + named_apply(init_weights_vit_timm, self) + + def interpolate_pos_encoding(self, x, h, w): + previous_dtype = x.dtype + npatch = x.shape[1] - 1 + batch_size = x.shape[0] + N = self.pos_embed.shape[1] - 1 + if not self.onnx_compatible_mode and npatch == N and w == h: + return self.pos_embed + pos_embed = self.pos_embed.float() + class_pos_embed = pos_embed[:, 0, :] + patch_pos_embed = pos_embed[:, 1:, :] + dim = x.shape[-1] + h0, w0 = h // self.patch_size, w // self.patch_size + M = int(math.sqrt(N)) # Recover the number of patches in each dimension + assert N == M * M + kwargs = {} + if not self.onnx_compatible_mode and self.interpolate_offset > 0: + # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8 + # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors + sx = float(w0 + self.interpolate_offset) / M + sy = float(h0 + self.interpolate_offset) / M + kwargs["scale_factor"] = (sy, sx) + else: + # Simply specify an output size instead of a scale factor + kwargs["size"] = (h0, w0) + + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2), + mode="bicubic", + antialias=self.interpolate_antialias, + **kwargs, + ) + + assert (h0, w0) == patch_pos_embed.shape[-2:] + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).flatten(1, 2) + return torch.cat((class_pos_embed[:, None, :].expand(patch_pos_embed.shape[0], -1, -1), patch_pos_embed), dim=1).to(previous_dtype) + + def prepare_tokens_with_masks(self, x, masks=None): + B, nc, h, w = x.shape + x = self.patch_embed(x) + + if masks is not None: + x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x) + + x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1) + x = x + self.interpolate_pos_encoding(x, h, w) + + if self.register_tokens is not None: + x = torch.cat( + ( + x[:, :1], + self.register_tokens.expand(x.shape[0], -1, -1), + x[:, 1:], + ), + dim=1, + ) + + return x + + def forward_features_list(self, x_list, masks_list): + x = [self.prepare_tokens_with_masks(x, masks) for x, masks, ar in zip(x_list, masks_list)] + for blk in self.blocks: + x = blk(x) + + all_x = x + output = [] + for x, masks in zip(all_x, masks_list): + x_norm = self.norm(x) + output.append( + { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], + "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :], + "x_prenorm": x, + "masks": masks, + } + ) + return output + + def forward_features(self, x, masks=None): + if isinstance(x, list): + return self.forward_features_list(x, masks) + + x = self.prepare_tokens_with_masks(x, masks) + + for blk in self.blocks: + x = blk(x) + + x_norm = self.norm(x) + return { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], + "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :], + "x_prenorm": x, + "masks": masks, + } + + def _get_intermediate_layers_not_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + # If n is an int, take the n last blocks. If it's a list, take them + output, total_block_len = [], len(self.blocks) + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + for i, blk in enumerate(self.blocks): + x = blk(x) + if i in blocks_to_take: + output.append(x) + assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def _get_intermediate_layers_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + output, i, total_block_len = [], 0, len(self.blocks[-1]) + # If n is an int, take the n last blocks. If it's a list, take them + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + for block_chunk in self.blocks: + for blk in block_chunk[i:]: # Passing the nn.Identity() + x = blk(x) + if i in blocks_to_take: + output.append(x) + i += 1 + assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def get_intermediate_layers( + self, + x: torch.Tensor, + n: Union[int, Sequence] = 1, # Layers or n last layers to take + reshape: bool = False, + return_class_token: bool = False, + norm=True, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]: + if self.chunked_blocks: + outputs = self._get_intermediate_layers_chunked(x, n) + else: + outputs = self._get_intermediate_layers_not_chunked(x, n) + if norm: + outputs = [self.norm(out) for out in outputs] + class_tokens = [out[:, 0] for out in outputs] + outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs] + if reshape: + B, _, w, h = x.shape + outputs = [ + out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous() + for out in outputs + ] + if return_class_token: + return tuple(zip(outputs, class_tokens)) + return tuple(outputs) + + def forward(self, *args, is_training=False, **kwargs): + ret = self.forward_features(*args, **kwargs) + if is_training: + return ret + else: + return self.head(ret["x_norm_clstoken"]) + + +def init_weights_vit_timm(module: nn.Module, name: str = ""): + """ViT weight initialization, original timm impl (for reproducibility)""" + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + + +def vit_small(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=384, + depth=12, + num_heads=6, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_base(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_large(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1024, + depth=24, + num_heads=16, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs): + """ + Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64 + """ + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1536, + depth=40, + num_heads=24, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model diff --git a/pxdepth/model/inference.py b/pxdepth/model/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..1ecc3ecd6497b352fcd41484723101dd4b63fbd6 --- /dev/null +++ b/pxdepth/model/inference.py @@ -0,0 +1,229 @@ +"""Metric-scale inference helpers used by :class:`PXDepth`. + +The network predicts normalized log-depth. This module keeps reference-model +loading, low-resolution log-space alignment, camera reconstruction, and mask +application outside the architecture file while preserving the released +``model.infer`` behavior. +""" + +from numbers import Number +from typing import Dict, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import utils3d + +from ..utils.alignment import align_depth_affine +from .precision import full_precision + + +def _reference_model(model: nn.Module) -> nn.Module: + """Load and cache the optional MoGe-2 reference model. + + Args: + model: PXDepth-like module exposing ``device``, ``dtype``, and a + mutable ``_reference_model`` attribute. + + Returns: + An evaluation-mode MoGe-2 module placed on the same device and storage + dtype as ``model``. + + Raises: + RuntimeError: If the optional MoGe-2 dependency is unavailable. + """ + if model._reference_model is None: + try: + from moge.model.v2 import MoGeModel + except ImportError as exc: + raise RuntimeError( + "Metric-scale visualization requires MoGe-2. Install the optional `reference` dependencies " + "or pass gt_depth and intrinsics to infer()." + ) from exc + reference = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal") + model._reference_model = reference.to(device=model.device, dtype=model.dtype).eval() + return model._reference_model + + +def _patch_size(model: nn.Module) -> Optional[int]: + """Resolve a reference model's scalar image patch size. + + Args: + model: Reference model potentially exposing ``patch_size`` directly or + through ``encoder.backbone``. + + Returns: + A positive integer patch size, or ``None`` when it cannot be resolved. + """ + patch = getattr(model, "patch_size", None) + if patch is None: + patch = getattr(getattr(getattr(model, "encoder", None), "backbone", None), "patch_size", None) + if isinstance(patch, (tuple, list)): + patch = patch[0] + return int(patch) if isinstance(patch, Number) and int(patch) > 0 else None + + +def _prepare_reference_image(image: torch.Tensor, model: nn.Module) -> Tuple[torch.Tensor, Tuple[int, int]]: + """Resize an RGB batch to satisfy a reference model's patch constraint. + + Args: + image: RGB tensor ``[B, 3, H, W]``. + model: Reference depth model inspected for its patch size. + + Returns: + A pair containing the bilinearly resized RGB tensor and original + ``(H, W)``. The input is returned unchanged when already divisible. + """ + patch = _patch_size(model) + size = tuple(image.shape[-2:]) + if patch is None: + return image, size + height, width = size + target = ( + height if height % patch == 0 else max(patch, height // patch * patch), + width if width % patch == 0 else max(patch, width // patch * patch), + ) + if target == size: + return image, size + return F.interpolate(image, target, mode="bilinear", align_corners=False), size + + +@torch.inference_mode() +def infer( + model: nn.Module, + image: torch.Tensor, + gt_depth: Optional[torch.Tensor] = None, + intrinsics: Optional[torch.Tensor] = None, + fov_x: Optional[Union[Number, 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 aligned depth, validity, camera 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 is + estimated from a masked-nearest 64x64 representation, matching the released + evaluation and visualization behavior. + + Args: + model: PXDepth-like module exposing ``forward``, ``device``, + ``dtype``, and ``mask_threshold``. + image: RGB tensor ``[3,H,W]`` or ``[B,3,H,W]`` in ``[0,1]``. + gt_depth: Optional reference depth ``[H,W]`` or ``[B,H,W]``. Finite + positive values define log-space alignment. + intrinsics: Optional normalized camera matrix ``[3,3]`` or batch + ``[B,3,3]`` corresponding to the reference depth. + fov_x: Optional horizontal field of view in degrees, scalar or ``[B]``. + ref_image: Optional original-resolution RGB input used only by MoGe-2. + apply_mask: Replace invalid predicted depth and points with infinity. + use_fp16: Use FP16 in attention-heavy model regions. + use_fp32: Force full precision and disable reduced-precision autocast. + + Returns: + Dictionary containing ``depth`` ``[B,H,W]``, boolean ``mask`` + ``[B,H,W]``, ``points`` ``[B,H,W,3]``, normalized ``intrinsics`` + ``[B,3,3]``, and horizontal ``fov_x`` ``[B]``. The leading batch + dimension is removed when ``image`` is unbatched. + """ + squeeze = image.ndim == 3 + if squeeze: + image = image.unsqueeze(0) + image = image.to(device=model.device, dtype=model.dtype) + if ref_image is not None and ref_image.ndim == 3: + ref_image = ref_image.unsqueeze(0) + if ref_image is not None: + ref_image = ref_image.to(device=model.device, dtype=model.dtype) + if gt_depth is not None and gt_depth.ndim == 2: + gt_depth = gt_depth.unsqueeze(0) + if gt_depth is not None: + gt_depth = gt_depth.to(device=model.device, dtype=torch.float32) + if intrinsics is not None and intrinsics.ndim == 2: + intrinsics = intrinsics.unsqueeze(0) + if intrinsics is not None: + intrinsics = intrinsics.to(device=model.device, dtype=torch.float32) + + height, width = image.shape[-2:] + aspect = width / height + output = model.forward(image, use_fp16=use_fp16, use_fp32=use_fp32) + + with full_precision(model.device): + pred = output["depth"].float() + mask = output["mask"].float() + ref_depth, ref_intrinsics, ref_fov = gt_depth, intrinsics, fov_x + if ref_depth is None: + reference = _reference_model(model) + reference_input = image if ref_image is None else ref_image + reference_input, reference_size = _prepare_reference_image(reference_input, reference) + ref = reference.infer(reference_input, apply_mask=True, use_fp16=use_fp16 and not use_fp32) + ref_depth = ref["depth"].float() + if ref_depth.ndim == 2: + ref_depth = ref_depth.unsqueeze(0) + if ref_depth.shape[-2:] != reference_size: + ref_depth = F.interpolate(ref_depth.unsqueeze(1), reference_size, mode="nearest").squeeze(1) + ref_intrinsics = ref.get("intrinsics") + ref_fov = ref.get("fov_x") + if ref_intrinsics is not None: + ref_intrinsics = ref_intrinsics.float() + if ref_fov is not None: + ref_fov = ref_fov.float() + if ref_depth.shape[-2:] != pred.shape[-2:]: + ref_depth = F.interpolate(ref_depth.unsqueeze(1), pred.shape[-2:], mode="nearest").squeeze(1) + + ref_valid = torch.isfinite(ref_depth) & (ref_depth > 0) + ref_log = torch.where(ref_valid, torch.log1p(ref_depth), 0.0) + scale = torch.ones(pred.shape[0], device=pred.device, dtype=pred.dtype) + shift = torch.zeros_like(scale) + valid = torch.isfinite(pred) & ref_valid + for index in range(pred.shape[0]): + low_mask, nearest = utils3d.pt.masked_nearest_resize( + mask=valid[index], size=(64, 64), return_index=True + ) + if not low_mask.any(): + continue + pred_low = pred[index][nearest][low_mask] + ref_log_low = ref_log[index][nearest][low_mask] + ref_depth_low = ref_depth[index][nearest][low_mask] + a, b = align_depth_affine( + pred_low.unsqueeze(0), + ref_log_low.unsqueeze(0), + (1.0 / ref_depth_low.clamp_min(1e-5)).unsqueeze(0), + ) + scale[index], shift[index] = a.squeeze(0), b.squeeze(0) + depth = torch.expm1(scale[:, None, None] * pred + shift[:, None, None]) + + if ref_intrinsics is None: + if ref_fov is None: + fx = torch.ones(depth.shape[0], device=depth.device) + fy = torch.ones_like(fx) + ref_fov = 2.0 * torch.atan(0.5 / fx).rad2deg() + else: + ref_fov = torch.as_tensor(ref_fov, device=depth.device, dtype=depth.dtype) + focal = aspect / (1.0 + aspect**2) ** 0.5 / torch.tan(torch.deg2rad(ref_fov / 2.0)) + if focal.ndim == 0: + focal = focal[None].expand(depth.shape[0]) + fx = focal / 2.0 * (1.0 + aspect**2) ** 0.5 / aspect + fy = focal / 2.0 * (1.0 + aspect**2) ** 0.5 + ref_intrinsics = utils3d.pt.intrinsics_from_focal_center( + fx, + fy, + torch.tensor(0.5, device=depth.device), + torch.tensor(0.5, device=depth.device), + ) + else: + ref_fov = 2.0 * torch.atan(0.5 / ref_intrinsics[..., 0, 0]).rad2deg() + + mask_binary = (mask > model.mask_threshold) & torch.isfinite(depth) & (depth > 0) + points = utils3d.pt.depth_map_to_point_map(depth, intrinsics=ref_intrinsics) + if apply_mask: + depth = torch.where(mask_binary, depth, torch.inf) + points = torch.where(mask_binary[..., None], points, torch.inf) + result = { + "depth": depth, + "mask": mask_binary, + "points": points, + "intrinsics": ref_intrinsics, + "fov_x": ref_fov, + } + return {key: value.squeeze(0) for key, value in result.items()} if squeeze else result diff --git a/pxdepth/model/precision.py b/pxdepth/model/precision.py new file mode 100644 index 0000000000000000000000000000000000000000..f37fe229ad235872a06703c5e4e40e19f2e74778 --- /dev/null +++ b/pxdepth/model/precision.py @@ -0,0 +1,64 @@ +"""Precision-policy helpers shared by PXDepth model components. + +The context managers define where full precision is required and where +attention-heavy regions may use FP16 or BF16 autocast. Centralizing this policy +keeps numerical behavior consistent between evaluation and public inference +entry points. +""" + +from contextlib import nullcontext +from typing import Optional + +import torch + + +def reduced_precision(device: torch.device, dtype: Optional[torch.dtype] = torch.bfloat16): + """Create an autocast context for attention-heavy model regions. + + Args: + device: Device on which enclosed tensor operations execute. + dtype: CUDA autocast dtype, normally BF16 or FP16. ``None`` requests + full precision. + + Returns: + Context manager enabling CUDA autocast when applicable, otherwise a + no-op context manager. + """ + if device.type == "cuda" and dtype is not None: + return torch.autocast(device_type="cuda", dtype=dtype, enabled=True) + return nullcontext() + + +def full_precision(device: torch.device): + """Create a context that disables an enclosing autocast region. + + Args: + device: Device type used to construct the autocast context. + + Returns: + Context manager that executes enclosed operators in their explicit + dtypes, or a no-op context on unsupported devices. + """ + if device.type in {"cuda", "cpu"}: + return torch.autocast(device_type=device.type, enabled=False) + return nullcontext() + + +def inference_dtype(use_fp16: bool = False, use_fp32: bool = False) -> Optional[torch.dtype]: + """Resolve public inference precision flags to an autocast dtype. + + Args: + use_fp16: Select FP16 attention and encoder execution. + use_fp32: Disable reduced precision. Mutually exclusive with FP16. + + Returns: + ``torch.float16`` for FP16, ``None`` for FP32, and + ``torch.bfloat16`` for the default path. + """ + if use_fp16 and use_fp32: + raise ValueError("use_fp16 and use_fp32 are mutually exclusive") + if use_fp32: + return None + if use_fp16: + return torch.float16 + return torch.bfloat16 diff --git a/pxdepth/model/utils.py b/pxdepth/model/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3cc000141e32472717c7545f471ea5d8cb31fcb9 --- /dev/null +++ b/pxdepth/model/utils.py @@ -0,0 +1,137 @@ +"""Runtime wrappers for DINOv2 checkpointing and optimized attention. + +The vendored DINOv2 source is kept close to upstream. These two helpers apply +PXDepth-specific runtime behavior without editing every upstream block. +""" + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +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 + + +def wrap_module_with_gradient_checkpointing(module: nn.Module) -> nn.Module: + """Recompute a module's forward pass during backward to save memory. + + A dynamic subclass replaces only the supplied module instance's class. Its + original ``forward`` remains the computation being checkpointed, so + parameters and state-dict names are unchanged. + + Args: + module: DINO transformer block to modify in place. Its forward inputs + and output retain the upstream block contract, normally token + tensors ``[B, N, C]`` plus optional attention metadata. + + Returns: + The same module instance with a checkpointed ``forward`` method. + """ + from torch.utils.checkpoint import checkpoint + + class _CheckpointingWrapper(module.__class__): + """Per-instance subclass that checkpoints the inherited forward call. + + It introduces no parameters or buffers, preserving the wrapped DINO + block's state-dict schema and tensor interface. + + Instances are created by mutating the class of one existing block. + """ + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Run the inherited module under non-reentrant checkpointing. + + Args: + *args: Positional arguments accepted by the wrapped block, + typically a token tensor ``[B,N,C]``. + **kwargs: Keyword arguments accepted by the wrapped block. + + Returns: + The wrapped block's original output, normally ``[B,N,C]``. + """ + return checkpoint(super().forward, *args, use_reentrant=False, **kwargs) + + module.__class__ = _CheckpointingWrapper + return module + + +def wrap_dinov2_attention_with_sdpa(module: nn.Module) -> nn.Module: + """Replace one DINOv2 attention forward path with Flash-Attention or SDPA. + + Flash-Attention is selected only for CUDA FP16/BF16 tensors without an + additive attention bias. Other inputs use PyTorch scaled dot-product + attention. The wrapper reuses the upstream QKV and output projections, so + checkpoint keys and numerical semantics remain compatible. + + Args: + module: DINOv2 attention module with ``qkv``, ``num_heads``, ``scale``, + ``proj``, and ``proj_drop`` attributes. + + Returns: + The same attention module instance with an optimized forward method. + """ + if torch.__version__ < "2.0": + raise RuntimeError("SDPA requires PyTorch 2.0 or later") + + class _AttentionWrapper(module.__class__): + """Per-instance attention subclass dispatching to optimized kernels. + + It reuses all upstream projections and introduces no new checkpoint + state, so wrapping does not affect serialization compatibility. + + Instances are created by mutating one existing DINO attention module. + """ + + def forward(self, x: torch.Tensor, attn_bias: Any = None) -> torch.Tensor: + """Apply multi-head self-attention to DINO patch and special tokens. + + Args: + x: Input token tensor ``[B, N, C]``. + attn_bias: Optional bias broadcastable to ``[B,H,N,N]``. A + non-``None`` value disables the external Flash-Attention path. + + Returns: + Projected token tensor ``[B, N, C]``. + """ + batch, length, channels = x.shape + qkv = self.qkv(x).reshape( + batch, + length, + 3, + self.num_heads, + channels // self.num_heads, + ).permute(2, 0, 3, 1, 4) + query, key, value = torch.unbind(qkv, 0) + + use_flash = ( + FLASH_ATTN_AVAILABLE + and attn_bias is None + and query.is_cuda + and query.dtype in (torch.float16, torch.bfloat16) + ) + if use_flash: + query = query.permute(0, 2, 1, 3).contiguous() + key = key.permute(0, 2, 1, 3).contiguous() + value = value.permute(0, 2, 1, 3).contiguous() + out = flash_attn_func( + query, + key, + value, + dropout_p=0.0, + softmax_scale=self.scale, + causal=False, + ).reshape(batch, length, channels) + else: + out = F.scaled_dot_product_attention(query, key, value, attn_bias) + out = out.permute(0, 2, 1, 3).reshape(batch, length, channels) + return self.proj_drop(self.proj(out)) + + module.__class__ = _AttentionWrapper + return module diff --git a/pxdepth/registry.py b/pxdepth/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..b0085dea06028c2ddb23e8bd0f6a1e4037be0933 --- /dev/null +++ b/pxdepth/registry.py @@ -0,0 +1,126 @@ +"""Small component registries used by PXDepth configuration builders. + +The release intentionally keeps this mechanism minimal. A registry maps a +short string from a JSON config to a Python class or function. Third-party +projects can register their own implementation from a module listed in the +config's optional ``imports`` field, without modifying PXDepth source. +""" + +from importlib import import_module +from typing import Any, Callable, Dict, Optional, TypeVar + + +T = TypeVar("T") + + +class Registry: + """Map configurable names to callables and instantiate them from dictionaries. + + Args: + name: Human-readable component category used in error messages. + + The registry accepts both short registered names and dotted Python paths. + A dotted path such as ``my_package.models.CustomEncoder`` is imported lazily, + which makes small out-of-tree experiments possible without editing this repo. + """ + + def __init__(self, name: str) -> None: + """Create an empty registry for one component category. + + Args: + name: Human-readable category used in validation error messages. + + Returns: + ``None``. Registered items are stored in a new private mapping. + """ + self.name = name + self._items: Dict[str, Callable[..., Any]] = {} + + def register( + self, + value: Optional[T] = None, + name: Optional[str] = None, + ) -> Callable[[T], T] | T: + """Register a class or function, directly or as a decorator. + + Args: + value: Callable to register. Omit it when using decorator syntax. + name: Optional config name. The callable's ``__name__`` is used when + no explicit name is supplied. + + Returns: + The original callable, allowing ``@REGISTRY.register()`` usage. + """ + + def add(item: T) -> T: + """Insert one callable and return it unchanged. + + Args: + item: Class or function supplied directly or by decorator use. + + Returns: + The original object, preserving normal decorator semantics. + """ + key = name or getattr(item, "__name__", None) + if not key: + raise ValueError(f"A {self.name} registration needs an explicit name.") + if key in self._items and self._items[key] is not item: + raise KeyError(f"{self.name} '{key}' is already registered.") + self._items[key] = item # type: ignore[assignment] + return item + + return add if value is None else add(value) + + def get(self, name: str) -> Callable[..., Any]: + """Resolve a registered name or dotted import path to a callable. + + Args: + name: Registered short name or ``package.module.callable`` path. + + Returns: + Resolved class or function. + """ + if name in self._items: + return self._items[name] + if "." in name: + module_name, attribute = name.rsplit(".", 1) + value = getattr(import_module(module_name), attribute) + if not callable(value): + raise TypeError(f"Resolved {self.name} '{name}' is not callable.") + return value + available = ", ".join(sorted(self._items)) or "none" + raise KeyError(f"Unknown {self.name} '{name}'. Available: {available}") + + def build(self, config: Dict[str, Any], **defaults: Any) -> Any: + """Instantiate one component from a ``type`` plus constructor arguments. + + Args: + config: Dictionary containing a required ``type`` key. Remaining + entries are passed to the resolved callable as keyword arguments. + **defaults: Values used only when the config does not define the key. + + Returns: + Constructed component instance. + """ + if not isinstance(config, dict): + raise TypeError(f"{self.name} config must be a dictionary.") + params = dict(defaults) + params.update(config) + type_name = params.pop("type", None) + if not isinstance(type_name, str) or not type_name: + raise ValueError(f"{self.name} config requires a non-empty 'type'.") + return self.get(type_name)(**params) + + def names(self) -> tuple[str, ...]: + """List the short names currently registered in this category. + + Returns: + Lexicographically sorted tuple of names. Dotted paths are resolved + lazily and therefore do not appear unless explicitly registered. + """ + return tuple(sorted(self._items)) + + +MODELS = Registry("model") +ENCODERS = Registry("encoder") +PREDICTORS = Registry("predictor") diff --git a/pxdepth/utils/__init__.py b/pxdepth/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..35cd76f058f9c043784f5ce6f2f8d8cce090ad5d --- /dev/null +++ b/pxdepth/utils/__init__.py @@ -0,0 +1,6 @@ +"""Shared low-level utilities for PXDepth inference and evaluation. + +Submodules provide camera geometry, robust prediction alignment, +processed-benchmark readers, point-cloud export, and visualization. No broad +wildcard API is re-exported from this package. +""" diff --git a/pxdepth/utils/alignment.py b/pxdepth/utils/alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..fb559fefc24882bc3e514611fde69246452ac7f6 --- /dev/null +++ b/pxdepth/utils/alignment.py @@ -0,0 +1,324 @@ +"""Robust scalar, affine-depth, and point-cloud alignment primitives. + +These routines estimate scale and shift parameters under missing data and +outliers, using low-resolution robust solvers where required by the evaluation +protocol. Returned parameters are batched tensors that callers apply to +full-resolution depth maps or point clouds. +""" + +import math +from typing import Callable, Optional, Tuple, Union + +import torch + + +def scatter_min(size: int, dim: int, index: torch.LongTensor, src: torch.Tensor) -> torch.return_types.min: + """Scatter-reduce values by minimum and recover source indices. + + Args: + size: Length of the reduced output dimension. + dim: Dimension of ``src`` along which group indices apply. + index: Long tensor broadcast-compatible with ``src`` assigning each + source value to an output group. + src: Floating source tensor of arbitrary shape. + + Returns: + ``torch.return_types.min`` containing grouped minimum values and source + indices, each shaped like ``src`` with dimension ``dim`` replaced by + ``size``. + """ + shape = src.shape[:dim] + (size,) + src.shape[dim + 1:] + minimum = torch.full(shape, float('inf'), dtype=src.dtype, device=src.device).scatter_reduce(dim=dim, index=index, src=src, reduce='amin', include_self=False) + minimum_where = torch.where(src == torch.gather(minimum, dim=dim, index=index)) + indices = torch.full(shape, -1, dtype=torch.long, device=src.device) + indices[(*minimum_where[:dim], index[minimum_where], *minimum_where[dim + 1:])] = minimum_where[dim] + return torch.return_types.min((minimum, indices)) + + +def split_batch_fwd(fn: Callable, chunk_size: int, *args, **kwargs): + """Evaluate a tensor function in chunks along its leading batch axis. + + Args: + fn: Callable accepting the supplied positional and keyword arguments. + chunk_size: Maximum leading-axis tensor length per invocation. + *args: Tensor arguments sharing leading batch length, or constants + repeated for every chunk. + **kwargs: Keyword equivalents of ``args``. + + Returns: + Concatenated tensor result, or tuple of concatenated tensors when ``fn`` + returns a tuple. + """ + batch_size = next(x for x in (*args, *kwargs.values()) if isinstance(x, torch.Tensor)).shape[0] + n_chunks = batch_size // chunk_size + (batch_size % chunk_size > 0) + splited_args = tuple(arg.split(chunk_size, dim=0) if isinstance(arg, torch.Tensor) else [arg] * n_chunks for arg in args) + splited_kwargs = {k: [v.split(chunk_size, dim=0) if isinstance(v, torch.Tensor) else [v] * n_chunks] for k, v in kwargs.items()} + results = [] + for i in range(n_chunks): + chunk_args = tuple(arg[i] for arg in splited_args) + chunk_kwargs = {k: v[i] for k, v in splited_kwargs.items()} + results.append(fn(*chunk_args, **chunk_kwargs)) + + if isinstance(results[0], tuple): + return tuple(torch.cat(r, dim=0) for r in zip(*results)) + else: + return torch.cat(results, dim=0) + + +def _pad_inf(x_: torch.Tensor): + """Pad a sorted sequence with negative and positive infinity. + + Args: + x_: Tensor ``[...,N]``. + + Returns: + Tensor ``[...,N+2]`` with sentinels at both ends. + """ + return torch.cat([torch.full_like(x_[..., :1], -torch.inf), x_, torch.full_like(x_[..., :1], torch.inf)], dim=-1) + + +def _pad_cumsum(cumsum: torch.Tensor): + """Pad a cumulative sum with zero and its final total. + + Args: + cumsum: Cumulative values ``[...,N]``. + + Returns: + Tensor ``[...,N+2]`` equal to ``[0,cumsum,total]``. + """ + return torch.cat([torch.zeros_like(cumsum[..., :1]), cumsum, cumsum[..., -1:]], dim=-1) + + +def _compute_residual(a: torch.Tensor, xyw: torch.Tensor, trunc: float): + """Evaluate a truncated weighted absolute residual for candidate scales. + + Args: + a: Candidate scales ``[K,1]``. + xyw: Stacked source, target, and weight values ``[K,N,3]``. + trunc: Scalar upper bound applied to every weighted residual. + + Returns: + Objective values ``[K]``. + """ + return a.mul(xyw[..., 0]).sub_(xyw[..., 1]).abs_().mul_(xyw[..., 2]).clamp_max_(trunc).sum(dim=-1) + + +def align(x: torch.Tensor, y: torch.Tensor, w: torch.Tensor, trunc: Optional[Union[float, torch.Tensor]] = None, eps: float = 1e-7) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: + """Solve robust weighted L1 scale alignment without iterative optimization. + + Args: + x: Source values ``[...,N]``. + y: Target values ``[...,N]``. + w: Nonnegative correspondence weights ``[...,N]``. + trunc: Optional scalar/tensor cap for each weighted absolute residual. + eps: Lower bound protecting divisions by near-zero source values. + + Returns: + scale: Differentiable optimal scale tensor ``[...]``. + loss: Detached objective value tensor ``[...]``. + index: Long tensor ``[...]`` identifying the correspondence whose ratio + reproduces each selected optimum. + """ + if trunc is None: + x, y, w = torch.broadcast_tensors(x, y, w) + sign = torch.sign(x) + x, y = x * sign, y * sign + y_div_x = y / x.clamp_min(eps) + y_div_x, argsort = y_div_x.sort(dim=-1) + + wx = torch.gather(x * w, dim=-1, index=argsort) + derivatives = 2 * wx.cumsum(dim=-1) - wx.sum(dim=-1, keepdim=True) + search = torch.searchsorted(derivatives, torch.zeros_like(derivatives[..., :1]), side='left').clamp_max(derivatives.shape[-1] - 1) + + a = y_div_x.gather(dim=-1, index=search).squeeze(-1) + index = argsort.gather(dim=-1, index=search).squeeze(-1) + loss = (w * (a[..., None] * x - y).abs()).sum(dim=-1) + + else: + # Reshape to (batch_size, n) for simplicity + x, y, w = torch.broadcast_tensors(x, y, w) + batch_shape = x.shape[:-1] + batch_size = math.prod(batch_shape) + x, y, w = x.reshape(-1, x.shape[-1]), y.reshape(-1, y.shape[-1]), w.reshape(-1, w.shape[-1]) + + sign = torch.sign(x) + x, y = x * sign, y * sign + wx, wy = w * x, w * y + xyw = torch.stack([x, y, w], dim=-1) # Stacked for convenient gathering + + y_div_x = A = y / x.clamp_min(eps) + B = (wy - trunc) / wx.clamp_min(eps) + C = (wy + trunc) / wx.clamp_min(eps) + with torch.no_grad(): + # Caculate prefix sum by orders of A, B, C + A, A_argsort = A.sort(dim=-1) + Q_A = torch.cumsum(torch.gather(wx, dim=-1, index=A_argsort), dim=-1) + A, Q_A = _pad_inf(A), _pad_cumsum(Q_A) # Pad [-inf, A1, ..., An, inf] and [0, Q1, ..., Qn, Qn] to handle edge cases. + + B, B_argsort = B.sort(dim=-1) + Q_B = torch.cumsum(torch.gather(wx, dim=-1, index=B_argsort), dim=-1) + B, Q_B = _pad_inf(B), _pad_cumsum(Q_B) + + C, C_argsort = C.sort(dim=-1) + Q_C = torch.cumsum(torch.gather(wx, dim=-1, index=C_argsort), dim=-1) + C, Q_C = _pad_inf(C), _pad_cumsum(Q_C) + + # Caculate left and right derivative of A + j_A = torch.searchsorted(A, y_div_x, side='left').sub_(1) + j_B = torch.searchsorted(B, y_div_x, side='left').sub_(1) + j_C = torch.searchsorted(C, y_div_x, side='left').sub_(1) + left_derivative = 2 * torch.gather(Q_A, dim=-1, index=j_A) - torch.gather(Q_B, dim=-1, index=j_B) - torch.gather(Q_C, dim=-1, index=j_C) + j_A = torch.searchsorted(A, y_div_x, side='right').sub_(1) + j_B = torch.searchsorted(B, y_div_x, side='right').sub_(1) + j_C = torch.searchsorted(C, y_div_x, side='right').sub_(1) + right_derivative = 2 * torch.gather(Q_A, dim=-1, index=j_A) - torch.gather(Q_B, dim=-1, index=j_B) - torch.gather(Q_C, dim=-1, index=j_C) + + # Find extrema + is_extrema = (left_derivative < 0) & (right_derivative >= 0) + is_extrema[..., 0] |= ~is_extrema.any(dim=-1) # In case all derivatives are zero, take the first one as extrema. + where_extrema_batch, where_extrema_index = torch.where(is_extrema) + + # Calculate objective value at extrema + extrema_a = y_div_x[where_extrema_batch, where_extrema_index] # (num_extrema,) + MAX_ELEMENTS = 4096 ** 2 # Split into small batches to avoid OOM in case there are too many extrema.(~1G) + SPLIT_SIZE = MAX_ELEMENTS // x.shape[-1] + extrema_value = torch.cat([ + _compute_residual(extrema_a_split[:, None], xyw[extrema_i_split, :, :], trunc) + for extrema_a_split, extrema_i_split in zip(extrema_a.split(SPLIT_SIZE), where_extrema_batch.split(SPLIT_SIZE)) + ]) # (num_extrema,) + + # Find minima among corresponding extrema + minima, indices = scatter_min(size=batch_size, dim=0, index=where_extrema_batch, src=extrema_value) # (batch_size,) + index = where_extrema_index[indices] + + a = torch.gather(y, dim=-1, index=index[..., None]) / torch.gather(x, dim=-1, index=index[..., None]).clamp_min(eps) + a = a.reshape(batch_shape) + loss = minima.reshape(batch_shape) + index = index.reshape(batch_shape) + + return a, loss, index + + +def align_depth_affine(depth_src: torch.Tensor, depth_tgt: torch.Tensor, weight: Optional[torch.Tensor], trunc: Optional[Union[float, torch.Tensor]] = None): + """Fit robust affine scale and shift between paired scalar values. + + Args: + depth_src: Source depth/log-depth values ``[...,N]``. + depth_tgt: Target values ``[...,N]``. + weight: Nonnegative correspondence weights ``[...,N]``. + trunc: Optional robust residual cap forwarded to :func:`align`. + + Returns: + scale: Scalar multiplier tensor ``[...]``. + shift: Additive offset tensor ``[...]``. + """ + + # Flatten batch dimensions for simplicity + batch_shape, n = depth_src.shape[:-1], depth_src.shape[-1] + batch_size = math.prod(batch_shape) + depth_src, depth_tgt, weight = depth_src.reshape(batch_size, n), depth_tgt.reshape(batch_size, n), weight.reshape(batch_size, n) + + # Here, we take anchors only for non-zero weights. + # Although the results will be still correct even anchor points have zero weight, + # it is wasting computation and may cause instability in some cases, e.g. too many extrema. + anchors_where_batch, anchors_where_n = torch.where(weight > 0) + + # Stop gradient when solving optimal anchors + with torch.no_grad(): + depth_src_anchor = depth_src[anchors_where_batch, anchors_where_n] # (anchors) + depth_tgt_anchor = depth_tgt[anchors_where_batch, anchors_where_n] # (anchors) + + depth_src_anchored = depth_src[anchors_where_batch, :] - depth_src_anchor[..., None] # (anchors, n) + depth_tgt_anchored = depth_tgt[anchors_where_batch, :] - depth_tgt_anchor[..., None] # (anchors, n) + weight_anchored = weight[anchors_where_batch, :] # (anchors, n) + + scale, loss, index = align(depth_src_anchored, depth_tgt_anchored, weight_anchored, trunc) # (anchors) + + loss, index_anchor = scatter_min(size=batch_size, dim=0, index=anchors_where_batch, src=loss) # (batch_size,) + + # Reproduce by indexing for shorter compute graph + index_1 = anchors_where_n[index_anchor] # (batch_size,) + index_2 = index[index_anchor] # (batch_size,) + + tgt_1, src_1 = torch.gather(depth_tgt, dim=1, index=index_1[..., None]).squeeze(-1), torch.gather(depth_src, dim=1, index=index_1[..., None]).squeeze(-1) + tgt_2, src_2 = torch.gather(depth_tgt, dim=1, index=index_2[..., None]).squeeze(-1), torch.gather(depth_src, dim=1, index=index_2[..., None]).squeeze(-1) + + scale = (tgt_2 - tgt_1) / torch.where(src_2 != src_1, src_2 - src_1, 1e-7) + shift = tgt_1 - scale * src_1 + + scale, shift = scale.reshape(batch_shape), shift.reshape(batch_shape) + + return scale, shift + + +def align_points_scale_xyz_shift(points_src: torch.Tensor, points_tgt: torch.Tensor, weight: Optional[torch.Tensor], trunc: Optional[Union[float, torch.Tensor]] = None, max_iters: int = 30, eps: float = 1e-6): + """Fit one isotropic scale and a three-axis translation to point pairs. + + Args: + points_src: Source camera-space points ``[...,N,3]``. + points_tgt: Target points ``[...,N,3]`` with one-to-one correspondence. + weight: Nonnegative correspondence weights ``[...,N]``. + trunc: Optional robust residual cap. + max_iters: Compatibility argument retained from the research API. + eps: Compatibility stabilizer retained from the research API. + + Returns: + scale: Isotropic scale tensor ``[...]``. + shift: XYZ translation tensor ``[...,3]``. + """ + + # Flatten batch dimensions for simplicity + batch_shape, n = points_src.shape[:-2], points_src.shape[-2] + batch_size = math.prod(batch_shape) + points_src, points_tgt, weight = points_src.reshape(batch_size, n, 3), points_tgt.reshape(batch_size, n, 3), weight.reshape(batch_size, n) + + # Take anchors + anchor_where_batch, anchor_where_n = torch.where(weight > 0) + + with torch.no_grad(): + points_src_anchor = points_src[anchor_where_batch, anchor_where_n] # (anchors, 3) + points_tgt_anchor = points_tgt[anchor_where_batch, anchor_where_n] # (anchors, 3) + + points_src_anchored = points_src[anchor_where_batch, :, :] - points_src_anchor[..., None, :] # (anchors, n, 3) + points_tgt_anchored = points_tgt[anchor_where_batch, :, :] - points_tgt_anchor[..., None, :] # (anchors, n, 3) + weight_anchored = weight[anchor_where_batch, :, None].expand(-1, -1, 3) # (anchors, n, 3) + + # Solve optimal scale and shift for each anchor + MAX_ELEMENTS = 2 ** 20 + scale, loss, index = split_batch_fwd(align, MAX_ELEMENTS // 2, points_src_anchored.flatten(-2), points_tgt_anchored.flatten(-2), weight_anchored.flatten(-2), trunc) # (anchors,) + + # Get optimal scale and shift for each batch element + loss, index_anchor = scatter_min(size=batch_size, dim=0, index=anchor_where_batch, src=loss) # (batch_size,) + + index_2 = index[index_anchor] # (batch_size,) [0, 3n) + index_1 = anchor_where_n[index_anchor] * 3 + index_2 % 3 # (batch_size,) [0, 3n) + + src_1, tgt_1 = torch.gather(points_src.flatten(-2), dim=1, index=index_1[..., None]).squeeze(-1), torch.gather(points_tgt.flatten(-2), dim=1, index=index_1[..., None]).squeeze(-1) + src_2, tgt_2 = torch.gather(points_src.flatten(-2), dim=1, index=index_2[..., None]).squeeze(-1), torch.gather(points_tgt.flatten(-2), dim=1, index=index_2[..., None]).squeeze(-1) + + scale = (tgt_2 - tgt_1) / torch.where(src_2 != src_1, src_2 - src_1, 1.0) + shift = torch.gather(points_tgt, dim=1, index=(index_1 // 3)[..., None, None].expand(-1, -1, 3)).squeeze(-2) - scale[..., None] * torch.gather(points_src, dim=1, index=(index_1 // 3)[..., None, None].expand(-1, -1, 3)).squeeze(-2) + + scale, shift = scale.reshape(batch_shape), shift.reshape(*batch_shape, 3) + + return scale, shift + + +def align_affine_lstsq(x: torch.Tensor, y: torch.Tensor, w: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]: + """Fit weighted affine scale and shift by linear least squares. + + Args: + x: Source scalar values ``[...,N]``. + y: Target scalar values ``[...,N]``. + w: Optional nonnegative least-squares weights ``[...,N]``. ``None`` + assigns unit weight. + + Returns: + scale: Least-squares multiplier tensor ``[...]``. + shift: Least-squares additive offset tensor ``[...]``. + """ + w_sqrt = torch.ones_like(x) if w is None else w.sqrt() + A = torch.stack([w_sqrt * x, torch.ones_like(x)], dim=-1) + B = (w_sqrt * y)[..., None] + a, b = torch.linalg.lstsq(A, B)[0].squeeze(-1).unbind(-1) + return a, b diff --git a/pxdepth/utils/data_augmentation.py b/pxdepth/utils/data_augmentation.py new file mode 100644 index 0000000000000000000000000000000000000000..5f4dfd088766051240d2eddd46ed46b0345e396a --- /dev/null +++ b/pxdepth/utils/data_augmentation.py @@ -0,0 +1,609 @@ +"""Geometry-aware RGB-depth-camera transforms and blur primitives. + +Every geometric function updates normalized camera intrinsics together with RGB +and depth. Mixed depth interpolation uses bilinear disparity on smooth surfaces +and nearest sampling around geometric discontinuities to avoid flying points. +The evaluation loader uses the geometry-preserving transforms in this module. +""" + +from typing import Literal, Optional, Tuple + +import numpy as np +import cv2 +from PIL import Image +import utils3d +from scipy.signal import fftconvolve + +def sample_perspective( + src_intrinsics: np.ndarray, + tgt_aspect: float, + center_augmentation: float, + fov_range_absolute: Tuple[float, float], + fov_range_relative: Tuple[float, float], + rng: np.random.Generator = None +) -> Tuple[np.ndarray, np.ndarray]: + """Sample a valid target pinhole view inside a source camera frustum. + + Args: + src_intrinsics: Normalized source camera matrix ``float [3,3]``. + tgt_aspect: Target width divided by target height. + center_augmentation: Fraction controlling random optical-axis movement. + fov_range_absolute: Minimum/maximum target FoV in degrees. + fov_range_relative: Multipliers limiting target FoV relative to source. + rng: NumPy random generator used for FoV and center sampling. + + Returns: + tgt_intrinsics: Normalized target camera matrix ``float32 [3,3]``. + rotation: Camera-space rotation ``float32 [3,3]`` mapping source rays + into the sampled target view. + """ + raw_horizontal, raw_vertical = abs(1.0 / src_intrinsics[0, 0]), abs(1.0 / src_intrinsics[1, 1]) + raw_fov_x, raw_fov_y = utils3d.np.intrinsics_to_fov(src_intrinsics) + + # 1. set target fov + fov_range_absolute_min, fov_range_absolute_max = fov_range_absolute + fov_range_relative_min, fov_range_relative_max = fov_range_relative + tgt_fov_x_min = min(fov_range_relative_min * raw_fov_x, utils3d.focal_to_fov(utils3d.fov_to_focal(fov_range_relative_min * raw_fov_y) / tgt_aspect)) + tgt_fov_x_max = min(fov_range_relative_max * raw_fov_x, utils3d.focal_to_fov(utils3d.fov_to_focal(fov_range_relative_max * raw_fov_y) / tgt_aspect)) + tgt_fov_x_min, tgt_fov_max = max(np.deg2rad(fov_range_absolute_min), tgt_fov_x_min), min(np.deg2rad(fov_range_absolute_max), tgt_fov_x_max) + tgt_fov_x = rng.uniform(min(tgt_fov_x_min, tgt_fov_x_max), tgt_fov_x_max) + tgt_fov_y = utils3d.focal_to_fov(utils3d.np.fov_to_focal(tgt_fov_x) * tgt_aspect) + + # 2. set target image center (principal point) and the corresponding z-direction in raw camera space + center_dtheta = center_augmentation * rng.uniform(-0.5, 0.5) * (raw_fov_x - tgt_fov_x) + center_dphi = center_augmentation * rng.uniform(-0.5, 0.5) * (raw_fov_y - tgt_fov_y) + cu, cv = 0.5 + 0.5 * np.tan(center_dtheta) / np.tan(raw_fov_x / 2), 0.5 + 0.5 * np.tan(center_dphi) / np.tan(raw_fov_y / 2) + direction = utils3d.np.unproject_cv(np.array([[cu, cv]], dtype=np.float32), np.array([1.0], dtype=np.float32), intrinsics=src_intrinsics)[0] + + # 3. obtain the rotation matrix for homography warping (new_ext = R * old_ext) + R = utils3d.np.rotation_matrix_from_vectors(direction, np.array([0, 0, 1], dtype=np.float32)) + + # 4. shrink the target view to fit into the warped image + corners = np.array([[0, 0], [0, 1], [1, 1], [1, 0]], dtype=np.float32) + corners = np.concatenate([corners, np.ones((4, 1), dtype=np.float32)], axis=1) @ (np.linalg.inv(src_intrinsics).T @ R.T) # corners in viewport's camera plane + corners = corners[:, :2] / corners[:, 2:3] + tgt_horizontal, tgt_vertical = np.tan(tgt_fov_x / 2) * 2, np.tan(tgt_fov_y / 2) * 2 + warp_horizontal, warp_vertical = float('inf'), float('inf') + for i in range(4): + intersection, _ = utils3d.np.ray_intersection( + np.array([0., 0.]), np.array([[tgt_aspect, 1.0], [tgt_aspect, -1.0]]), + corners[i - 1], corners[i] - corners[i - 1], + ) + warp_horizontal, warp_vertical = min(warp_horizontal, 2 * np.abs(intersection[:, 0]).min()), min(warp_vertical, 2 * np.abs(intersection[:, 1]).min()) + tgt_horizontal, tgt_vertical = min(tgt_horizontal, warp_horizontal), min(tgt_vertical, warp_vertical) + + # 5. obtain the target intrinsics + fx, fy = 1 / tgt_horizontal, 1 / tgt_vertical + tgt_intrinsics = utils3d.np.intrinsics_from_focal_center(fx, fy, 0.5, 0.5).astype(np.float32) + + return tgt_intrinsics, R + + +def warp_perspective( + src_map: Optional[np.ndarray] = None, + transform: Optional[np.ndarray] = None, + tgt_size: Optional[Tuple[int, int]] = None, + interpolation: Literal['nearest', 'bilinear', 'lanczos'] = 'nearest', + sparse_mask: Optional[np.ndarray] = None, +) -> np.ndarray: + """Warp an image-like array through a normalized planar homography. + + Lanczos downsampling first reduces the source to avoid aliasing. Sparse + nearest-neighbor input optionally uses mask-aware pre-resizing so isolated + samples are not discarded. + + Args: + src_map: Source array ``[H,W]`` or ``[H,W,C]``. + transform: Normalized 3x3 homography satisfying + ``p_target = transform @ p_source``. + tgt_size: Output tuple ``(height,width)``. + interpolation: ``nearest``, ``bilinear``, or ``lanczos``. + sparse_mask: Optional boolean source support ``[H,W]`` for sparse + nearest-neighbor maps. + + Returns: + Warped array ``[H_t,W_t]`` or ``[H_t,W_t,C]``. + """ + + tgt_height, tgt_width = tgt_size + src_height, src_width = src_map.shape[:2] + + # source to target transform + transform_pixel = np.array([[tgt_width, 0, -0.5], [0, tgt_height, -0.5], [0, 0, 1]], dtype=np.float32) @ transform @ np.array([[1 / src_width, 0, 0.5 / src_width], [0, 1 / src_height, 0.5 / src_height], [0, 0, 1]], dtype=np.float32) + # Get scale factor at the target center + w = np.dot(np.linalg.inv(transform_pixel)[2, :], np.array([tgt_width / 2, tgt_height / 2, 1], dtype=np.float32)) + scale_x, scale_y = w * np.linalg.norm(transform_pixel[:2, :2], axis=0) + + if interpolation == 'lanczos' and (scale_x < 0.8 or scale_y < 0.8): + # If lanczos & downsampling, use PIL to resize first to reduce aliasing + src_height, src_width = max(round(src_height * scale_y * 1.25), 16), max(round(src_width * scale_x * 1.25), 16) + src_map = np.array(Image.fromarray(src_map).resize((src_width, src_height), Image.Resampling.LANCZOS)) + elif interpolation == 'nearest' and sparse_mask is not None and (scale_x < 1 or scale_y < 1): + # If nearest and sparse, use mask-aware nearest resize first to avoid losing points + src_height, src_width = max(round(src_height * scale_y), 16), max(round(src_width * scale_x), 16) + src_map, _ = utils3d.np.masked_nearest_resize(src_map, mask=sparse_mask, size=(src_height, src_width)) + + # Recompute the pixel-space transform after resizing + transform_pixel = np.array([[tgt_width, 0, -0.5], [0, tgt_height, -0.5], [0, 0, 1]], dtype=np.float32) @ transform @ np.array([[1 / src_width, 0, 0.5 / src_width], [0, 1 / src_height, 0.5 / src_height], [0, 0, 1]], dtype=np.float32) + + # Remap + cv2_interpolation = {'nearest': cv2.INTER_NEAREST, 'bilinear': cv2.INTER_LINEAR, 'lanczos': cv2.INTER_LANCZOS4}[interpolation] + tgt_map = cv2.warpPerspective(src_map, transform_pixel, (tgt_width, tgt_height), flags=cv2_interpolation) + + return tgt_map + + +def crop_resize_view( + src_image: np.ndarray, + src_depth: np.ndarray, + src_intrinsics: np.ndarray, + tgt_size: Tuple[int, int], + rng: Optional[np.random.Generator] = None, + random_crop: bool = True, + image_interpolation: Literal['nearest', 'bilinear', 'lanczos'] = 'lanczos', + depth_interpolation: Literal['nearest', 'mixed'] = 'nearest', +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Crop a source view and resize it to an exact target resolution. + + Large images are cropped directly at target size. Smaller images first take + the largest crop matching target aspect ratio and then resize. Normalized + intrinsics are updated for crop geometry; a full-image resize alone does not + alter normalized values. + + Args: + src_image: RGB uint8 image ``[H_s,W_s,3]``. + src_depth: Depth map ``[H_s,W_s]`` with NaN/Inf semantics. + src_intrinsics: Normalized source matrix ``[3,3]``. + tgt_size: Output tuple ``(height,width)``. + rng: Optional random generator for crop offsets. + random_crop: Randomize offsets instead of center cropping. + image_interpolation: RGB resampling method. + depth_interpolation: ``nearest`` or edge-aware ``mixed``. + + Returns: + image: RGB array ``[H_t,W_t,3]``. + depth: Depth array ``float32 [H_t,W_t]``. + intrinsics: Updated normalized matrix ``float32 [3,3]``. + """ + if depth_interpolation not in {'nearest', 'mixed'}: + raise ValueError(f"crop_resize_view only supports nearest/mixed depth interpolation, got {depth_interpolation}.") + + tgt_height, tgt_width = tgt_size + src_height, src_width = src_image.shape[:2] + if src_height <= 0 or src_width <= 0 or tgt_height <= 0 or tgt_width <= 0: + raise ValueError( + f"Invalid source/target size: src=({src_height}, {src_width}), tgt=({tgt_height}, {tgt_width})." + ) + + if src_width >= tgt_width and src_height >= tgt_height: + crop_width = tgt_width + crop_height = tgt_height + else: + tgt_aspect = tgt_width / tgt_height + if src_width / src_height >= tgt_aspect: + crop_height = src_height + crop_width = max(1, min(src_width, int(np.floor(src_height * tgt_aspect)))) + else: + crop_width = src_width + crop_height = max(1, min(src_height, int(np.floor(src_width / tgt_aspect)))) + + max_x = max(0, src_width - crop_width) + max_y = max(0, src_height - crop_height) + if random_crop: + if rng is None: + rng = np.random.default_rng() + x0 = int(rng.integers(0, max_x + 1)) if max_x > 0 else 0 + y0 = int(rng.integers(0, max_y + 1)) if max_y > 0 else 0 + else: + x0 = max_x // 2 + y0 = max_y // 2 + + x1, y1 = x0 + crop_width, y0 + crop_height + cropped_image = src_image[y0:y1, x0:x1] + cropped_depth = src_depth[y0:y1, x0:x1] + + if crop_width != tgt_width or crop_height != tgt_height: + if image_interpolation == 'lanczos': + tgt_image = np.array( + Image.fromarray(cropped_image).resize((tgt_width, tgt_height), Image.Resampling.LANCZOS) + ) + else: + cv2_interpolation = { + 'nearest': cv2.INTER_NEAREST, + 'bilinear': cv2.INTER_LINEAR, + }[image_interpolation] + tgt_image = cv2.resize(cropped_image, (tgt_width, tgt_height), interpolation=cv2_interpolation) + + cropped_valid = np.isfinite(cropped_depth) + cropped_depth_values = np.where(cropped_valid, cropped_depth, 0).astype(np.float32) + tgt_depth_nearest = cv2.resize(cropped_depth_values, (tgt_width, tgt_height), interpolation=cv2.INTER_NEAREST) + tgt_depth_valid = cv2.resize(cropped_valid.astype(np.uint8), (tgt_width, tgt_height), interpolation=cv2.INTER_NEAREST).astype(bool) + + if depth_interpolation == 'mixed': + depth_edge_mask = utils3d.np.depth_map_edge(cropped_depth, mask=cropped_valid, kernel_size=5, ltol=0.005) + depth_bilinear_mask = cropped_valid & ~depth_edge_mask + tgt_depth_bilinear_mask = cv2.resize( + depth_bilinear_mask.astype(np.float32), + (tgt_width, tgt_height), + interpolation=cv2.INTER_LINEAR, + ) + cropped_disp = np.where(cropped_valid & (cropped_depth > 0), 1.0 / cropped_depth, 0.0).astype(np.float32) + tgt_disp_bilinear = cv2.resize(cropped_disp, (tgt_width, tgt_height), interpolation=cv2.INTER_LINEAR) + tgt_depth_bilinear = np.where(tgt_disp_bilinear > 0, 1.0 / tgt_disp_bilinear, np.inf).astype(np.float32) + tgt_depth = np.where(tgt_depth_bilinear_mask == 1.0, tgt_depth_bilinear, tgt_depth_nearest) + else: + tgt_depth = tgt_depth_nearest + + tgt_depth = np.where(tgt_depth_valid, tgt_depth, np.inf).astype(np.float32) + else: + tgt_image = cropped_image.copy() + tgt_depth = cropped_depth.copy().astype(np.float32) + + tgt_intrinsics = src_intrinsics.astype(np.float32).copy() + tgt_intrinsics[0, 0] = src_intrinsics[0, 0] * src_width / crop_width + tgt_intrinsics[0, 1] = src_intrinsics[0, 1] * src_width / crop_width + tgt_intrinsics[0, 2] = (src_intrinsics[0, 2] * src_width - x0) / crop_width + tgt_intrinsics[1, 0] = 0.0 + tgt_intrinsics[1, 1] = src_intrinsics[1, 1] * src_height / crop_height + tgt_intrinsics[1, 2] = (src_intrinsics[1, 2] * src_height - y0) / crop_height + tgt_intrinsics[2, 0] = 0.0 + tgt_intrinsics[2, 1] = 0.0 + tgt_intrinsics[2, 2] = 1.0 + + return tgt_image, tgt_depth, tgt_intrinsics + + +def resize_then_crop_view( + src_image: np.ndarray, + src_depth: np.ndarray, + src_intrinsics: np.ndarray, + resize_size: Tuple[int, int], + crop_size: Tuple[int, int], + rng: Optional[np.random.Generator] = None, + random_crop: bool = True, + image_interpolation: Literal['nearest', 'bilinear', 'area'] = 'area', + depth_interpolation: Literal['nearest'] = 'nearest', +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Resize the full view first, then crop an exact output window. + + Args: + src_image: RGB uint8 image ``[H_s,W_s,3]``. + src_depth: Depth map ``[H_s,W_s]``. + src_intrinsics: Normalized source matrix ``[3,3]``. + resize_size: Intermediate tuple ``(height,width)``. + crop_size: Final tuple ``(height,width)`` no larger than resize size. + rng: Optional random generator for crop offsets. + random_crop: Randomize offsets instead of center cropping. + image_interpolation: Intermediate RGB resampling method. + depth_interpolation: Depth method; only nearest is supported. + + Returns: + image: Cropped RGB ``[H_c,W_c,3]``. + depth: Cropped depth ``float32 [H_c,W_c]``. + intrinsics: Crop-adjusted normalized matrix ``float32 [3,3]``. + """ + if depth_interpolation != 'nearest': + raise ValueError(f"resize_then_crop_view only supports nearest depth interpolation, got {depth_interpolation}.") + + resize_height, resize_width = resize_size + crop_height, crop_width = crop_size + src_height, src_width = src_image.shape[:2] + if src_height <= 0 or src_width <= 0 or resize_height <= 0 or resize_width <= 0 or crop_height <= 0 or crop_width <= 0: + raise ValueError( + f"Invalid source/resize/crop size: src=({src_height}, {src_width}), " + f"resize=({resize_height}, {resize_width}), crop=({crop_height}, {crop_width})." + ) + if crop_height > resize_height or crop_width > resize_width: + raise ValueError( + f"Crop size ({crop_height}, {crop_width}) must not exceed resize size ({resize_height}, {resize_width})." + ) + + image_cv2_interpolation = { + 'nearest': cv2.INTER_NEAREST, + 'bilinear': cv2.INTER_LINEAR, + 'area': cv2.INTER_AREA, + }[image_interpolation] + resized_image = cv2.resize(src_image, (resize_width, resize_height), interpolation=image_cv2_interpolation) + + resized_depth = cv2.resize(src_depth.astype(np.float32), (resize_width, resize_height), interpolation=cv2.INTER_NEAREST) + + max_x = max(0, resize_width - crop_width) + max_y = max(0, resize_height - crop_height) + if random_crop: + if rng is None: + rng = np.random.default_rng() + x0 = int(rng.integers(0, max_x + 1)) if max_x > 0 else 0 + y0 = int(rng.integers(0, max_y + 1)) if max_y > 0 else 0 + else: + x0 = max_x // 2 + y0 = max_y // 2 + + x1, y1 = x0 + crop_width, y0 + crop_height + tgt_image = resized_image[y0:y1, x0:x1].copy() + tgt_depth = resized_depth[y0:y1, x0:x1].copy().astype(np.float32) + + tgt_intrinsics = src_intrinsics.astype(np.float32).copy() + tgt_intrinsics[0, 0] = src_intrinsics[0, 0] * resize_width / crop_width + tgt_intrinsics[0, 1] = src_intrinsics[0, 1] * resize_width / crop_width + tgt_intrinsics[0, 2] = (src_intrinsics[0, 2] * resize_width - x0) / crop_width + tgt_intrinsics[1, 0] = 0.0 + tgt_intrinsics[1, 1] = src_intrinsics[1, 1] * resize_height / crop_height + tgt_intrinsics[1, 2] = (src_intrinsics[1, 2] * resize_height - y0) / crop_height + tgt_intrinsics[2, 0] = 0.0 + tgt_intrinsics[2, 1] = 0.0 + tgt_intrinsics[2, 2] = 1.0 + + return tgt_image, tgt_depth, tgt_intrinsics + + +def resize_to_cover_center_crop_view( + src_image: np.ndarray, + src_depth: np.ndarray, + src_intrinsics: np.ndarray, + tgt_size: Tuple[int, int], + image_interpolation: Literal['nearest', 'bilinear', 'lanczos'] = 'lanczos', + depth_interpolation: Literal['nearest', 'mixed'] = 'mixed', +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Resize a view to cover the target, then take a centered crop. + + Args: + src_image: RGB uint8 image ``[H_s,W_s,3]``. + src_depth: Depth map ``[H_s,W_s]`` with invalid sentinels. + src_intrinsics: Normalized source matrix ``[3,3]``. + tgt_size: Output tuple ``(height,width)``. + image_interpolation: RGB resampling method. + depth_interpolation: ``nearest`` or edge-aware ``mixed``. + + Returns: + image: Center-cropped RGB ``[H_t,W_t,3]``. + depth: Center-cropped depth ``float32 [H_t,W_t]``. + intrinsics: Resize/crop-adjusted normalized matrix ``float32 [3,3]``. + """ + if depth_interpolation not in {'nearest', 'mixed'}: + raise ValueError( + f"resize_to_cover_center_crop_view only supports nearest/mixed depth interpolation, got {depth_interpolation}." + ) + + tgt_height, tgt_width = tgt_size + src_height, src_width = src_image.shape[:2] + if src_height <= 0 or src_width <= 0 or tgt_height <= 0 or tgt_width <= 0: + raise ValueError( + f"Invalid source/target size: src=({src_height}, {src_width}), tgt=({tgt_height}, {tgt_width})." + ) + + scale = max(tgt_width / src_width, tgt_height / src_height) + resized_width = max(tgt_width, int(round(src_width * scale))) + resized_height = max(tgt_height, int(round(src_height * scale))) + + if image_interpolation == 'lanczos': + resized_image = np.array( + Image.fromarray(src_image).resize((resized_width, resized_height), Image.Resampling.LANCZOS) + ) + else: + cv2_interpolation = { + 'nearest': cv2.INTER_NEAREST, + 'bilinear': cv2.INTER_LINEAR, + }[image_interpolation] + resized_image = cv2.resize(src_image, (resized_width, resized_height), interpolation=cv2_interpolation) + + resized_valid = np.isfinite(src_depth) + resized_depth_values = cv2.resize( + np.where(resized_valid, src_depth, 0).astype(np.float32), + (resized_width, resized_height), + interpolation=cv2.INTER_NEAREST, + ) + resized_depth_valid = cv2.resize( + resized_valid.astype(np.uint8), + (resized_width, resized_height), + interpolation=cv2.INTER_NEAREST, + ).astype(bool) + resized_depth = np.where(resized_depth_valid, resized_depth_values, np.inf).astype(np.float32) + + if depth_interpolation == 'mixed': + depth_edge_mask = utils3d.np.depth_map_edge(src_depth, mask=np.isfinite(src_depth), kernel_size=5, ltol=0.005) + depth_bilinear_mask = np.isfinite(src_depth) & ~depth_edge_mask + resized_bilinear_mask = cv2.resize( + depth_bilinear_mask.astype(np.float32), + (resized_width, resized_height), + interpolation=cv2.INTER_LINEAR, + ) + disp = np.where(np.isfinite(src_depth) & (src_depth > 0), 1.0 / src_depth, 0.0).astype(np.float32) + resized_disp = cv2.resize(disp, (resized_width, resized_height), interpolation=cv2.INTER_LINEAR) + resized_depth_bilinear = np.where(resized_disp > 0, 1.0 / resized_disp, np.inf).astype(np.float32) + resized_depth = np.where(resized_bilinear_mask == 1.0, resized_depth_bilinear, resized_depth) + resized_depth = np.where(resized_depth_valid, resized_depth, np.inf).astype(np.float32) + + x0 = max(0, (resized_width - tgt_width) // 2) + y0 = max(0, (resized_height - tgt_height) // 2) + 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_intrinsics = src_intrinsics.astype(np.float32).copy() + tgt_intrinsics[0, 0] = src_intrinsics[0, 0] * resized_width / tgt_width + tgt_intrinsics[0, 1] = src_intrinsics[0, 1] * resized_width / tgt_width + tgt_intrinsics[0, 2] = (src_intrinsics[0, 2] * resized_width - x0) / tgt_width + tgt_intrinsics[1, 0] = 0.0 + tgt_intrinsics[1, 1] = src_intrinsics[1, 1] * resized_height / tgt_height + tgt_intrinsics[1, 2] = (src_intrinsics[1, 2] * resized_height - y0) / tgt_height + tgt_intrinsics[2, 0] = 0.0 + tgt_intrinsics[2, 1] = 0.0 + tgt_intrinsics[2, 2] = 1.0 + + return tgt_image, tgt_depth, tgt_intrinsics + + +def resize_view( + src_image: np.ndarray, + src_depth: np.ndarray, + src_intrinsics: np.ndarray, + tgt_size: Tuple[int, int], + image_interpolation: Literal['nearest', 'bilinear', 'lanczos'] = 'lanczos', + depth_interpolation: Literal['nearest', 'mixed'] = 'mixed', +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Resize RGB and depth directly to a target width and height. + + Args: + src_image: RGB uint8 image ``[H_s,W_s,3]``. + src_depth: Depth map ``[H_s,W_s]`` with invalid sentinels. + src_intrinsics: Normalized source matrix ``[3,3]``. + tgt_size: Exact output tuple ``(height,width)``. + image_interpolation: RGB resampling method. + depth_interpolation: ``nearest`` or edge-aware ``mixed``. + + Returns: + image: Resized RGB ``[H_t,W_t,3]``. + depth: Resized depth ``float32 [H_t,W_t]``. + intrinsics: Normalized matrix ``float32 [3,3]``. Direct full-image + resizing leaves normalized focal lengths and center unchanged. + """ + if depth_interpolation not in {'nearest', 'mixed'}: + raise ValueError(f"resize_view only supports nearest/mixed depth interpolation, got {depth_interpolation}.") + + tgt_height, tgt_width = tgt_size + src_height, src_width = src_image.shape[:2] + if src_height <= 0 or src_width <= 0 or tgt_height <= 0 or tgt_width <= 0: + raise ValueError( + f"Invalid source/target size: src=({src_height}, {src_width}), tgt=({tgt_height}, {tgt_width})." + ) + + if image_interpolation == 'lanczos': + tgt_image = np.array(Image.fromarray(src_image).resize((tgt_width, tgt_height), Image.Resampling.LANCZOS)) + else: + cv2_interpolation = { + 'nearest': cv2.INTER_NEAREST, + 'bilinear': cv2.INTER_LINEAR, + }[image_interpolation] + tgt_image = cv2.resize(src_image, (tgt_width, tgt_height), interpolation=cv2_interpolation) + + src_valid = np.isfinite(src_depth) + tgt_depth_values = cv2.resize( + np.where(src_valid, src_depth, 0).astype(np.float32), + (tgt_width, tgt_height), + interpolation=cv2.INTER_NEAREST, + ) + tgt_depth_valid = cv2.resize( + src_valid.astype(np.uint8), + (tgt_width, tgt_height), + interpolation=cv2.INTER_NEAREST, + ).astype(bool) + tgt_depth = np.where(tgt_depth_valid, tgt_depth_values, np.inf).astype(np.float32) + + if depth_interpolation == 'mixed': + depth_edge_mask = utils3d.np.depth_map_edge(src_depth, mask=src_valid, kernel_size=5, ltol=0.005) + depth_bilinear_mask = src_valid & ~depth_edge_mask + tgt_depth_bilinear_mask = cv2.resize( + depth_bilinear_mask.astype(np.float32), + (tgt_width, tgt_height), + interpolation=cv2.INTER_LINEAR, + ) + src_disp = np.where(src_valid & (src_depth > 0), 1.0 / src_depth, 0.0).astype(np.float32) + tgt_disp_bilinear = cv2.resize(src_disp, (tgt_width, tgt_height), interpolation=cv2.INTER_LINEAR) + tgt_depth_bilinear = np.where(tgt_disp_bilinear > 0, 1.0 / tgt_disp_bilinear, np.inf).astype(np.float32) + tgt_depth = np.where(tgt_depth_bilinear_mask == 1.0, tgt_depth_bilinear, tgt_depth) + tgt_depth = np.where(tgt_depth_valid, tgt_depth, np.inf).astype(np.float32) + + tgt_intrinsics = src_intrinsics.astype(np.float32).copy() + tgt_intrinsics[1, 0] = 0.0 + tgt_intrinsics[2, 0] = 0.0 + tgt_intrinsics[2, 1] = 0.0 + tgt_intrinsics[2, 2] = 1.0 + + return tgt_image, tgt_depth, tgt_intrinsics + + +def disk_kernel(radius: int) -> np.ndarray: + """Generate a normalized circular convolution kernel. + + Args: + radius: Nonnegative disk radius in pixels. + + Returns: + Float32 kernel ``[2*radius+1,2*radius+1]`` summing to one. + """ + # Create coordinate grid centered at (0,0) + L = np.arange(-radius, radius + 1) + X, Y = np.meshgrid(L, L) + # Generate disk: region inside circle with radius R is 1 + kernel = ((X**2 + Y**2) <= radius**2).astype(np.float32) + # Normalize the kernel + kernel /= np.sum(kernel) + return kernel + + +def disk_blur(image: np.ndarray, radius: int) -> np.ndarray: + """Apply a circular point-spread function with FFT convolution. + + Args: + image: Scalar ``[H,W]`` or channel image ``[H,W,C]``. + radius: Nonnegative blur radius in pixels. + + Returns: + Blurred floating array with the same shape as ``image``. + """ + if radius == 0: + return image + kernel = disk_kernel(radius) + if image.ndim == 2: + blurred = fftconvolve(image, kernel, mode='same') + elif image.ndim == 3: + channels = [] + for i in range(image.shape[2]): + blurred_channel = fftconvolve(image[..., i], kernel, mode='same') + channels.append(blurred_channel) + blurred = np.stack(channels, axis=-1) + else: + raise ValueError("Image must be 2D or 3D.") + return blurred + + +def depth_of_field( + img: np.ndarray, + disp: np.ndarray, + focus_disp : float, + max_blur_radius : int = 10, +) -> np.ndarray: + """Synthesize depth of field from a disparity map and focus plane. + + Args: + img: RGB image ``[H,W,3]``. + disp: Positive disparity map ``[H,W]`` aligned to ``img``. + focus_disp: Disparity value lying on the simulated focus plane. + max_blur_radius: Largest circular blur radius in pixels. + + Returns: + Depth-of-field image with shape ``[H,W,3]`` and ``img`` dtype. + """ + # Precalculate dialated depth map for each blur radius + max_disp = np.max(disp) + disp = disp / max_disp + focus_disp = focus_disp / max_disp + dilated_disp = [] + for radius in range(max_blur_radius + 1): + dilated_disp.append(cv2.dilate(disp, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1, 2 * radius + 1)), iterations=1)) + + # Determine the blur radius for each pixel based on the depth map + blur_radii = np.clip(np.abs(disp - focus_disp) * max_blur_radius, 0, max_blur_radius).astype(np.int32) + for radius in range(max_blur_radius + 1): + dialted_blur_radii = np.clip(np.abs(dilated_disp[radius] - focus_disp) * max_blur_radius, 0, max_blur_radius).astype(np.int32) + mask = (dialted_blur_radii >= radius) & (dialted_blur_radii >= blur_radii) & (dilated_disp[radius] > disp) + blur_radii[mask] = dialted_blur_radii[mask] + blur_radii = np.clip(blur_radii, 0, max_blur_radius) + blur_radii = cv2.blur(blur_radii, (5, 5)) + + # Precalculate the blured image for each blur radius + unique_radii = np.unique(blur_radii) + precomputed = {} + for radius in range(max_blur_radius + 1): + if radius not in unique_radii: + continue + precomputed[radius] = disk_blur(img, radius) + + # Composit the blured image for each pixel + output = np.zeros_like(img) + for r in unique_radii: + mask = blur_radii == r + output[mask] = precomputed[r][mask] + + return output diff --git a/pxdepth/utils/io.py b/pxdepth/utils/io.py new file mode 100644 index 0000000000000000000000000000000000000000..41b00af7be80e27c673462fc0e427525f0efe7a1 --- /dev/null +++ b/pxdepth/utils/io.py @@ -0,0 +1,209 @@ +"""I/O for the processed benchmark format used by PXDepth evaluation. + +Benchmark samples store RGB as JPEG or PNG, depth as a logarithmically encoded +16-bit PNG, optional semantic labels as PNG metadata, and camera information in +JSON. The matching writers are shared by the released evaluation-dataset +converters so their outputs can be consumed directly by the benchmark loader. +""" + +import io +import json +import os +from pathlib import Path +from typing import Any, Dict, IO, List, Optional, Tuple, Union + +import cv2 +import numpy as np +from PIL import Image, PngImagePlugin + + +PathOrBinary = Union[str, os.PathLike, IO[bytes]] +JsonValue = Union[str, int, float, bool, None, Dict[str, Any], List[Any]] + + +def _read_bytes(path: PathOrBinary) -> bytes: + """Read encoded data from a filesystem path or binary stream. + + Args: + path: File path or a binary stream exposing ``read()``. + + Returns: + Encoded file contents as ``bytes``. + """ + if isinstance(path, (str, os.PathLike)): + return Path(path).read_bytes() + return path.read() + + +def read_image(path: PathOrBinary) -> np.ndarray: + """Decode an RGB image. + + Args: + path: JPEG/PNG path or readable binary stream. + + Returns: + RGB uint8 array with shape ``[H, W, 3]``. + + Raises: + ValueError: If OpenCV cannot decode the input. + """ + image = cv2.imdecode(np.frombuffer(_read_bytes(path), np.uint8), cv2.IMREAD_COLOR) + if image is None: + raise ValueError(f"Unable to decode image: {path}") + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + +def write_image(path: Union[str, os.PathLike, IO[bytes]], image: np.ndarray, quality: int = 95) -> None: + """Encode an RGB image as JPEG. + + Args: + path: Destination path or writable binary stream. + image: RGB uint8 array with shape ``[H, W, 3]``. + quality: JPEG quality passed to OpenCV. + """ + encoded = cv2.imencode( + ".jpg", + cv2.cvtColor(image, cv2.COLOR_RGB2BGR), + [cv2.IMWRITE_JPEG_QUALITY, int(quality)], + )[1].tobytes() + if isinstance(path, (str, os.PathLike)): + Path(path).write_bytes(encoded) + else: + path.write(encoded) + + +def read_depth(path: PathOrBinary) -> np.ndarray: + """Decode the logarithmic 16-bit PNG depth representation. + + Args: + path: Encoded depth PNG path or readable binary stream. The PNG must + contain ``near`` and ``far`` text metadata. + + Returns: + Float32 depth array ``[H, W]``. Code 0 maps to NaN, code 65535 maps to + positive infinity, and codes 1 through 65534 map to finite depth. + """ + image = Image.open(io.BytesIO(_read_bytes(path))) + near = float(image.info["near"]) + far = float(image.info["far"]) + encoded = np.asarray(image) + mask_nan = encoded == 0 + mask_inf = encoded == 65535 + value = (encoded.astype(np.float32) - 1.0) / 65533.0 + depth = near ** (1.0 - value) * far**value + if "unit" in image.info: + depth *= float(image.info["unit"]) + depth[mask_nan] = np.nan + depth[mask_inf] = np.inf + return depth + + +def write_depth( + path: Union[str, os.PathLike, IO[bytes]], + depth: np.ndarray, + max_range: float = 1e5, + compression_level: int = 7, +) -> None: + """Encode depth as logarithmic 16-bit PNG with NaN/Inf sentinels. + + Args: + path: Destination path or writable binary stream. + depth: Float depth array ``[H, W]``. NaN stores unknown geometry and + positive infinity stores known infinite geometry. + max_range: Maximum finite ``far / near`` encoding ratio. + compression_level: PNG compression level from zero through nine. + """ + depth = np.asarray(depth, dtype=np.float32) + finite = np.isfinite(depth) + mask_nan = np.isnan(depth) + mask_inf = np.isinf(depth) + if not np.any(finite): + raise ValueError("Depth encoding requires at least one finite value.") + near = max(float(depth[finite].min()), 1e-5) + far = max(near * 1.1, min(float(depth[finite].max()), near * float(max_range))) + clipped = np.nan_to_num(depth, nan=near, posinf=far, neginf=near).clip(near, far) + encoded = 1 + np.round(np.log(clipped / near) / np.log(far / near) * 65533).astype(np.uint16) + encoded[mask_nan] = 0 + encoded[mask_inf] = 65535 + pnginfo = PngImagePlugin.PngInfo() + pnginfo.add_text("near", str(near)) + pnginfo.add_text("far", str(far)) + Image.fromarray(encoded).save(path, pnginfo=pnginfo, compress_level=int(compression_level)) + + +def read_segmentation(path: PathOrBinary) -> Tuple[np.ndarray, Optional[Dict[str, int]]]: + """Decode an integer segmentation PNG and its optional label mapping. + + Args: + path: Segmentation PNG path or readable binary stream. + + Returns: + A pair of ``mask`` and ``labels``. ``mask`` has shape ``[H, W]`` and + retains the PNG integer dtype. ``labels`` maps names to IDs when the + PNG contains label metadata, otherwise it is ``None``. + """ + image = Image.open(io.BytesIO(_read_bytes(path))) + labels = json.loads(image.info["labels"]) if "labels" in image.info else None + return np.asarray(image), labels + + +def write_segmentation( + path: Union[str, os.PathLike, IO[bytes]], + mask: np.ndarray, + labels: Optional[Dict[str, int]] = None, + compression_level: int = 7, +) -> None: + """Write an integer segmentation PNG and optional label mapping. + + Args: + path: Destination path or writable binary stream. + mask: Integer label array ``[H, W]`` with uint8 or uint16 dtype. + labels: Optional mapping from label names to integer IDs. + compression_level: PNG compression level from zero through nine. + """ + mask = np.asarray(mask) + if mask.dtype not in (np.uint8, np.uint16): + raise TypeError(f"Segmentation must be uint8 or uint16, got {mask.dtype}.") + pnginfo = PngImagePlugin.PngInfo() + if labels is not None: + pnginfo.add_text("labels", json.dumps(labels, ensure_ascii=True, separators=(",", ":"))) + Image.fromarray(mask).save(path, pnginfo=pnginfo, compress_level=int(compression_level)) + + +def read_json(path: Union[str, os.PathLike, IO[str]]) -> JsonValue: + """Parse JSON from a path or readable text stream. + + Args: + path: JSON path or text stream exposing ``read()``. + + Returns: + Parsed JSON-compatible Python value. + """ + text = Path(path).read_text() if isinstance(path, (str, os.PathLike)) else path.read() + return json.loads(text) + + +def write_json(path: Union[str, os.PathLike, IO[str]], content: JsonValue) -> None: + """Serialize a JSON-compatible value. + + Args: + path: Destination path or writable text stream. + content: JSON-compatible scalar, list, or dictionary. + """ + text = json.dumps(content) + if isinstance(path, (str, os.PathLike)): + Path(path).write_text(text) + else: + path.write(text) + + +__all__ = [ + "read_depth", + "read_image", + "read_json", + "read_segmentation", + "write_depth", + "write_image", + "write_json", + "write_segmentation", +] diff --git a/pxdepth/utils/ply.py b/pxdepth/utils/ply.py new file mode 100644 index 0000000000000000000000000000000000000000..0dd555fc070c40819ae38a5b22cbdd6187ce7770 --- /dev/null +++ b/pxdepth/utils/ply.py @@ -0,0 +1,74 @@ +"""Minimal binary little-endian PLY export for dense colored point clouds. + +Finite XYZ samples and their corresponding RGB values are flattened, filtered, +and serialized with a standards-compliant vertex header. The implementation is +dependency-light and is used by inference and evaluation dumps. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + + +def write_point_cloud_ply(path: Path, points: np.ndarray, colors: np.ndarray) -> None: + """Write XYZ points and RGB colors as a binary PLY vertex list. + + Args: + path: Destination file. Parent directories are created automatically. + points: Floating point positions ``[N,3]``. + colors: RGB values ``[N,3]``. Floating arrays are interpreted in + ``[0,1]``; integer arrays are interpreted in ``[0,255]``. + + Returns: + ``None``. A binary little-endian PLY file is written. + """ + points = np.asarray(points, dtype=np.float32) + colors = np.asarray(colors) + if points.ndim != 2 or points.shape[1] != 3: + raise ValueError(f"points must have shape [N, 3], got {points.shape}") + if colors.ndim != 2 or colors.shape[1] != 3: + raise ValueError(f"colors must have shape [N, 3], got {colors.shape}") + if points.shape[0] != colors.shape[0]: + raise ValueError(f"points/colors length mismatch: {points.shape[0]} vs {colors.shape[0]}") + + if colors.dtype.kind == "f": + colors = np.clip(colors, 0.0, 1.0) * 255.0 + colors = np.clip(colors, 0, 255).astype(np.uint8) + + vertices = np.empty( + points.shape[0], + dtype=[ + ("x", " Dict[str, Any]: + """Recursively average matching finite numeric dictionary leaves. + + Args: + items: List of nested dictionaries. Missing keys are ignored, nested + dictionaries recurse, numeric NaNs are excluded, and unsupported + leaf types are omitted. + + Returns: + Dictionary with the union of supported keys and arithmetic means at + numeric leaves. Returns an empty dictionary for an empty list. + """ + if not items: + return {} + keys = set().union(*(item.keys() for item in items)) + result: Dict[str, Any] = {} + for key in keys: + values = [item[key] for item in items if key in item] + if not values: + continue + if isinstance(values[0], dict): + result[key] = key_average(values) + elif isinstance(values[0], Number): + finite_values = [value for value in values if not math.isnan(float(value))] + result[key] = sum(finite_values) / len(finite_values) if finite_values else float("nan") + return result diff --git a/pxdepth/utils/vis.py b/pxdepth/utils/vis.py new file mode 100644 index 0000000000000000000000000000000000000000..c66a774f617f74bd58c51c3b59a78cb1a3a34130 --- /dev/null +++ b/pxdepth/utils/vis.py @@ -0,0 +1,172 @@ +"""NumPy visualization mappings for geometry and semantic predictions. + +Depth, disparity, normals, labels, and scalar error arrays are normalized with +explicit valid masks and converted to display-ready RGB images. Unknown or +infinite regions use stable colors shared by training and evaluation outputs. +""" + +from typing import Optional, Tuple + +import numpy as np +import matplotlib + + +def colorize_depth(depth: np.ndarray, mask: Optional[np.ndarray] = None, normalize: bool = True, cmap: str = 'Spectral') -> np.ndarray: + """Colorize positive depth through inverse-depth ordering. + + Args: + depth: Depth map ``[H,W]``. + mask: Optional boolean validity mask ``[H,W]``. + normalize: Quantile-normalize disparity before colormap lookup. + cmap: Matplotlib colormap name. + + Returns: + RGB uint8 visualization ``[H,W,3]``; invalid pixels are black. + """ + if mask is None: + depth = np.where(depth > 0, depth, np.nan) + else: + depth = np.where((depth > 0) & mask, depth, np.nan) + disp = 1 / depth + if normalize: + min_disp, max_disp = np.nanquantile(disp, 0.001), np.nanquantile(disp, 0.99) + disp = (disp - min_disp) / (max_disp - min_disp) + colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - disp)[..., :3], 0) + colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8)) + return colored + + +def colorize_depth_affine(depth: np.ndarray, mask: Optional[np.ndarray] = None, cmap: str = 'Spectral') -> np.ndarray: + """Colorize depth after direct affine quantile normalization. + + Args: + depth: Scalar depth-like map ``[H,W]``. + mask: Optional boolean validity mask ``[H,W]``. + cmap: Matplotlib colormap name. + + Returns: + RGB uint8 visualization ``[H,W,3]``. + """ + if mask is not None: + depth = np.where(mask, depth, np.nan) + + min_depth, max_depth = np.nanquantile(depth, 0.001), np.nanquantile(depth, 0.999) + depth = (depth - min_depth) / (max_depth - min_depth) + colored = np.nan_to_num(matplotlib.colormaps[cmap](depth)[..., :3], 0) + colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8)) + return colored + + +def colorize_depth_shifted_disparity( + depth: np.ndarray, + mask: Optional[np.ndarray] = None, + normalize: bool = True, + cmap: str = 'Spectral', + eps: float = 1.0, +) -> np.ndarray: + """Colorize depth using disparity shifted by the nearest finite value. + + Args: + depth: Depth-like map ``[H,W]`` that may include negative values. + mask: Optional boolean validity mask ``[H,W]``. + normalize: Quantile-normalize shifted disparity. + cmap: Matplotlib colormap name. + eps: Positive offset preventing division by zero at minimum depth. + + Returns: + RGB uint8 visualization ``[H,W,3]``. + """ + if mask is not None: + depth = np.where(mask, depth, np.nan) + else: + depth = np.where(np.isfinite(depth), depth, np.nan) + if not np.isfinite(depth).any(): + return np.zeros((*depth.shape, 3), dtype=np.uint8) + + min_depth = np.nanmin(depth) + disp = 1.0 / (depth - min_depth + eps) + if normalize: + min_disp, max_disp = np.nanquantile(disp, 0.001), np.nanquantile(disp, 0.99) + if max_disp > min_disp: + disp = (disp - min_disp) / (max_disp - min_disp) + else: + disp = np.zeros_like(disp) + colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - disp)[..., :3], 0) + colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8)) + return colored + + +def colorize_disparity(disparity: np.ndarray, mask: Optional[np.ndarray] = None, normalize: bool = True, cmap: str = 'Spectral') -> np.ndarray: + """Colorize a disparity map with optional quantile normalization. + + Args: + disparity: Disparity array ``[H,W]``. + mask: Optional boolean validity mask ``[H,W]``. + normalize: Normalize the 0.1%--99.9% quantile interval. + cmap: Matplotlib colormap name. + + Returns: + RGB uint8 visualization ``[H,W,3]``. + """ + if mask is not None: + disparity = np.where(mask, disparity, np.nan) + + if normalize: + min_disp, max_disp = np.nanquantile(disparity, 0.001), np.nanquantile(disparity, 0.999) + disparity = (disparity - min_disp) / (max_disp - min_disp) + colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - disparity)[..., :3], 0) + colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8)) + return colored + + +def colorize_segmentation(segmentation: np.ndarray, cmap: str = 'Set1') -> np.ndarray: + """Assign repeating categorical colors to integer segmentation IDs. + + Args: + segmentation: Integer label map ``[H,W]``. + cmap: Matplotlib categorical colormap name. + + Returns: + RGB uint8 visualization ``[H,W,3]``. + """ + colored = matplotlib.colormaps[cmap]((segmentation % 20) / 20)[..., :3] + colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8)) + return colored + + +def colorize_normal(normal: np.ndarray, mask: Optional[np.ndarray] = None) -> np.ndarray: + """Map camera-space unit normals to conventional RGB colors. + + Args: + normal: Normal map ``[H,W,3]`` with components near ``[-1,1]``. + mask: Optional boolean validity mask ``[H,W]``. + + Returns: + RGB uint8 normal visualization ``[H,W,3]``. + """ + if mask is not None: + normal = np.where(mask[..., None], normal, 0) + normal = normal * [0.5, -0.5, -0.5] + 0.5 + normal = (normal.clip(0, 1) * 255).astype(np.uint8) + return normal + + +def colorize_error_map(error_map: np.ndarray, mask: Optional[np.ndarray] = None, cmap: str = 'plasma', value_range: Optional[Tuple[float, float]] = None) -> np.ndarray: + """Colorize a scalar error map over an explicit or observed value range. + + Args: + error_map: Scalar error array ``[H,W]``. + mask: Optional boolean validity mask ``[H,W]``. + cmap: Matplotlib colormap name. + value_range: Optional ``(minimum,maximum)`` normalization bounds. + + Returns: + RGB uint8 error visualization ``[H,W,3]``. + """ + vmin, vmax = value_range if value_range is not None else (np.nanmin(error_map), np.nanmax(error_map)) + cmap = matplotlib.colormaps[cmap] + colorized_error_map = cmap(((error_map - vmin) / (vmax - vmin)).clip(0, 1))[..., :3] + if mask is not None: + colorized_error_map = np.where(mask[..., None], colorized_error_map, 0) + colorized_error_map = np.ascontiguousarray((colorized_error_map.clip(0, 1) * 255).astype(np.uint8)) + return colorized_error_map diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..d22181dd1b606cb1ccfc3ab1e8cdecb3ebbf6c43 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,13 @@ +torch>=2.8,<2.12 +torchvision>=0.23,<0.27 +gradio>=6.22,<7 +spaces +numpy>=1.26 +einops>=0.8 +huggingface_hub>=0.28 +matplotlib>=3.8 +pillow>=10.0 +scipy>=1.12 +trimesh>=4.5 +git+https://github.com/EasternJournalist/utils3d.git@3fab839f0be9931dac7c8488eb0e1600c236e183 +git+https://github.com/microsoft/MoGe.git@42acd8f46e974f5e0548ecd72315d1d7df2cb6f4