src/lerobot/datasets/depth_utils.py

#!/usr/bin/env python

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Depth encoding/decoding helpers for :class:`DepthEncoderConfig`.
"""

import math
from typing import Literal

import av
import numpy as np
import torch
from numpy.typing import NDArray

from lerobot.configs.video import (
    DEFAULT_DEPTH_MAX,
    DEFAULT_DEPTH_MIN,
    DEFAULT_DEPTH_PIX_FMT,
    DEFAULT_DEPTH_SHIFT,
    DEFAULT_DEPTH_USE_LOG,
    DEPTH_METER_UNIT,
    DEPTH_MILLIMETER_UNIT,
    DEPTH_QMAX,
    infer_depth_unit,
)

from .image_writer import squeeze_single_channel
from .pyav_utils import write_u16_plane

MM_PER_METRE = 1000.0
_UINT16_MAX = 65535


def _validate_log_quant_params(depth_min: float, shift: float) -> None:
    """Ensure ``log(depth_min + shift)`` is finite."""
    if depth_min + shift <= 0:
        raise ValueError(
            f"depth_min + shift must be positive for logarithmic quantization, "
            f"got depth_min={depth_min} + shift={shift} = {depth_min + shift}"
        )


def _depth_input_to_float32_and_unit(
    depth: NDArray[np.integer] | NDArray[np.floating],
    input_unit: Literal["auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT],
) -> tuple[NDArray[np.float32], Literal[DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT]]:
    """Convert depth to float32 in the chosen unit, and return the resolved unit."""
    resolved_unit = infer_depth_unit(depth.dtype) if input_unit == "auto" else input_unit
    return depth.astype(np.float32, order="K"), resolved_unit


def quantize_depth(
    depth: NDArray[np.uint16] | NDArray[np.float32] | torch.Tensor,
    depth_min: float = DEFAULT_DEPTH_MIN,
    depth_max: float = DEFAULT_DEPTH_MAX,
    shift: float = DEFAULT_DEPTH_SHIFT,
    use_log: bool = DEFAULT_DEPTH_USE_LOG,
    pix_fmt: str = DEFAULT_DEPTH_PIX_FMT,
    video_backend: str | None = "pyav",
    input_unit: Literal["auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT] = "auto",
) -> NDArray[np.uint16] | av.VideoFrame:
    """Quantize depth to 12-bit codes (``uint16``, values ``0…DEPTH_QMAX``).

    Depth maps are packed into 12-bit integer frames so they fit in standard
    high-bit-depth pixel formats (e.g. ``yuv420p12le`` / ``gray12le``)
    and can be encoded by widely supported video codecs (e.g. HEVC Main 12).
    Logarithmic quantization is the default because it allocates more quanta
    to near-range depth, which matches the (1/depth) error profile of typical
    depth sensors. Math is ported from BEHAVIOR-1K's ``obs_utils.py``.

    **Input units**:

    - ``input_unit="auto"`` (default): infer from dtype (floating = m, non-floating = mm).
    - ``input_unit="mm"``: interpret input values as millimetres.
    - ``input_unit="m"``: interpret input values as metres.

    Quantization math runs in the **resolved input unit**.

    ``depth_min``, ``depth_max``, and ``shift`` are always in **metres**.

    Args:
        depth: Depth map; ``torch.Tensor`` is moved to CPU for conversion.
        depth_min: Depth (metres) at quantum ``0``.
        depth_max: Depth (metres) at quantum :data:`DEPTH_QMAX`.
        shift: Depth shift (metres); used in log mode. Must satisfy ``depth_min + shift > 0``.
        use_log: If ``True`` (default), quantize in log space.
        video_backend: Video backend to use for encoding. Defaults to "pyav".
        input_unit: Input unit policy (``"auto"``, ``"mm"``, ``"m"``).

    Returns:
        ``numpy.ndarray``, ``dtype=uint16``, same shape as ``depth``, values in
        ``[0, DEPTH_QMAX]``.

    Raises:
        ValueError: If ``input_unit`` is not ``"auto"``, ``"mm"``, or ``"m"``.
        ValueError: If ``use_log=True`` and ``depth_min + shift <= 0``.
    """
    if input_unit not in ("auto", DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT):
        raise ValueError(
            f"input_unit must be 'auto', '{DEPTH_METER_UNIT}', or '{DEPTH_MILLIMETER_UNIT}', got {input_unit!r}"
        )

    if isinstance(depth, torch.Tensor):
        depth = depth.detach().cpu().numpy()

    # Squeeze single-channel dim: (H, W, 1) or (1, H, W) → (H, W)
    depth = squeeze_single_channel(depth)

    depth_f, resolved_unit = _depth_input_to_float32_and_unit(depth, input_unit=input_unit)

    # Convert depth_min, depth_max, and shift to the resolved input unit.
    depth_min_u = (
        np.float32(depth_min) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_min * MM_PER_METRE)
    )
    depth_max_u = (
        np.float32(depth_max) if resolved_unit == DEPTH_METER_UNIT else np.float32(depth_max * MM_PER_METRE)
    )
    shift_u = np.float32(shift) if resolved_unit == DEPTH_METER_UNIT else np.float32(shift * MM_PER_METRE)

    # Normalization and quantization is performed in the resolved input unit.
    if use_log:
        _validate_log_quant_params(depth_min, shift)
        log_min = math.log(float(depth_min_u + shift_u))
        log_max = math.log(float(depth_max_u + shift_u))
        norm = (np.log(depth_f + shift_u) - log_min) / (log_max - log_min)
    else:
        norm = (depth_f - depth_min_u) / (depth_max_u - depth_min_u)

    quantized = np.rint(norm * DEPTH_QMAX).clip(0, DEPTH_QMAX).astype(np.uint16, copy=False)

    if video_backend == "pyav":
        frame = av.VideoFrame.from_ndarray(quantized, format=pix_fmt)
        write_u16_plane(frame.planes[0], quantized)
        return frame
    else:
        return quantized


def dequantize_depth(
    quantized: NDArray[np.uint16] | av.VideoFrame | torch.Tensor,
    depth_min: float = DEFAULT_DEPTH_MIN,
    depth_max: float = DEFAULT_DEPTH_MAX,
    shift: float = DEFAULT_DEPTH_SHIFT,
    use_log: bool = DEFAULT_DEPTH_USE_LOG,
    pix_fmt: str = DEFAULT_DEPTH_PIX_FMT,
    output_unit: Literal[DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT] = DEPTH_MILLIMETER_UNIT,
    output_tensor: bool = True,
    output_channel_last: bool = False,
) -> NDArray[np.uint16] | NDArray[np.float32] | torch.Tensor:
    """Inverse of :func:`quantize_depth`.

    Decoding inverts the same normalized code mapping as :func:`quantize_depth`
    using ``depth_min`` / ``depth_max`` / ``shift`` (in metres), then returns
    the requested output unit. Tuning arguments **must match** :func:`quantize_depth`.

    Accepted input layouts :

    - ``(H, W, 1)`` or ``(H, W)`` — single frame with channel-last.
    - ``(..., 1, H, W)`` — batched frames with channel-first.
    - ``(..., H, W, 1)`` — batched frames with channel-last.
    Output layout is determined by ``output_channel_last``.

    Args:
        quantized: 12-bit codes in ``[0, DEPTH_QMAX]``. ``np.ndarray``,
            ``av.VideoFrame``, or ``torch.Tensor`` (any integer or float dtype).
        depth_min, depth_max, shift, use_log: Same as :func:`quantize_depth` (metres).
        pix_fmt: Pixel format used to extract the plane from an ``av.VideoFrame``.
        output_unit: ``"mm"`` returns ``uint16`` millimetres (rint, clip
            ``[0, 65535]``) when returning a numpy array, or ``float32`` mm when
            ``output_tensor=True``. ``"m"`` returns ``float32`` metres in
            ``[depth_min, depth_max]``.
        output_tensor: If True, return a ``torch.Tensor`` instead of a numpy array.

    Returns:
        Depth map in the requested unit and dtype.

    Raises:
        ValueError: If ``output_unit`` is not ``"m"`` or ``"mm"``.
        ValueError: If ``use_log=True`` and ``depth_min + shift <= 0``.
    """
    if output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT):
        raise ValueError(
            f"output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {output_unit!r}"
        )
    if use_log:
        _validate_log_quant_params(depth_min, shift)

    if isinstance(quantized, av.VideoFrame):
        quantized = quantized.to_ndarray(format=pix_fmt)

    # Compute the scale and offset first.
    depth_min_m = float(depth_min)
    depth_max_m = float(depth_max)
    shift_m = float(shift)
    if use_log:
        log_min = math.log(depth_min_m + shift_m)
        log_max = math.log(depth_max_m + shift_m)
        scale = (log_max - log_min) / DEPTH_QMAX
        offset = log_min
    else:
        scale = (depth_max_m - depth_min_m) / DEPTH_QMAX
        offset = depth_min_m

    # ── Torch path: stay on the input device, single fp32 allocation. ────────
    if isinstance(quantized, torch.Tensor):
        if quantized.ndim >= 3:
            # Drop the single-channel dimension so the math runs on (..., H, W).
            quantized = quantized.squeeze(-3) if quantized.shape[-3] == 1 else quantized.squeeze(-1)

        # Single allocation we own; everything else is in-place.
        buf = quantized.to(dtype=torch.float32, copy=True)
        buf.mul_(scale).add_(offset)
        if use_log:
            buf.exp_().sub_(shift_m)
        buf.clamp_(depth_min_m, depth_max_m)
        buf.unsqueeze_(-1) if output_channel_last else buf.unsqueeze_(-3)

        if output_unit == DEPTH_METER_UNIT:
            return buf if output_tensor else buf.cpu().numpy()

        # mm path: round + clamp in float32, skipping the uint16 round-trip
        # when returning a tensor (torch.uint16 is poorly supported).
        buf.mul_(MM_PER_METRE).round_().clamp_(0.0, _UINT16_MAX)
        if output_tensor:
            return buf
        return buf.cpu().numpy().astype(np.uint16, copy=False)

    # ── NumPy path: single fp32 allocation, ``out=`` for in-place math. ─────
    arr = np.asarray(quantized)
    if arr.ndim >= 3:
        # Drop the single-channel dimension so the math runs on (..., H, W).
        arr = np.squeeze(arr, axis=-3) if arr.shape[-3] == 1 else np.squeeze(arr, axis=-1)

    buf = np.empty(arr.shape, dtype=np.float32)
    np.multiply(arr, scale, out=buf)
    np.add(buf, offset, out=buf)
    if use_log:
        np.exp(buf, out=buf)
        np.subtract(buf, shift_m, out=buf)
    np.clip(buf, depth_min_m, depth_max_m, out=buf)
    buf = np.expand_dims(buf, axis=-1) if output_channel_last else np.expand_dims(buf, axis=-3)

    if output_unit == DEPTH_METER_UNIT:
        return torch.from_numpy(buf) if output_tensor else buf

    np.multiply(buf, MM_PER_METRE, out=buf)
    np.rint(buf, out=buf)
    np.clip(buf, 0.0, _UINT16_MAX, out=buf)
    if output_tensor:
        # torch.uint16 support is very limited; return float32 millimetres.
        return torch.from_numpy(buf)
    return buf.astype(np.uint16, copy=False)