# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# 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 quantization helpers adapted from lerobot v0.6.0
# (``src/lerobot/datasets/depth_utils.py``, ``pyav_utils.py``,
# ``image_writer.py`` and ``configs/video.py``), Apache-2.0.
#
# Vendored verbatim in math so that sidecar depth videos produced here are
# binary-compatible with the official lerobot 0.6.0 reader (see issue #424).
"""Depth quantization/dequantization for compressed depth sidecar videos.
Depth maps are packed into 12-bit integer codes (``uint16``, values
``0…DEPTH_QMAX``) so they fit the ``gray12le`` pixel format and can be encoded
losslessly by HEVC (``libx265``). Logarithmic quantization is the default
because it allocates more quanta to near-range depth, matching the
``1/depth`` error profile of typical depth sensors.
The math is ported from lerobot 0.6.0 (itself adapted from BEHAVIOR-1K's
``obs_utils.py``) and depends only on :mod:`av`, :mod:`numpy` and
:mod:`torch` -- all of which run on Python 3.10--3.12. See issue #424.
"""
from __future__ import annotations
import math
from typing import Literal
import av
import numpy as np
import torch
from numpy.typing import NDArray
__all__ = [
"DEPTH_QUANT_BITS",
"DEPTH_QMAX",
"DEPTH_METER_UNIT",
"DEPTH_MILLIMETER_UNIT",
"DEFAULT_DEPTH_MIN",
"DEFAULT_DEPTH_MAX",
"DEFAULT_DEPTH_SHIFT",
"DEFAULT_DEPTH_USE_LOG",
"DEFAULT_DEPTH_PIX_FMT",
"MM_PER_METRE",
"infer_depth_unit",
"squeeze_single_channel",
"write_u16_plane",
"quantize_depth",
"dequantize_depth",
]
# ---------------------------------------------------------------------------
# Constants (from lerobot ``configs/video.py``).
# ---------------------------------------------------------------------------
DEPTH_QUANT_BITS: int = 12
DEPTH_QMAX: int = (1 << DEPTH_QUANT_BITS) - 1 # 4095
DEFAULT_DEPTH_MIN: float = 0.01
DEFAULT_DEPTH_MAX: float = 10.0
DEFAULT_DEPTH_SHIFT: float = 3.5
DEFAULT_DEPTH_USE_LOG: bool = True
DEFAULT_DEPTH_PIX_FMT: str = "gray12le"
DEPTH_METER_UNIT: str = "m"
DEPTH_MILLIMETER_UNIT: str = "mm"
MM_PER_METRE = 1000.0
_UINT16_MAX = 65535
[docs]
def infer_depth_unit(dtype: np.dtype | type) -> str:
"""Infer the depth unit from the array dtype.
Floating-point arrays are interpreted as metres, integer arrays as
millimetres -- the lerobot convention.
Args:
dtype: NumPy dtype (or anything coercible to one).
Returns:
``"m"`` or ``"mm"``.
"""
return (
DEPTH_METER_UNIT
if np.issubdtype(np.dtype(dtype), np.floating)
else DEPTH_MILLIMETER_UNIT
)
[docs]
def squeeze_single_channel(array: np.ndarray) -> np.ndarray:
"""Drop a leading or trailing singleton channel dim: ``(1, H, W)`` / ``(H, W, 1)`` -> ``(H, W)``.
Unlike ``array.squeeze()``, this only removes the channel axis, never an
``H`` or ``W`` of size 1.
"""
if array.ndim == 3:
if array.shape[0] == 1:
return array[0]
if array.shape[-1] == 1:
return array[..., 0]
return array
[docs]
def write_u16_plane(
plane: av.video.plane.VideoPlane, src: np.ndarray, fill_value: int | None = None
) -> None:
"""Copy a 2D ``uint16`` image into the plane's memory buffer, row by row.
For speed, each row is padded to a wider size than ``width``, so the true row width in
memory is ``plane.line_size`` (bytes), not ``width``. Copying as one straight stream
would skew the image, so we write only the first ``width`` columns of each row and
leave the padding untouched.
Args:
plane: Destination 16-bit plane.
src: Source image, shape ``(height, width)``, dtype ``uint16``.
fill_value: If given, every pixel (padding included) is set to this first, so the
padding holds clean data instead of garbage.
"""
height, width = src.shape
stride_u16 = plane.line_size // np.dtype(np.uint16).itemsize
dst = np.frombuffer(plane, dtype=np.uint16).reshape(height, stride_u16)
if fill_value is not None:
dst.fill(fill_value)
dst[:, :width] = src
# ---------------------------------------------------------------------------
# Quantization / dequantization (from lerobot ``depth_utils.py``).
# ---------------------------------------------------------------------------
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
[docs]
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: If ``"pyav"`` (default), return an :class:`av.VideoFrame`
ready for encoding; otherwise return the raw ``uint16`` code array.
input_unit: Input unit policy (``"auto"``, ``"mm"``, ``"m"``).
Returns:
``av.VideoFrame`` (when ``video_backend="pyav"``) or ``numpy.ndarray``
of ``dtype=uint16``, same spatial 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
[docs]
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.
output_channel_last: Channel layout of the output (``(H, W, 1)`` vs ``(1, H, W)``).
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)