apairo-preprocess 0.2.0__py3-none-any.whl

apairo_preprocess/pointcloud/voxelise.py

@@ -0,0 +1,281 @@
+"""Voxel-grid downsampling of LiDAR point clouds and their semantic labels.
+
+Three preprocessors share a common voxel assignment and must be run with the
+same ``voxel_size`` and ``max_range`` to stay index-aligned:
+
+* :class:`VoxelisePointCloud` — one representative float point per voxel.
+* :class:`VoxeliseLabels`     — one aggregated label per voxel (majority/max).
+* :class:`VoxeliseCoords`     — integer voxel coordinates (N', 3) int32,
+                                 ready for sparse convolution networks.
+
+After preprocessing, ``voxelised[i]``, ``voxelised_labels[i]``, and
+``voxel_coords[i]`` all refer to the same voxel cell.
+
+Typical usage::
+
+    for prep in [
+        VoxelisePointCloud(voxel_size=0.1, max_range=50.0),
+        VoxeliseLabels(voxel_size=0.1, max_range=50.0),
+        VoxeliseCoords(voxel_size=0.1, max_range=50.0),
+    ]:
+        dataset.run_preprocess(prep, split_dir)
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar, Literal
+
+import numpy as np
+
+from apairo.core.preprocessor import FramePreprocessor
+from apairo.core.sample import Sample
+
+_Reduction = Literal["centroid", "first", "random"]
+_LabelAggregation = Literal["majority", "max"]
+
+
+def _filter_and_quantize(
+    xyz: np.ndarray,
+    voxel_size: float,
+    max_range: float | None,
+) -> tuple[np.ndarray | None, np.ndarray, np.ndarray, int]:
+    """Apply range filter and compute voxel assignment.
+
+    Returns:
+        mask:          Boolean keep-mask (or None if no filtering was done).
+        unique_coords: Integer voxel coordinates, shape (n_voxels, 3) int32.
+        inverse:       Per-point voxel index array, shape (N_kept,).
+        n_voxels:      Number of occupied voxels.
+    """
+    mask: np.ndarray | None = None
+    if max_range is not None:
+        mask = np.linalg.norm(xyz, axis=1) < max_range
+        xyz = xyz[mask]
+
+    if len(xyz) == 0:
+        empty_coords = np.empty((0, 3), dtype=np.int32)
+        return mask, empty_coords, np.empty(0, dtype=np.int64), 0
+
+    coords = np.floor(xyz / voxel_size).astype(np.int32)
+    unique_coords, inverse = np.unique(coords, axis=0, return_inverse=True)
+    return mask, unique_coords, inverse, len(unique_coords)
+
+
+class VoxelisePointCloud(FramePreprocessor):
+    """Downsample a point cloud with a voxel grid.
+
+    Each voxel retains exactly one point according to ``reduction``:
+
+    * ``"centroid"`` — mean of all points falling in the voxel (all channels).
+    * ``"first"``    — first point encountered (input order).
+    * ``"random"``   — one uniformly random point per voxel.
+
+    Args:
+        lidar_key:   Input channel for point cloud data.
+        voxel_size:  Edge length of each cubic voxel cell (metres).
+        max_range:   Discard points farther than this distance from the sensor
+                     origin (metres).  ``None`` keeps all points.
+        reduction:   How to pick the representative point per voxel.
+        output_key:  Override the default output channel name ``"voxelised"``.
+    """
+
+    output_key: ClassVar[str] = "voxelised"
+    input_keys: ClassVar[list[str]] = ["lidar"]
+    timestamps_from: ClassVar[str] = "lidar"
+    sources: ClassVar[list[str]] = ["lidar"]
+
+    def __init__(
+        self,
+        lidar_key: str = "lidar",
+        voxel_size: float = 0.1,
+        max_range: float | None = None,
+        reduction: _Reduction = "centroid",
+        output_key: str | None = None,
+    ) -> None:
+        if voxel_size <= 0:
+            raise ValueError(f"voxel_size must be positive, got {voxel_size}")
+        if max_range is not None and max_range <= 0:
+            raise ValueError(f"max_range must be positive, got {max_range}")
+        if reduction not in ("centroid", "first", "random"):
+            raise ValueError(
+                f"reduction must be 'centroid', 'first', or 'random', got {reduction!r}"
+            )
+
+        self._lidar_key = lidar_key
+        self._voxel_size = float(voxel_size)
+        self._max_range = max_range
+        self._reduction = reduction
+
+        self.input_keys = [lidar_key]
+        self.sources = [lidar_key]
+        if output_key is not None:
+            self.output_key = output_key
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        pc = np.asarray(sample.data[self._lidar_key], dtype=np.float64)
+        if pc.ndim != 2 or pc.shape[1] < 3:
+            raise ValueError(f"Point cloud must be (N, D>=3), got {pc.shape}")
+
+        mask, _, inverse, n_voxels = _filter_and_quantize(
+            pc[:, :3], self._voxel_size, self._max_range
+        )
+        if mask is not None:
+            pc = pc[mask]
+
+        if n_voxels == 0:
+            return np.empty((0, pc.shape[1]), dtype=np.float64)
+
+        if self._reduction == "centroid":
+            out = np.zeros((n_voxels, pc.shape[1]), dtype=np.float64)
+            counts = np.zeros(n_voxels, dtype=np.int64)
+            np.add.at(out, inverse, pc)
+            np.add.at(counts, inverse, 1)
+            out /= counts[:, None]
+            return out
+
+        if self._reduction == "first":
+            out = np.empty((n_voxels, pc.shape[1]), dtype=np.float64)
+            # assign in reverse so the first occurrence wins
+            out[inverse] = pc
+            return out
+
+        # random: shuffle so a random point ends up last (and thus selected)
+        rng = np.random.default_rng()
+        chosen = np.empty(n_voxels, dtype=np.int64)
+        perm = rng.permutation(len(pc))
+        chosen[inverse[perm]] = perm
+        return pc[chosen]
+
+
+class VoxeliseLabels(FramePreprocessor):
+    """Aggregate per-point semantic labels onto the same voxel grid.
+
+    Must be run with the **same** ``voxel_size`` and ``max_range`` as
+    :class:`VoxelisePointCloud` so the output channels remain index-aligned.
+
+    * ``"majority"`` — most frequent label per voxel (semantically correct).
+    * ``"max"``      — highest label ID per voxel (fast, mirrors common
+                       sparse-convolution data loaders).
+
+    Args:
+        lidar_key:    Input channel for point cloud data (geometry only).
+        labels_key:   Input channel for per-point semantic labels.
+        voxel_size:   Edge length of each cubic voxel cell (metres).
+        max_range:    Same range filter as applied to the point cloud.
+        aggregation:  Label merging strategy per voxel.
+        output_key:   Override the default output channel name
+                      ``"voxelised_labels"``.
+    """
+
+    output_key: ClassVar[str] = "voxelised_labels"
+    input_keys: ClassVar[list[str]] = ["lidar", "labels"]
+    timestamps_from: ClassVar[str] = "lidar"
+    sources: ClassVar[list[str]] = ["lidar", "labels"]
+
+    def __init__(
+        self,
+        lidar_key: str = "lidar",
+        labels_key: str = "labels",
+        voxel_size: float = 0.1,
+        max_range: float | None = None,
+        aggregation: _LabelAggregation = "majority",
+        output_key: str | None = None,
+    ) -> None:
+        if voxel_size <= 0:
+            raise ValueError(f"voxel_size must be positive, got {voxel_size}")
+        if max_range is not None and max_range <= 0:
+            raise ValueError(f"max_range must be positive, got {max_range}")
+        if aggregation not in ("majority", "max"):
+            raise ValueError(f"aggregation must be 'majority' or 'max', got {aggregation!r}")
+
+        self._lidar_key = lidar_key
+        self._labels_key = labels_key
+        self._voxel_size = float(voxel_size)
+        self._max_range = max_range
+        self._aggregation = aggregation
+
+        self.input_keys = [lidar_key, labels_key]
+        self.sources = [lidar_key, labels_key]
+        if output_key is not None:
+            self.output_key = output_key
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        xyz = np.asarray(sample.data[self._lidar_key], dtype=np.float64)[:, :3]
+        labels = np.asarray(sample.data[self._labels_key])
+
+        mask, _, inverse, n_voxels = _filter_and_quantize(
+            xyz, self._voxel_size, self._max_range
+        )
+        if mask is not None:
+            labels = labels[mask]
+
+        if n_voxels == 0:
+            return np.empty(0, dtype=labels.dtype)
+
+        if self._aggregation == "max":
+            out = np.zeros(n_voxels, dtype=labels.dtype)
+            np.maximum.at(out, inverse, labels)
+            return out
+
+        # majority: encode (voxel_idx, label) as a single int and bincount
+        n_classes = int(labels.max()) + 1
+        keys = inverse.astype(np.int64) * n_classes + labels.astype(np.int64)
+        pair_counts = np.bincount(keys, minlength=n_voxels * n_classes)
+        return pair_counts.reshape(n_voxels, n_classes).argmax(axis=1).astype(labels.dtype)
+
+
+class VoxeliseCoords(FramePreprocessor):
+    """Save integer voxel coordinates for each occupied cell.
+
+    Produces an ``(N', 3)`` int32 array where row ``i`` is the quantised
+    coordinate ``floor(xyz / voxel_size)`` of the ``i``-th voxel.  This is the
+    format expected by sparse convolution frameworks (torchsparse,
+    MinkowskiEngine) and eliminates the ``floor`` recomputation from every
+    ``__getitem__``.
+
+    Must be run with the **same** ``voxel_size`` and ``max_range`` as
+    :class:`VoxelisePointCloud` and :class:`VoxeliseLabels` so all three
+    channels stay index-aligned.
+
+    Args:
+        lidar_key:   Input channel for point cloud data.
+        voxel_size:  Edge length of each cubic voxel cell (metres).
+        max_range:   Same range filter as applied to the point cloud.
+        output_key:  Override the default output channel name ``"voxel_coords"``.
+    """
+
+    output_key: ClassVar[str] = "voxel_coords"
+    input_keys: ClassVar[list[str]] = ["lidar"]
+    timestamps_from: ClassVar[str] = "lidar"
+    sources: ClassVar[list[str]] = ["lidar"]
+
+    def __init__(
+        self,
+        lidar_key: str = "lidar",
+        voxel_size: float = 0.1,
+        max_range: float | None = None,
+        output_key: str | None = None,
+    ) -> None:
+        if voxel_size <= 0:
+            raise ValueError(f"voxel_size must be positive, got {voxel_size}")
+        if max_range is not None and max_range <= 0:
+            raise ValueError(f"max_range must be positive, got {max_range}")
+
+        self._lidar_key = lidar_key
+        self._voxel_size = float(voxel_size)
+        self._max_range = max_range
+
+        self.input_keys = [lidar_key]
+        self.sources = [lidar_key]
+        if output_key is not None:
+            self.output_key = output_key
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        xyz = np.asarray(sample.data[self._lidar_key], dtype=np.float64)[:, :3]
+        if xyz.ndim != 2 or xyz.shape[1] != 3:
+            raise ValueError(f"LiDAR channel must be (N, D>=3), got {xyz.shape}")
+
+        _, unique_coords, _, _ = _filter_and_quantize(
+            xyz, self._voxel_size, self._max_range
+        )
+        return unique_coords  # (N', 3) int32

apairo_preprocess/priors/__init__.py

@@ -0,0 +1,7 @@
+from apairo_preprocess.priors.ground_height import GroundHeightFromLabels
+from apairo_preprocess.priors.trajectory_distance import TrajectoryDistance
+
+__all__ = [
+    "GroundHeightFromLabels",
+    "TrajectoryDistance",
+]

apairo_preprocess/priors/ground_height.py

@@ -0,0 +1,60 @@
+"""Height above ground surface from binary ground segmentation labels."""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import numpy as np
+from scipy.spatial import KDTree
+
+from apairo.core.preprocessor import FramePreprocessor
+from apairo.core.sample import Sample
+
+
+class GroundHeightFromLabels(FramePreprocessor):
+    """Per-voxel height above ground from any binary ground segmentation.
+
+    Takes the output of any ground segmentation preprocessor
+    (``GroundSegmentationCSF``, ``GroundSegmentationRANSAC``,
+    ``GroundSegmentationFromLabels``) and computes the signed height (m) of
+    each voxel above the nearest ground point in XY.
+
+    Args:
+        ground_key:    Input channel for binary ground labels (0=ground,
+                       1=non-ground).  Typically ``"ground_csf"``,
+                       ``"ground_ransac"``, or ``"ground_labels"``.
+        voxelised_key: Input channel for the voxelised point cloud
+                       (shape N x D≥3).
+        output_key:    Override the default channel name ``"ground_height"``.
+    """
+
+    output_key: ClassVar[str] = "ground_height"
+    input_keys: ClassVar[list[str]] = ["voxelised", "ground_labels"]
+    timestamps_from: ClassVar[str] = "voxelised"
+    sources: ClassVar[list[str]] = ["voxelised", "ground_labels"]
+
+    def __init__(
+        self,
+        ground_key: str,
+        voxelised_key: str = "voxelised",
+        output_key: str | None = None,
+    ) -> None:
+        self._ground_key = ground_key
+        self._voxelised_key = voxelised_key
+        self.input_keys = [voxelised_key, ground_key]
+        self.sources = [voxelised_key, ground_key]
+        if output_key is not None:
+            self.output_key = output_key
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        xyz = np.asarray(sample.data[self._voxelised_key], dtype=np.float64)[:, :3]
+        labels = np.asarray(sample.data[self._ground_key])
+        ground_pts = xyz[labels == 0]
+
+        if len(ground_pts) == 0:
+            z_ground = float(np.percentile(xyz[:, 2], 5))
+            return (xyz[:, 2] - z_ground).astype(np.float32)
+
+        tree = KDTree(ground_pts[:, :2])
+        _, nn = tree.query(xyz[:, :2])
+        return (xyz[:, 2] - ground_pts[nn, 2]).astype(np.float32)

apairo_preprocess/priors/trajectory_distance.py

@@ -0,0 +1,97 @@
+"""Trajectory-distance preprocessor for Weakly Informed PUNCE.
+
+Computes the Euclidean distance (metres) from each voxelised point to the
+nearest waypoint on the robot's recorded trajectory.  Points that lie close
+to the path the robot actually drove are likely traversable.
+
+The output float32 array of shape (N,) is consumed by
+:class:`WeaklyInformedPUNCELoss` to compute the Gaussian prior
+
+    π(p) = exp(−d(p, traj)² / (2σ²))
+
+All trajectory waypoints must be passed at construction time
+(as world-frame positions) so the preprocessor can build a single KD-tree
+shared across frames.
+
+Typical usage in a preprocess script::
+
+    ds = Rellis3DDataset(root, keys=["poses"]).transform("poses", PoseTo4x4())
+    all_poses = np.stack([ds[i].data["poses"] for i in range(len(ds))])
+    trajectory = all_poses[:, :3, 3]          # (N_frames, 3) world positions
+
+    Rellis3DDataset.run_preprocess(
+        TrajectoryDistance(trajectory=trajectory),
+        root,
+    )
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import numpy as np
+from scipy.spatial import KDTree
+
+from apairo.core.preprocessor import FramePreprocessor
+from apairo.core.sample import Sample
+
+
+
+class TrajectoryDistance(FramePreprocessor):
+    """Per-voxel Euclidean distance to the nearest trajectory waypoint.
+
+    For each frame, voxelised points are transformed to the world frame using
+    the current pose, then queried against a KD-tree built from all trajectory
+    waypoints.
+
+    Args:
+        trajectory:    (M, 3) array of robot positions in world frame.
+                       Typically ``all_poses[:, :3, 3]`` from the sequence.
+        voxelised_key: Input channel for the voxelised point cloud.
+        poses_key:     Input channel for the current-frame pose — must be ``(4, 4)``
+                       float64.  Apply ``PoseTo4x4()`` from ``apairo_transform`` if needed.
+        output_key:    Override the default channel name
+                       ``"trajectory_distance"``.
+    """
+
+    output_key: ClassVar[str] = "trajectory_distance"
+    input_keys: ClassVar[list[str]] = ["voxelised", "poses"]
+    timestamps_from: ClassVar[str] = "voxelised"
+    sources: ClassVar[list[str]] = ["voxelised", "poses"]
+
+    def __init__(
+        self,
+        trajectory: np.ndarray,
+        voxelised_key: str = "voxelised",
+        poses_key: str = "poses",
+        output_key: str | None = None,
+    ) -> None:
+        traj = np.asarray(trajectory, dtype=np.float64)
+        if traj.ndim != 2 or traj.shape[1] != 3:
+            raise ValueError(
+                f"trajectory must be (M, 3), got {traj.shape}"
+            )
+        self._tree = KDTree(traj)
+        self._voxelised_key = voxelised_key
+        self._poses_key = poses_key
+        self.input_keys = [voxelised_key, poses_key]
+        self.sources = [voxelised_key, poses_key]
+        if output_key is not None:
+            self.output_key = output_key
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        pc = np.asarray(sample.data[self._voxelised_key], dtype=np.float64)
+        pose = np.asarray(sample.data[self._poses_key], dtype=np.float64)
+        if pose.shape != (4, 4):
+            raise ValueError(
+                f"pose must be (4, 4) — apply PoseTo4x4() from apairo_transform first. "
+                f"Got {pose.shape}"
+            )
+
+        n = len(pc)
+        xyz_sensor = pc[:, :3]
+        xyz_h = np.column_stack([xyz_sensor, np.ones(n)])
+        xyz_world = (pose @ xyz_h.T).T[:, :3]
+
+        dist, _ = self._tree.query(xyz_world)
+        return dist.astype(np.float32)

apairo_preprocess/projection/__init__.py

@@ -0,0 +1,12 @@
+from apairo_preprocess.projection.bev import BEVRasterisation, to_uint8_image
+from apairo_preprocess.projection.image_mask import ImageMaskFromPointLabels
+from apairo_preprocess.projection.lidar_to_camera import LidarCameraProjection
+from apairo_preprocess.projection.point_features import PointFeaturesFromImage
+
+__all__ = [
+    "BEVRasterisation",
+    "to_uint8_image",
+    "ImageMaskFromPointLabels",
+    "LidarCameraProjection",
+    "PointFeaturesFromImage",
+]

apairo_preprocess/projection/bev.py

@@ -0,0 +1,166 @@
+"""Rasterise per-point values into a top-down (bird's-eye-view) grid.
+
+Unlike :class:`LidarCameraProjection` + :class:`ImageMaskFromPointLabels`
+(perspective projection into a camera, nearest-depth-wins occlusion), a BEV
+grid needs no camera at all: points are binned by their XY position onto an
+orthographic top-down grid, so it works on any lidar scan regardless of
+whether a calibrated camera exists for it -- useful for datasets that ship
+point clouds without a registered camera/lidar extrinsic.
+
+Row 0 of the output is north (``ymax``), so plotting the grid directly reads
+top-down.
+
+Typical usage::
+
+    TartanKittiDataset.run_preprocess(
+        BEVRasterisation(bounds=(-20, 20, -20, 20), resolution=0.2,
+                          points_key="lidar", values_key="trav_traj"),
+        seq_dir,
+    )
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import numpy as np
+
+from apairo.core.preprocessor import FramePreprocessor
+from apairo.core.sample import Sample
+
+
+def _rasterise(points, values, *, bounds, resolution, reduce, background):
+    xmin, xmax, ymin, ymax = bounds
+    w = int(round((xmax - xmin) / resolution))
+    h = int(round((ymax - ymin) / resolution))
+
+    xy = np.asarray(points)[:, :2]
+    vals = np.asarray(values, dtype=np.float64)
+
+    col = np.floor((xy[:, 0] - xmin) / resolution).astype(np.int64)
+    row = np.floor((ymax - xy[:, 1]) / resolution).astype(np.int64)
+    inside = (col >= 0) & (col < w) & (row >= 0) & (row < h)
+    cell = row[inside] * w + col[inside]
+    vals = vals[inside]
+
+    grid = np.full(h * w, background, dtype=np.float64)
+    if reduce == "max":
+        # A plain np.maximum.at into `grid` would use `background` as an
+        # implicit floor -- wrong whenever real values can be below it (e.g.
+        # negative heights). Accumulate from -inf, then fill only the cells
+        # no point touched.
+        acc = np.full(h * w, -np.inf)
+        np.maximum.at(acc, cell, vals)
+        hit = np.isfinite(acc)
+        grid[hit] = acc[hit]
+    elif reduce == "last":
+        grid[cell] = vals
+    else:  # "mean"
+        count = np.zeros(h * w)
+        total = np.zeros(h * w)
+        np.add.at(count, cell, 1.0)
+        np.add.at(total, cell, vals)
+        hit = count > 0
+        grid[hit] = total[hit] / count[hit]
+    return grid.reshape(h, w)
+
+
+class BEVRasterisation(FramePreprocessor):
+    """Per-point values splatted onto a top-down ``(H, W)`` grid.
+
+    Output is float64, row-major with row 0 at ``ymax`` (north up) and
+    column 0 at ``xmin``. Cells no point falls into hold ``background``.
+
+    Args:
+        bounds:      ``(xmin, xmax, ymin, ymax)`` in metres, in the point
+                     cloud's own frame (typically sensor- or robot-centred).
+        resolution:  Metres per pixel.
+        points_key:  Input channel for the point cloud, ``(N, >=2)``; only
+                     the X/Y columns are used.
+        values_key:  Input channel for the per-point scalar to rasterise
+                     (e.g. height, a traversability label), ``(N,)``,
+                     row-aligned with ``points_key``.
+        reduce:      How to combine points landing in the same cell --
+                     ``"max"`` (default), ``"mean"`` or ``"last"``.
+        background:  Value for cells no point falls into.
+        output_key:  Override the default output channel name ``"bev"``.
+    """
+
+    output_key: ClassVar[str] = "bev"
+    input_keys: ClassVar[list[str]] = ["lidar", "values"]
+    timestamps_from: ClassVar[str] = "lidar"
+    sources: ClassVar[list[str]] = ["lidar", "values"]
+
+    def __init__(
+        self,
+        bounds: tuple[float, float, float, float],
+        resolution: float,
+        points_key: str = "lidar",
+        values_key: str = "values",
+        reduce: str = "max",
+        background: float = 0.0,
+        output_key: str | None = None,
+    ) -> None:
+        xmin, xmax, ymin, ymax = bounds
+        if not (xmax > xmin and ymax > ymin):
+            raise ValueError(f"bounds must have xmax > xmin and ymax > ymin, got {bounds}")
+        if resolution <= 0:
+            raise ValueError(f"resolution must be > 0, got {resolution}")
+        if reduce not in ("max", "mean", "last"):
+            raise ValueError(f"reduce must be 'max', 'mean' or 'last', got {reduce!r}")
+
+        self._bounds = (float(xmin), float(xmax), float(ymin), float(ymax))
+        self._resolution = float(resolution)
+        self._reduce = reduce
+        self._background = float(background)
+        self._points_key = points_key
+        self._values_key = values_key
+
+        self.input_keys = [points_key, values_key]
+        self.sources = [points_key, values_key]
+        self.timestamps_from = points_key
+        if output_key is not None:
+            self.output_key = output_key
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        """``(height, width)`` of the output grid, from ``bounds``/``resolution``."""
+        xmin, xmax, ymin, ymax = self._bounds
+        w = int(round((xmax - xmin) / self._resolution))
+        h = int(round((ymax - ymin) / self._resolution))
+        return h, w
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        points = np.asarray(sample.data[self._points_key])
+        values = np.asarray(sample.data[self._values_key])
+        if points.ndim != 2 or points.shape[1] < 2:
+            raise ValueError(f"points must be (N, >=2), got {points.shape}")
+        if len(values) != len(points):
+            raise ValueError(
+                f"points ({len(points)}) and values ({len(values)}) are not "
+                f"row-aligned -- they must derive from the same scan."
+            )
+        return _rasterise(
+            points, values,
+            bounds=self._bounds, resolution=self._resolution,
+            reduce=self._reduce, background=self._background,
+        )
+
+
+def to_uint8_image(
+    grid: np.ndarray, *, vmin: float | None = None, vmax: float | None = None
+) -> np.ndarray:
+    """Normalise a ``(H, W)`` grid to an ``(H, W, 3)`` uint8 image.
+
+    Pass fixed ``vmin``/``vmax`` for anything beyond one-off visualisation:
+    left at ``None``, each call normalises to its *own* min/max, so the same
+    physical value maps to a different pixel intensity from one frame to the
+    next -- fine to eyeball a single grid, wrong for a dataset a model trains
+    on (a given value must mean the same thing in every frame).
+    """
+    grid = np.asarray(grid, dtype=np.float64)
+    lo = grid.min() if vmin is None else vmin
+    hi = grid.max() if vmax is None else vmax
+    norm = (grid - lo) / (hi - lo) if hi > lo else np.zeros_like(grid)
+    gray = np.clip(norm * 255.0, 0, 255).astype(np.uint8)
+    return np.repeat(gray[:, :, None], 3, axis=2)