From cdff588c71c34045ad3412203231076bb2ddc8aa Mon Sep 17 00:00:00 2001 From: Farbod Motlagh Date: Wed, 29 Jul 2026 12:48:17 -0700 Subject: [PATCH] feat(camera_viz): calibrated lens undistortion (cameras..calib) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Point a camera at its calibration JSON and the feed is undistorted on the GPU before display. A per-eye sampling LUT is precomputed once from the intrinsics + distortion (+ optional stereo rectification, folded in so pairs come out row-aligned); each frame is then a single bilinear remap kernel on the producer's CUDA stream. Kannala-Brandt fisheye and Brown-Conrady models; the gr00t shw5g ChArUco JSON is accepted verbatim; buffers materialize lazily on the frame's device (multi-GPU hosts). The remap targets the projection matching the display shape — quad → rectilinear, cylinder → cylindrical, equirect → equirectangular — so the shown image is angle-correct end to end. A calibrated camera defaults to shape: cylinder (rectilinear stretches badly at wide FOV and is refused beyond 160°; cylindrical keeps the full FOV at uniform angular resolution — the gr00t WebXR client's in-shader undistort is rectilinear-only). Output FOV is derived from the calibration (Newton inversion of the fisheye polynomial at the image borders) and the layer geometry follows automatically: cylinder arc + aspect = the camera's true FOV, quads sized to the FOV at the configured distance — 1:1 visual angle with no manual tuning; undistort.hfov_deg/vfov_deg/ out_width/out_height override. UndistortSource wraps any FrameSource (all camera types + RTP), so the remap runs exactly once per new frame. 10 new tests (LUT math on CPU, GPU remap round-trip); validated live against CloudXR with the real shw5g stereo calibration (auto FOV 127.4x101.8 deg, zero pipeline loss at 2560x1984 stereo). Signed-off-by: Farbod Motlagh --- docs/source/references/camera_streaming.rst | 32 ++ examples/camera_viz/README.md | 2 + examples/camera_viz/camera_viz.py | 173 +++++-- examples/camera_viz/configs/zed.yaml | 3 + examples/camera_viz/pipeline/__init__.py | 17 +- examples/camera_viz/pipeline/undistort.py | 524 ++++++++++++++++++++ examples/camera_viz/tests/test_undistort.py | 218 ++++++++ 7 files changed, 940 insertions(+), 29 deletions(-) create mode 100644 examples/camera_viz/pipeline/undistort.py create mode 100644 examples/camera_viz/tests/test_undistort.py diff --git a/docs/source/references/camera_streaming.rst b/docs/source/references/camera_streaming.rst index 3a1ea230c..e07406996 100644 --- a/docs/source/references/camera_streaming.rst +++ b/docs/source/references/camera_streaming.rst @@ -216,6 +216,38 @@ CloudXR stream them efficiently (see :ref:`OpenXR composition layers `). For flat planes only, ``compositor: televiz`` opts a camera back into Televiz's built-in compositor. +Lens undistortion +----------------- + +Give a camera its calibration and the feed is undistorted on the GPU before display — +one precomputed sampling LUT per eye, one bilinear remap per frame: + +.. code-block:: yaml + + cameras: + - name: head + type: zed + stereo: true + calib: calib/head_fisheye.json # relative to this YAML + # undistort: # optional overrides (defaults from the calibration) + # hfov_deg: 120 # crop the output field of view + # vfov_deg: 90 + # out_width: 1920 # output resolution (default: source size) + # out_height: 1080 + +The calibration JSON carries ``model`` (``fisheye`` for Kannala-Brandt, ``brown``/``pinhole`` +for Brown-Conrady), ``image_size``, and per-eye ``K``/``D`` (plus optional stereo-rectification +``R_rect``/``R_rect_inv``, folded into the LUT so stereo pairs come out row-aligned) — the +gr00t shw5g ChArUco format is accepted verbatim; top-level ``K``/``D`` works for mono. + +The remap targets the projection that matches the display shape, so the shown image is +angle-correct end to end: ``quad`` → rectilinear, ``cylinder`` → cylindrical, ``equirect`` → +equirectangular. A calibrated camera defaults to ``shape: cylinder`` — rectilinear output +stretches badly at wide FOV (and is refused beyond 160°), while the cylinder shows the full +FOV at uniform angular resolution. Layer geometry is derived from the calibration +automatically (arc = the camera's true FOV, quads sized to the FOV at the configured +distance), so calibrated feeds render at a 1:1 visual angle with no manual tuning. + CloudXR runtime flags --------------------- diff --git a/examples/camera_viz/README.md b/examples/camera_viz/README.md index 8a34b1db8..d62626491 100644 --- a/examples/camera_viz/README.md +++ b/examples/camera_viz/README.md @@ -106,6 +106,8 @@ cameras: bitrate_mbps: 15 # gop: 150 # default fps*5 # gpu_id: 0 # multi-GPU pin + # calib: calib.json # lens calibration (per-eye K/D fisheye|brown); + # enables GPU undistortion + auto layer geometry display: # camera_viz only mode: xr | window # default: xr diff --git a/examples/camera_viz/camera_viz.py b/examples/camera_viz/camera_viz.py index f6817b604..0d9831004 100755 --- a/examples/camera_viz/camera_viz.py +++ b/examples/camera_viz/camera_viz.py @@ -31,7 +31,14 @@ import isaacteleop.viz as viz from isaacteleop.cloudxr import CloudXRLauncher -from pipeline import FrameSource, VizRunner +from pipeline import ( + FrameSource, + LensCalibration, + UndistortSettings, + UndistortSource, + VizRunner, + resolve_settings, +) from placements import PlacementConfig, PlacementStrategy, build as build_placement from sources import ( PairedFrameSource, @@ -59,6 +66,10 @@ class SourceEntry: # Cylinder shape parameters (display.placements.). cylinder_radius_m: float = 2.0 cylinder_angle_deg: float = 90.0 + # Lens undistortion (cameras..calib) — when set, the source is + # wrapped in an UndistortSource and the layer geometry is derived + # from the calibrated FOV. + undistort: Optional[UndistortSettings] = None _VALID_SHAPES = ("quad", "cylinder", "equirect") @@ -104,14 +115,18 @@ def _warn_unknown_placement_keys(cam_name: str, pspec: dict) -> None: ) -def _shape_for(cam_name: str, placements_cfg: dict) -> Tuple[str, bool, float, float]: +def _shape_for( + cam_name: str, placements_cfg: dict, has_calib: bool = False +) -> Tuple[str, str, float, float]: """Per-camera surface config from ``display.placements.``: - ``shape`` (quad | cylinder | equirect, default quad), ``compositor`` - (openxr — the default — or televiz; quads only), ``cylinder_radius_m`` - / ``cylinder_angle_deg`` (cylinder only).""" + ``shape`` (quad | cylinder | equirect; default quad, or cylinder when + the camera has a ``calib:`` — the wide-FOV-correct pairing), + ``compositor`` (openxr — the default — or televiz; quads only), + ``cylinder_radius_m`` / ``cylinder_angle_deg`` (cylinder only).""" pspec = placements_cfg.get(cam_name) or {} _warn_unknown_placement_keys(cam_name, pspec) - shape = str(pspec.get("shape", "quad")).lower() + default_shape = "cylinder" if has_calib else "quad" + shape = str(pspec.get("shape", default_shape)).lower() if shape not in _VALID_SHAPES: raise ValueError( f"camera_viz: placements.{cam_name}.shape must be one of " @@ -134,6 +149,53 @@ def _shape_for(cam_name: str, placements_cfg: dict) -> Tuple[str, bool, float, f return shape, compositor, radius_m, angle_deg +def _resolve_calib_paths(cfg: dict, base_dir: Path) -> None: + """Anchor relative ``calib:`` values to the YAML's directory (in + place), mirroring resolve_video_paths.""" + for cam in cfg.get("cameras", []): + if "calib" in cam: + p = Path(str(cam["calib"])).expanduser() + if not p.is_absolute(): + p = base_dir / p + cam["calib"] = str(p) + + +# shape -> remap target projection: the pairing that keeps the displayed +# image angle-correct (see pipeline/undistort.py). +_SHAPE_PROJECTION = { + "quad": "rectilinear", + "cylinder": "cylindrical", + "equirect": "equirect", +} + + +def _undistort_for( + cam: dict, shape: str, stereo: bool +) -> Optional[Tuple[LensCalibration, UndistortSettings]]: + """Load ``cameras..calib`` and resolve the remap settings for the + chosen display shape. Returns None when the camera has no calib.""" + if "calib" not in cam: + return None + calib = LensCalibration.load(cam["calib"]) + if stereo and calib.right is None: + print( + f"camera_viz: warning: {cam['name']}: stereo camera but calibration " + f"{cam['calib']} has no right eye — applying the left model to both eyes", + file=sys.stderr, + flush=True, + ) + overrides = dict(cam.get("undistort") or {}) + settings = resolve_settings(calib, _SHAPE_PROJECTION[shape], overrides) + print( + f"camera_viz: {cam['name']}: undistort {calib.model} -> {settings.projection} " + f"{settings.out_width}x{settings.out_height}, " + f"fov {math.degrees(settings.hfov_rad):.1f}° x " + f"{math.degrees(settings.vfov_rad):.1f}°", + flush=True, + ) + return calib, settings + + _VALID_LOCK_MODES = ("world", "head", "lazy", "gimbal") @@ -178,16 +240,28 @@ def _enabled_cameras(cfg: dict) -> List[dict]: def _placement_with_aspect( - spec: Optional[dict], width: int, height: int, is_xr: bool + spec: Optional[dict], + width: int, + height: int, + is_xr: bool, + undistort: Optional[UndistortSettings] = None, ) -> Optional[PlacementStrategy]: """Build the placement, filling in ``size`` from the source's aspect ratio when the YAML doesn't pin it. Width defaults to 1.0 m so a - 16:9 source lands at 1.0 x 0.5625, a 3.55:1 SBS at 1.0 x 0.281.""" + 16:9 source lands at 1.0 x 0.5625, a 3.55:1 SBS at 1.0 x 0.281. + Calibrated quads instead get the ANGULAR size at the configured + distance (width = 2 d tan(hfov/2)), so the rectilinear image spans + exactly the field of view the camera captured.""" if spec is not None and "size" not in spec: - spec = { - **spec, - "size": [_DEFAULT_PLANE_WIDTH_M, _DEFAULT_PLANE_WIDTH_M * height / width], - } + if undistort is not None: + d = float(spec.get("distance", PlacementConfig.distance)) + size = [ + 2.0 * d * math.tan(undistort.hfov_rad / 2.0), + 2.0 * d * math.tan(undistort.vfov_rad / 2.0), + ] + else: + size = [_DEFAULT_PLANE_WIDTH_M, _DEFAULT_PLANE_WIDTH_M * height / width] + spec = {**spec, "size": size} return _build_placement(spec, is_xr) @@ -205,14 +279,26 @@ def _build_local_entries(cfg: dict, is_xr: bool) -> List[SourceEntry]: entries: List[SourceEntry] = [] for cam in _enabled_cameras(cfg): cam_sources = build_local_camera(cam) + stereo, baseline_mm = _stereo_for(cam, placements_cfg) + shape, compositor, radius_m, angle_deg = _shape_for( + cam["name"], placements_cfg, has_calib="calib" in cam + ) + und = _undistort_for(cam, shape, stereo) + settings = None + if und is not None: + calib, settings = und + cam_sources = [UndistortSource(src, calib, settings) for src in cam_sources] # Aspect comes from the built source's spec, not the YAML — video - # sources may omit width/height and size themselves from the file. + # sources may omit width/height and size themselves from the file, + # and undistortion may change the output dimensions. first = cam_sources[0].spec placement = _placement_with_aspect( - placements_cfg.get(cam["name"]), first.width, first.height, is_xr + placements_cfg.get(cam["name"]), + first.width, + first.height, + is_xr, + undistort=settings, ) - stereo, baseline_mm = _stereo_for(cam, placements_cfg) - shape, compositor, radius_m, angle_deg = _shape_for(cam["name"], placements_cfg) for source in cam_sources: entries.append( SourceEntry( @@ -224,6 +310,7 @@ def _build_local_entries(cfg: dict, is_xr: bool) -> List[SourceEntry]: compositor=compositor, cylinder_radius_m=radius_m, cylinder_angle_deg=angle_deg, + undistort=settings, ) ) return entries @@ -247,12 +334,6 @@ def _build_rtp_entries(cfg: dict, is_xr: bool) -> List[SourceEntry]: "width/height when source: rtp — the receiver sizes its " "decoder from the YAML, not from the wire" ) - placement = _placement_with_aspect( - placements_cfg.get(cam["name"]), - int(cam["width"]), - int(cam["height"]), - is_xr, - ) stereo, baseline_mm = _stereo_for(cam, placements_cfg) if stereo: @@ -290,7 +371,21 @@ def _build_rtp_entries(cfg: dict, is_xr: bool) -> List[SourceEntry]: gpu_id=int(rtp.get("gpu_id", 0)), ) - shape, compositor, radius_m, angle_deg = _shape_for(cam["name"], placements_cfg) + shape, compositor, radius_m, angle_deg = _shape_for( + cam["name"], placements_cfg, has_calib="calib" in cam + ) + und = _undistort_for(cam, shape, stereo) + settings = None + if und is not None: + calib, settings = und + source = UndistortSource(source, calib, settings) + placement = _placement_with_aspect( + placements_cfg.get(cam["name"]), + source.spec.width, + source.spec.height, + is_xr, + undistort=settings, + ) entries.append( SourceEntry( source=source, @@ -301,6 +396,7 @@ def _build_rtp_entries(cfg: dict, is_xr: bool) -> List[SourceEntry]: compositor=compositor, cylinder_radius_m=radius_m, cylinder_angle_deg=angle_deg, + undistort=settings, ) ) return entries @@ -345,17 +441,29 @@ def _add_layer(session: viz.VizSession, entry: SourceEntry): to be an equirect panorama). Runtime-composited always. """ spec = entry.source.spec + und = entry.undistort if entry.shape == "cylinder": layer_cfg = viz.CylinderLayerConfig() layer_cfg.name = spec.name layer_cfg.resolution = viz.Resolution(spec.width, spec.height) layer_cfg.stereo = entry.stereo layer_cfg.stereo_baseline_mm = entry.stereo_baseline_mm - # aspect_ratio 0 = derived from the source resolution (square texels). - layer_cfg.placement = viz.CylinderLayerPlacement( - radius_m=entry.cylinder_radius_m, - central_angle_rad=math.radians(entry.cylinder_angle_deg), - ) + if und is not None: + # Calibrated: the remapped image is equal-angle horizontally + # over hfov and planar vertically over 2 tan(vfov/2) — set the + # arc and its aspect so every texel lands at the exact visual + # angle the camera captured (1:1 angular mapping). + layer_cfg.placement = viz.CylinderLayerPlacement( + radius_m=entry.cylinder_radius_m, + central_angle_rad=und.hfov_rad, + aspect_ratio=und.hfov_rad / (2.0 * math.tan(und.vfov_rad / 2.0)), + ) + else: + # aspect_ratio 0 = derived from the source resolution. + layer_cfg.placement = viz.CylinderLayerPlacement( + radius_m=entry.cylinder_radius_m, + central_angle_rad=math.radians(entry.cylinder_angle_deg), + ) return session.add_cylinder_layer(layer_cfg) if entry.shape == "equirect": layer_cfg = viz.EquirectLayerConfig() @@ -365,6 +473,14 @@ def _add_layer(session: viz.VizSession, entry: SourceEntry): # Baseline only matters at finite sphere radius; harmless at the # default infinite-radius placement (full 360x180 sphere). layer_cfg.stereo_baseline_mm = entry.stereo_baseline_mm + if und is not None: + # Calibrated: span exactly the remapped FOV instead of a full + # sphere, centered on the horizon. + layer_cfg.placement = viz.EquirectLayerPlacement( + central_horizontal_angle_rad=und.hfov_rad, + upper_vertical_angle_rad=und.vfov_rad / 2.0, + lower_vertical_angle_rad=-und.vfov_rad / 2.0, + ) return session.add_equirect_layer(layer_cfg) layer_cfg = viz.QuadLayerConfig() @@ -405,6 +521,7 @@ def main(argv: Optional[list[str]] = None) -> int: # Top-level ``verbose:`` enables per-source periodic breadcrumbs. set_verbose(bool(cfg.get("verbose", False))) resolve_video_paths(cfg, args.config.parent) + _resolve_calib_paths(cfg, args.config.parent) source_mode = cfg.get("source", "local").lower() if source_mode not in ("local", "rtp"): diff --git a/examples/camera_viz/configs/zed.yaml b/examples/camera_viz/configs/zed.yaml index eb5d97335..55ff5221a 100644 --- a/examples/camera_viz/configs/zed.yaml +++ b/examples/camera_viz/configs/zed.yaml @@ -25,6 +25,9 @@ cameras: height: 720 fps: 60 bus_type: usb # usb | gmsl + # calib: calib/zed_fisheye.json # per-eye K/D (+R_rect) calibration: + # # undistorts on the GPU and derives the + # # display geometry from the true FOV serial_number: 0 # 0 → first available stereo: true rtp: diff --git a/examples/camera_viz/pipeline/__init__.py b/examples/camera_viz/pipeline/__init__.py index 114817b57..b4bb4f4cc 100644 --- a/examples/camera_viz/pipeline/__init__.py +++ b/examples/camera_viz/pipeline/__init__.py @@ -9,10 +9,25 @@ """ from .interface import Frame, FrameSource, SourceSpec +from .undistort import ( + LensCalibration, + UndistortSettings, + UndistortSource, + resolve_settings, +) # VizRunner is exposed via a lazy attribute so importing ``pipeline`` # doesn't drag in ``isaacteleop.viz`` for sender-side code paths. -__all__ = ["Frame", "FrameSource", "SourceSpec", "VizRunner"] +__all__ = [ + "Frame", + "FrameSource", + "LensCalibration", + "SourceSpec", + "UndistortSettings", + "UndistortSource", + "VizRunner", + "resolve_settings", +] def __getattr__(name: str): diff --git a/examples/camera_viz/pipeline/undistort.py b/examples/camera_viz/pipeline/undistort.py new file mode 100644 index 000000000..b76227cff --- /dev/null +++ b/examples/camera_viz/pipeline/undistort.py @@ -0,0 +1,524 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Lens undistortion stage: calibrated remap into a display-matched projection. + +Pipeline: a per-eye sampling LUT is precomputed ONCE on the CPU from the +camera calibration (intrinsics + distortion + optional stereo +rectification), then every camera frame is remapped on the GPU with a +single bilinear gather before it reaches the layer. The remap target +projection is chosen to match the display surface, so the shown image is +angle-correct end to end: + + ========== =================== ============================= + shape remap projection display parameterization + ========== =================== ============================= + quad rectilinear planar (pinhole) + cylinder cylindrical equal-angle x, planar y + equirect equirectangular equal-angle x and y + ========== =================== ============================= + +Rectilinear stretches toward the edges and cannot reach 180 degrees; +for wide-FOV fisheye heads (Tianji shw5g ~130 degrees) the cylindrical +target keeps the full FOV at uniform angular resolution. + +Calibration JSON schema (the gr00t shw5g ChArUco format is accepted +verbatim): top-level ``model`` ("fisheye" for Kannala-Brandt theta +polynomial, "brown"/"pinhole" for Brown-Conrady), ``image_size`` +[W, H], and either per-eye ``left``/``right`` blocks or top-level +``K``/``D`` for mono. Each eye block: ``K`` (3x3), ``D`` (4 fisheye or +4/5/8 Brown coefficients), optional ``R_rect_inv``/``R_rect`` (stereo +rectification rotation; folded into the LUT so stereo pairs come out +row-aligned). + +Coordinates follow OpenCV camera conventions (x right, y down, +z forward). +""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Tuple + +import numpy as np + +from .interface import Frame, FrameSource, SourceSpec + +PROJECTIONS = ("rectilinear", "cylindrical", "equirect") + + +# ---------------------------------------------------------------------- +# Calibration parsing +# ---------------------------------------------------------------------- + + +@dataclass(frozen=True) +class EyeCalibration: + K: np.ndarray # (3, 3) float64 + D: np.ndarray # distortion coefficients, model-dependent length + R_cam_from_rect: np.ndarray # (3, 3): rectified-frame ray -> camera-frame ray + + +@dataclass(frozen=True) +class LensCalibration: + model: str # "fisheye" | "brown" + image_size: Tuple[int, int] # (W, H) + left: EyeCalibration + right: Optional[EyeCalibration] + + @staticmethod + def load(path: str | Path) -> "LensCalibration": + with open(path) as f: + data = json.load(f) + model = str(data.get("model", "fisheye")).lower() + if model in ("pinhole", "brown", "brown-conrady", "plumb_bob"): + model = "brown" + elif model in ("fisheye", "kannala-brandt", "kb4", "equidistant"): + model = "fisheye" + else: + raise ValueError( + f"undistort: unknown calibration model {model!r} in {path}" + ) + if "image_size" not in data: + raise ValueError(f"undistort: calibration {path} missing image_size [W, H]") + w, h = (int(v) for v in data["image_size"]) + + def eye(block: dict) -> EyeCalibration: + K = np.asarray(block["K"], dtype=np.float64).reshape(3, 3) + D = np.asarray(block.get("D", []), dtype=np.float64).reshape(-1) + if "R_rect_inv" in block: + R = np.asarray(block["R_rect_inv"], dtype=np.float64).reshape(3, 3) + elif "R_rect" in block: + # R_rect maps camera rays -> rectified rays (OpenCV + # convention); the LUT walks the other way. + R = np.asarray(block["R_rect"], dtype=np.float64).reshape(3, 3).T + else: + R = np.eye(3) + return EyeCalibration(K=K, D=D, R_cam_from_rect=R) + + if "left" in data: + left = eye(data["left"]) + right = eye(data["right"]) if "right" in data else None + elif "K" in data: + left = eye(data) + right = None + else: + raise ValueError( + f"undistort: calibration {path} has neither left/right blocks nor top-level K" + ) + return LensCalibration(model=model, image_size=(w, h), left=left, right=right) + + +# ---------------------------------------------------------------------- +# Lens model: camera-frame rays -> source pixel coordinates +# ---------------------------------------------------------------------- + + +def _project_fisheye(dirs: np.ndarray, K: np.ndarray, D: np.ndarray) -> np.ndarray: + """Kannala-Brandt (OpenCV fisheye): theta_d = theta(1 + k1 t^2 + ... + k4 t^8).""" + x, y, z = dirs[..., 0], dirs[..., 1], dirs[..., 2] + r_xy = np.sqrt(x * x + y * y) + theta = np.arctan2(r_xy, z) + k = np.zeros(4) + k[: min(4, D.size)] = D[:4] + t2 = theta * theta + theta_d = theta * (1.0 + k[0] * t2 + k[1] * t2**2 + k[2] * t2**3 + k[3] * t2**4) + with np.errstate(invalid="ignore", divide="ignore"): + scale = np.where(r_xy > 1e-9, theta_d / r_xy, 1.0) + xd, yd = x * scale, y * scale + u = K[0, 0] * (xd + (K[0, 1] / K[0, 0]) * yd) + K[0, 2] + v = K[1, 1] * yd + K[1, 2] + # theta >= pi/2 is behind-the-lens for practical FOVs derived below; + # keep whatever the polynomial produced — out-of-image lands black. + return np.stack([u, v], axis=-1) + + +def _project_brown(dirs: np.ndarray, K: np.ndarray, D: np.ndarray) -> np.ndarray: + """Brown-Conrady with OpenCV coefficient order (k1 k2 p1 p2 [k3 [k4 k5 k6]]).""" + x, y, z = dirs[..., 0], dirs[..., 1], dirs[..., 2] + valid = z > 1e-6 + zs = np.where(valid, z, 1.0) + xp, yp = x / zs, y / zs + d = np.zeros(8) + d[: min(8, D.size)] = D[:8] + k1, k2, p1, p2, k3, k4, k5, k6 = d + r2 = xp * xp + yp * yp + radial = (1.0 + k1 * r2 + k2 * r2**2 + k3 * r2**3) / ( + 1.0 + k4 * r2 + k5 * r2**2 + k6 * r2**3 + ) + xd = xp * radial + 2.0 * p1 * xp * yp + p2 * (r2 + 2.0 * xp * xp) + yd = yp * radial + p1 * (r2 + 2.0 * yp * yp) + 2.0 * p2 * xp * yp + u = K[0, 0] * xd + K[0, 1] * yd + K[0, 2] + v = K[1, 1] * yd + K[1, 2] + # Rays at/behind the image plane can't be sampled: push out of range. + u = np.where(valid, u, -1e9) + v = np.where(valid, v, -1e9) + return np.stack([u, v], axis=-1) + + +def _fisheye_theta_from_radius(r_norm: float, D: np.ndarray) -> float: + """Invert theta_d(theta) = r for the KB polynomial (Newton, monotonic range).""" + k = np.zeros(4) + k[: min(4, D.size)] = D[:4] + theta = min(r_norm, math.pi / 2.0) # r is a decent seed for small distortion + for _ in range(20): + t2 = theta * theta + f = ( + theta * (1.0 + k[0] * t2 + k[1] * t2**2 + k[2] * t2**3 + k[3] * t2**4) + - r_norm + ) + fp = ( + 1.0 + + 3.0 * k[0] * t2 + + 5.0 * k[1] * t2**2 + + 7.0 * k[2] * t2**3 + + 9.0 * k[3] * t2**4 + ) + step = f / fp + theta -= step + if abs(step) < 1e-10: + break + return max(theta, 0.0) + + +def derive_fov(calib: LensCalibration, eye: EyeCalibration) -> Tuple[float, float]: + """(hfov, vfov) in radians covered by the source image, symmetric about + the optical axis (min of the two half-angles per axis, so a symmetric + output grid never samples outside the image).""" + w, h = calib.image_size + fx, fy = eye.K[0, 0], eye.K[1, 1] + cx, cy = eye.K[0, 2], eye.K[1, 2] + half_x_px = min(cx, (w - 1) - cx) + half_y_px = min(cy, (h - 1) - cy) + if calib.model == "fisheye": + half_h = _fisheye_theta_from_radius(half_x_px / fx, eye.D) + half_v = _fisheye_theta_from_radius(half_y_px / fy, eye.D) + else: + # Brown distortion is small at typical pinhole FOVs; the ideal + # angle is a good bound (out-of-image rays land black anyway). + half_h = math.atan(half_x_px / fx) + half_v = math.atan(half_y_px / fy) + return 2.0 * half_h, 2.0 * half_v + + +# ---------------------------------------------------------------------- +# Output-projection ray grids +# ---------------------------------------------------------------------- + + +def _ray_grid( + projection: str, out_w: int, out_h: int, hfov: float, vfov: float +) -> np.ndarray: + """(out_h, out_w, 3) unit-scale ray directions in the RECTIFIED frame, + OpenCV convention (x right, y down, z forward), pixel centers.""" + u = (np.arange(out_w, dtype=np.float64) + 0.5) / out_w # 0..1 left->right + v = (np.arange(out_h, dtype=np.float64) + 0.5) / out_h # 0..1 top->bottom + uu, vv = np.meshgrid(u, v) + if projection == "rectilinear": + x = (uu - 0.5) * 2.0 * math.tan(hfov / 2.0) + y = (vv - 0.5) * 2.0 * math.tan(vfov / 2.0) + z = np.ones_like(x) + elif projection == "cylindrical": + # Equal-angle horizontally, planar vertically: matches how an + # XrCompositionLayerCylinderKHR of central_angle == hfov samples + # its texture (height h at azimuth psi -> direction (sin, h, cos)). + psi = (uu - 0.5) * hfov + x = np.sin(psi) + y = (vv - 0.5) * 2.0 * math.tan(vfov / 2.0) + z = np.cos(psi) + elif projection == "equirect": + # Equal-angle both axes: matches XrCompositionLayerEquirect2KHR + # with central_horizontal_angle == hfov, vertical span == vfov. + psi = (uu - 0.5) * hfov + phi = (vv - 0.5) * vfov # positive = down (OpenCV y-down) + x = np.cos(phi) * np.sin(psi) + y = np.sin(phi) + z = np.cos(phi) * np.cos(psi) + else: + raise ValueError(f"undistort: unknown projection {projection!r}") + return np.stack([x, y, z], axis=-1) + + +def build_maps( + calib: LensCalibration, + eye: EyeCalibration, + projection: str, + out_w: int, + out_h: int, + hfov: float, + vfov: float, +) -> Tuple[np.ndarray, np.ndarray]: + """Sampling LUT: for each output pixel, the (sub-pixel) source + coordinate to bilinearly fetch. Returns (map_x, map_y) float32 of + shape (out_h, out_w); out-of-image entries stay out of range and the + remap kernel writes opaque black there.""" + rays = _ray_grid(projection, out_w, out_h, hfov, vfov) + cam_rays = rays @ eye.R_cam_from_rect.T + if calib.model == "fisheye": + uv = _project_fisheye(cam_rays, eye.K, eye.D) + else: + uv = _project_brown(cam_rays, eye.K, eye.D) + return ( + np.ascontiguousarray(uv[..., 0], dtype=np.float32), + np.ascontiguousarray(uv[..., 1], dtype=np.float32), + ) + + +# ---------------------------------------------------------------------- +# GPU remap (CuPy raw kernel, bilinear RGBA8) +# ---------------------------------------------------------------------- + +_REMAP_KERNEL_SRC = r""" +extern "C" __global__ void remap_rgba8( + const unsigned char* __restrict__ src, int src_w, int src_h, long long src_pitch, + const float* __restrict__ map_x, const float* __restrict__ map_y, + unsigned char* __restrict__ dst, int dst_w, int dst_h) +{ + int x = blockIdx.x * blockDim.x + threadIdx.x; + int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= dst_w || y >= dst_h) return; + long long o = ((long long)y * dst_w + x) * 4; + float sx = map_x[(long long)y * dst_w + x]; + float sy = map_y[(long long)y * dst_w + x]; + if (!(sx >= 0.0f && sy >= 0.0f && sx <= (float)(src_w - 1) && sy <= (float)(src_h - 1))) { + dst[o] = 0; dst[o + 1] = 0; dst[o + 2] = 0; dst[o + 3] = 255; + return; + } + int x0 = (int)sx, y0 = (int)sy; + int x1 = min(x0 + 1, src_w - 1), y1 = min(y0 + 1, src_h - 1); + float fx = sx - (float)x0, fy = sy - (float)y0; + const unsigned char* p00 = src + (long long)y0 * src_pitch + (long long)x0 * 4; + const unsigned char* p01 = src + (long long)y0 * src_pitch + (long long)x1 * 4; + const unsigned char* p10 = src + (long long)y1 * src_pitch + (long long)x0 * 4; + const unsigned char* p11 = src + (long long)y1 * src_pitch + (long long)x1 * 4; + #pragma unroll + for (int c = 0; c < 3; ++c) { + float top = (float)p00[c] + fx * ((float)p01[c] - (float)p00[c]); + float bot = (float)p10[c] + fx * ((float)p11[c] - (float)p10[c]); + dst[o + c] = (unsigned char)(top + fy * (bot - top) + 0.5f); + } + dst[o + 3] = 255; +} +""" + + +class GpuRemapper: + """One eye's persistent LUT + output buffer; remap() is a single + kernel launch on the producer's CUDA stream.""" + + def __init__( + self, map_x: np.ndarray, map_y: np.ndarray, src_size: Tuple[int, int] + ) -> None: + import cupy as cp + + self._cp = cp + self._map_x_np = map_x + self._map_y_np = map_y + self._src_w, self._src_h = src_size + self._out_h, self._out_w = map_x.shape + # GPU buffers materialize lazily on the DEVICE the frames arrive + # on (multi-GPU hosts: the viz session may live on a different + # GPU than the capture source; launching a kernel with + # cross-device pointers is an illegal access). + self._device_id: Optional[int] = None + self._kernel = None + self._map_x = None + self._map_y = None + self._out = None + + def _ensure_device(self, device_id: int) -> None: + if self._device_id == device_id: + return + cp = self._cp + with cp.cuda.Device(device_id): + self._kernel = cp.RawKernel(_REMAP_KERNEL_SRC, "remap_rgba8") + self._map_x = cp.asarray(self._map_x_np) + self._map_y = cp.asarray(self._map_y_np) + self._out = cp.empty((self._out_h, self._out_w, 4), dtype=cp.uint8) + self._device_id = device_id + + def remap(self, image, stream: int = 0): + """image: HxWx4 uint8 __cuda_array_interface__ array (C-contiguous + rows; pitch honored). Returns the persistent output CuPy array — + valid until the next remap() call for this eye.""" + cp = self._cp + src = cp.asarray(image) + if ( + src.shape[0] != self._src_h + or src.shape[1] != self._src_w + or src.shape[2] != 4 + ): + raise ValueError( + f"undistort: frame {src.shape[1]}x{src.shape[0]} does not match " + f"calibration image_size {self._src_w}x{self._src_h}" + ) + self._ensure_device(src.device.id) + pitch = src.strides[0] + block = (16, 16, 1) + grid = ( + (self._out_w + block[0] - 1) // block[0], + (self._out_h + block[1] - 1) // block[1], + 1, + ) + with cp.cuda.Device(src.device.id): + ctx = cp.cuda.ExternalStream(stream) if stream else cp.cuda.Stream.null + with ctx: + self._kernel( + grid, + block, + ( + src, + np.int32(self._src_w), + np.int32(self._src_h), + np.int64(pitch), + self._map_x, + self._map_y, + self._out, + np.int32(self._out_w), + np.int32(self._out_h), + ), + ) + return self._out + + +# ---------------------------------------------------------------------- +# FrameSource wrapper +# ---------------------------------------------------------------------- + + +@dataclass(frozen=True) +class UndistortSettings: + """Resolved undistort parameters for one camera (see camera_viz's + ``calib:`` / ``undistort:`` YAML keys).""" + + calib_path: str + projection: str # one of PROJECTIONS + out_width: int + out_height: int + hfov_rad: float + vfov_rad: float + + +def resolve_settings( + calib: LensCalibration, + projection: str, + overrides: dict, +) -> UndistortSettings: + """Fill FOV / output size from the calibration unless overridden. + FOV comes from the LEFT eye (stereo rigs are near-identical and a + shared output grid keeps the pair rectified).""" + if projection not in PROJECTIONS: + raise ValueError( + f"undistort: projection must be one of {'|'.join(PROJECTIONS)}, got {projection!r}" + ) + auto_h, auto_v = derive_fov(calib, calib.left) + hfov = ( + math.radians(float(overrides["hfov_deg"])) + if "hfov_deg" in overrides + else auto_h + ) + vfov = ( + math.radians(float(overrides["vfov_deg"])) + if "vfov_deg" in overrides + else auto_v + ) + if projection == "rectilinear": + # tan() blows up approaching 180 degrees; refuse silly outputs. + limit = math.radians(160.0) + if hfov >= limit or vfov >= limit: + raise ValueError( + "undistort: rectilinear output beyond 160 degrees is degenerate — " + "use shape: cylinder (or equirect) for wide-FOV lenses, or pass " + "undistort.hfov_deg / vfov_deg to crop" + ) + src_w, src_h = calib.image_size + out_w = int(overrides.get("out_width", src_w)) + out_h = int(overrides.get("out_height", src_h)) + return UndistortSettings( + calib_path="", + projection=projection, + out_width=out_w, + out_height=out_h, + hfov_rad=hfov, + vfov_rad=vfov, + ) + + +class UndistortSource(FrameSource): + """Wraps a FrameSource and remaps every new frame through the + calibrated LUT (per eye for stereo). latest() preserves the inner + mailbox contract: the remap runs exactly once per new frame.""" + + def __init__( + self, inner: FrameSource, calib: LensCalibration, settings: UndistortSettings + ) -> None: + self._inner = inner + self._settings = settings + maps_l = build_maps( + calib, + calib.left, + settings.projection, + settings.out_width, + settings.out_height, + settings.hfov_rad, + settings.vfov_rad, + ) + self._left = GpuRemapper(*maps_l, calib.image_size) + self._right: Optional[GpuRemapper] = None + if calib.right is not None: + maps_r = build_maps( + calib, + calib.right, + settings.projection, + settings.out_width, + settings.out_height, + settings.hfov_rad, + settings.vfov_rad, + ) + self._right = GpuRemapper(*maps_r, calib.image_size) + self._spec = SourceSpec( + name=inner.spec.name, + width=settings.out_width, + height=settings.out_height, + pixel_format=inner.spec.pixel_format, + ) + + @property + def spec(self) -> SourceSpec: + return self._spec + + @property + def settings(self) -> UndistortSettings: + return self._settings + + def start(self) -> None: + self._inner.start() + + def stop(self) -> None: + self._inner.stop() + + def latest(self) -> Optional[Frame]: + frame = self._inner.latest() + if frame is None: + return None + image = self._left.remap(frame.image, stream=frame.stream) + right = None + if frame.image_right is not None: + # Stereo frame: right eye through its own LUT; without a + # right calibration block, fall back to the left LUT (mono + # calib on a stereo rig — better than nothing, warned at load). + remapper = self._right if self._right is not None else self._left + if remapper is self._left: + # The left remapper's output buffer would be overwritten; + # copy before reusing it for the right eye. + image = image.copy() + right = remapper.remap(frame.image_right, stream=frame.stream) + return Frame( + image=image, + timestamp_ns=frame.timestamp_ns, + source_id=frame.source_id, + stream=frame.stream, + image_right=right, + ) diff --git a/examples/camera_viz/tests/test_undistort.py b/examples/camera_viz/tests/test_undistort.py new file mode 100644 index 000000000..0aa16f3e6 --- /dev/null +++ b/examples/camera_viz/tests/test_undistort.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Undistort stage: LUT math (CPU, always) + GPU remap (skipped without CUDA).""" + +from __future__ import annotations + +import json +import math + +import numpy as np +import pytest + +from pipeline.undistort import ( + LensCalibration, + UndistortSource, + _fisheye_theta_from_radius, + _project_fisheye, + build_maps, + derive_fov, + resolve_settings, +) + + +def _fisheye_calib_dict( + w=640, h=480, f=200.0, d=(0.02, -0.005, 0.001, 0.0), stereo=False +): + K = [[f, 0.0, (w - 1) / 2.0], [0.0, f, (h - 1) / 2.0], [0.0, 0.0, 1.0]] + eye = {"K": K, "D": list(d)} + data = {"model": "fisheye", "image_size": [w, h]} + if stereo: + data["left"] = dict(eye) + data["right"] = dict(eye) + else: + data.update(eye) + return data + + +def _load(tmp_path, data) -> LensCalibration: + p = tmp_path / "calib.json" + p.write_text(json.dumps(data)) + return LensCalibration.load(p) + + +def test_parses_gr00t_schema_and_mono_fallback(tmp_path): + stereo = _load(tmp_path, _fisheye_calib_dict(stereo=True)) + assert stereo.model == "fisheye" + assert stereo.right is not None + mono = _load(tmp_path, _fisheye_calib_dict(stereo=False)) + assert mono.right is None + assert mono.left.K[0, 0] == pytest.approx(200.0) + # No R_rect -> identity rotation. + assert np.allclose(mono.left.R_cam_from_rect, np.eye(3)) + + +def test_r_rect_transposed_when_inverse_missing(tmp_path): + data = _fisheye_calib_dict(stereo=True) + # 90-degree yaw as R_rect (camera->rectified); loader must store its + # transpose (rectified->camera). + R = [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]] + data["left"]["R_rect"] = R + calib = _load(tmp_path, data) + assert np.allclose(calib.left.R_cam_from_rect, np.asarray(R).T) + + +def test_fisheye_projection_center_and_equidistant(): + K = np.array([[200.0, 0.0, 320.0], [0.0, 200.0, 240.0], [0.0, 0.0, 1.0]]) + D = np.zeros(4) + # On-axis ray -> principal point. + uv = _project_fisheye(np.array([[0.0, 0.0, 1.0]]), K, D) + assert np.allclose(uv, [[320.0, 240.0]]) + # With zero distortion the model is pure equidistant: r = f * theta. + theta = 0.5 + uv = _project_fisheye(np.array([[math.sin(theta), 0.0, math.cos(theta)]]), K, D) + assert uv[0, 0] == pytest.approx(320.0 + 200.0 * theta) + assert uv[0, 1] == pytest.approx(240.0) + + +def test_theta_solver_inverts_polynomial(): + D = np.array([0.03, -0.01, 0.002, -0.0004]) + for theta in (0.1, 0.6, 1.2): + t2 = theta * theta + r = theta * (1 + D[0] * t2 + D[1] * t2**2 + D[2] * t2**3 + D[3] * t2**4) + assert _fisheye_theta_from_radius(r, D) == pytest.approx(theta, abs=1e-8) + + +def test_derive_fov_matches_border_angle(tmp_path): + calib = _load(tmp_path, _fisheye_calib_dict(w=640, h=480, f=200.0, d=(0, 0, 0, 0))) + hfov, vfov = derive_fov(calib, calib.left) + # Pure equidistant: half-angle = half-width-pixels / f. + assert hfov == pytest.approx(2.0 * ((640 - 1) / 2.0) / 200.0) + assert vfov == pytest.approx(2.0 * ((480 - 1) / 2.0) / 200.0) + + +def test_maps_center_pixel_hits_principal_point(tmp_path): + calib = _load(tmp_path, _fisheye_calib_dict()) + for projection in ("rectilinear", "cylindrical", "equirect"): + # The test lens covers ~178 degrees; rectilinear needs a crop + # (the guard for that is tested separately below). + overrides = ( + {"hfov_deg": 100, "vfov_deg": 80} if projection == "rectilinear" else {} + ) + s = resolve_settings(calib, projection, overrides) + mx, my = build_maps( + calib, + calib.left, + projection, + s.out_width, + s.out_height, + s.hfov_rad, + s.vfov_rad, + ) + assert mx.shape == (s.out_height, s.out_width) + cy, cx = s.out_height // 2, s.out_width // 2 + # Output center looks down the optical axis -> principal point, + # to within half an output pixel of angular quantization. + assert abs(mx[cy, cx] - calib.left.K[0, 2]) < 2.0 + assert abs(my[cy, cx] - calib.left.K[1, 2]) < 2.0 + # Along the principal axes the derived FOV never samples outside + # the source (grid CORNERS may — the display shows them black, + # matching the source's own coverage limits). + mid_row_x = mx[cy, :] + mid_col_y = my[:, cx] + assert mid_row_x.min() >= -1.0 and mid_row_x.max() <= 640.0 + assert mid_col_y.min() >= -1.0 and mid_col_y.max() <= 480.0 + + +def test_cylindrical_map_is_equal_angle_horizontally(tmp_path): + # With zero distortion, equal azimuth steps in the OUTPUT must map to + # equal theta steps at the horizon row — i.e. equal pixel steps in an + # equidistant source. That's the property the CylinderLayer relies on. + calib = _load(tmp_path, _fisheye_calib_dict(d=(0, 0, 0, 0))) + s = resolve_settings(calib, "cylindrical", {}) + mx, _ = build_maps( + calib, + calib.left, + "cylindrical", + s.out_width, + s.out_height, + s.hfov_rad, + s.vfov_rad, + ) + row = mx[s.out_height // 2, :] + steps = np.diff(row) + assert steps.std() / steps.mean() < 1e-3 + + +def test_rectilinear_rejects_degenerate_fov(tmp_path): + calib = _load(tmp_path, _fisheye_calib_dict(f=80.0)) # ~4.0 rad hfov + with pytest.raises(ValueError, match="rectilinear"): + resolve_settings(calib, "rectilinear", {}) + # Cylindrical accepts the same lens. + s = resolve_settings(calib, "cylindrical", {}) + assert s.hfov_rad > math.radians(160.0) + + +def test_overrides(tmp_path): + calib = _load(tmp_path, _fisheye_calib_dict()) + s = resolve_settings( + calib, + "cylindrical", + {"hfov_deg": 90, "vfov_deg": 60, "out_width": 800, "out_height": 500}, + ) + assert s.hfov_rad == pytest.approx(math.radians(90)) + assert s.vfov_rad == pytest.approx(math.radians(60)) + assert (s.out_width, s.out_height) == (800, 500) + + +# ── GPU path ────────────────────────────────────────────────────────── + + +def _gpu_available() -> bool: + try: + import cupy as cp + + return cp.cuda.runtime.getDeviceCount() > 0 + except Exception: + return False + + +@pytest.mark.skipif(not _gpu_available(), reason="no CUDA GPU") +def test_gpu_remap_identity_and_source(tmp_path): + import cupy as cp + + from pipeline.interface import Frame, SourceSpec + + calib = _load(tmp_path, _fisheye_calib_dict(w=64, h=48, f=40.0, d=(0, 0, 0, 0))) + s = resolve_settings(calib, "cylindrical", {}) + + class OneShot: + spec = SourceSpec(name="t", width=64, height=48) + + def __init__(self): + img = cp.zeros((48, 64, 4), dtype=cp.uint8) + img[:, :, 0] = cp.arange(64, dtype=cp.uint8)[None, :] * 3 # R ramp + img[:, :, 3] = 255 + self._frame = Frame(image=img, timestamp_ns=0, source_id="t") + + def start(self): + pass + + def stop(self): + pass + + def latest(self): + f, self._frame = self._frame, None + return f + + src = UndistortSource(OneShot(), calib, s) + assert (src.spec.width, src.spec.height) == (s.out_width, s.out_height) + out = src.latest() + assert out is not None + arr = cp.asnumpy(out.image) + assert arr.shape == (s.out_height, s.out_width, 4) + assert (arr[:, :, 3] == 255).all() + # Center of a distortion-free remap shows the center of the ramp. + assert abs(int(arr[s.out_height // 2, s.out_width // 2, 0]) - 31 * 3) <= 6 + # Mailbox contract: no new inner frame -> None. + assert src.latest() is None