apairo-preprocess 0.2.0__py3-none-any.whl

apairo_preprocess/projection/image_mask.py

@@ -0,0 +1,181 @@
+"""Rasterize per-point labels into an image mask — point cloud → image.
+
+The forward direction of the projection composition: given any per-point
+label channel (e.g. ``trav_traj`` from
+:class:`~apairo_preprocess.TraversabilityFromTrajectory`) and the per-point
+pixel coordinates written by :class:`LidarCameraProjection`, paint each
+label at its point's pixel.  The result is a sparse image-space ground-truth
+mask — the traversability channel of an image dataset.
+
+Pixels hit by several points keep the **nearest** point's label (depth
+ordering), and an optional occlusion filter drops points that are
+significantly behind the nearest return in their neighbourhood, so
+traversable ground behind an obstacle does not bleed onto the obstacle's
+pixels.
+
+This preprocessor reads two channels, both on the lidar clock.  On profiled
+synchronous datasets it runs via ``run_preprocess`` directly; on
+asynchronous datasets, run it over a ``synchronize()`` view and persist with
+``ChannelWriter`` (see ``examples/traversability_image_mask.py``).
+
+Typical usage::
+
+    Rellis3DDataset.run_preprocess(
+        ImageMaskFromPointLabels(image_size=(1200, 1920), radius=2),
+        root,
+    )
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import numpy as np
+
+from apairo.core.preprocessor import FramePreprocessor
+from apairo.core.sample import Sample
+
+
+class ImageMaskFromPointLabels(FramePreprocessor):
+    """Paint per-point labels into an ``(H, W)`` uint8 mask.
+
+    Output is ``(height, width)`` uint8: ``unknown_value`` where no point
+    projects, the nearest projecting point's label elsewhere.  With the
+    default inputs (``trav_traj``) that is ``255`` = no data, ``1`` =
+    traversable, ``0`` = not traversable.
+
+    Args:
+        image_size:      ``(height, width)`` of the output mask — must match
+                         the ``image_size`` the projection was computed for.
+        labels_key:      Input channel with per-point integer labels ``(N,)``,
+                         row-aligned with the projection.
+        projection_key:  Input channel with ``(N, 3)`` ``[u, v, depth]`` rows
+                         from :class:`LidarCameraProjection`.
+        radius:          Splat radius in pixels.  ``0`` (default) paints one
+                         pixel per point; ``r`` paints a disc of radius ``r``,
+                         densifying the sparse lidar coverage.
+        unknown_value:   Value for pixels no point projects to (default 255).
+                         Labels must not collide with it.
+        occlusion_bin_px:      Optional occlusion filtering: bin size in pixels
+                               of a coarse depth buffer.  Points deeper than
+                               the nearest return in their bin by more than
+                               ``occlusion_depth_margin`` are dropped before
+                               painting.  ``None`` (default) disables it.
+        occlusion_depth_margin: Depth slack in metres for the occlusion filter.
+        output_key:      Override the default output channel name
+                         ``"trav_mask"``.
+    """
+
+    output_key: ClassVar[str] = "trav_mask"
+    input_keys: ClassVar[list[str]] = ["trav_traj", "lidar_uv"]
+    timestamps_from: ClassVar[str] = "trav_traj"
+    sources: ClassVar[list[str]] = ["trav_traj", "lidar_uv"]
+
+    def __init__(
+        self,
+        image_size: tuple[int, int],
+        labels_key: str = "trav_traj",
+        projection_key: str = "lidar_uv",
+        radius: int = 0,
+        unknown_value: int = 255,
+        occlusion_bin_px: int | None = None,
+        occlusion_depth_margin: float = 0.5,
+        output_key: str | None = None,
+    ) -> None:
+        height, width = image_size
+        if height <= 0 or width <= 0:
+            raise ValueError(f"image_size must be positive, got {image_size}")
+        if radius < 0:
+            raise ValueError(f"radius must be >= 0, got {radius}")
+        if not 0 <= unknown_value <= 255:
+            raise ValueError(f"unknown_value must fit uint8, got {unknown_value}")
+        if occlusion_bin_px is not None and occlusion_bin_px <= 0:
+            raise ValueError(f"occlusion_bin_px must be > 0, got {occlusion_bin_px}")
+
+        self._height = int(height)
+        self._width = int(width)
+        self._labels_key = labels_key
+        self._projection_key = projection_key
+        self._radius = int(radius)
+        self._unknown = int(unknown_value)
+        self._occ_bin = occlusion_bin_px
+        self._occ_margin = occlusion_depth_margin
+
+        self.input_keys = [labels_key, projection_key]
+        self.timestamps_from = labels_key
+        self.sources = [labels_key, projection_key]
+        if output_key is not None:
+            self.output_key = output_key
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        labels = np.asarray(sample.data[self._labels_key]).ravel()
+        uv = np.asarray(sample.data[self._projection_key])
+        if uv.ndim != 2 or uv.shape[1] < 3:
+            raise ValueError(f"projection must be (N, 3) [u, v, depth], got {uv.shape}")
+        if len(labels) != len(uv):
+            raise ValueError(
+                f"labels ({len(labels)}) and projection ({len(uv)}) are not "
+                f"row-aligned — they must derive from the same scan."
+            )
+
+        height, width = self._height, self._width
+        mask = np.full((height, width), self._unknown, dtype=np.uint8)
+
+        valid = np.isfinite(uv[:, 0])
+        if not valid.any():
+            return mask
+
+        cols = uv[valid, 0].astype(np.intp)
+        rows = uv[valid, 1].astype(np.intp)
+        depth = uv[valid, 2].astype(np.float64)
+        vals = labels[valid]
+        if vals.min() < 0 or vals.max() > 255:
+            raise ValueError(
+                f"labels must fit uint8, got range [{vals.min()}, {vals.max()}]"
+            )
+        vals = vals.astype(np.uint8)
+
+        if self._occ_bin is not None:
+            keep = self._occlusion_keep(rows, cols, depth)
+            rows, cols, depth, vals = rows[keep], cols[keep], depth[keep], vals[keep]
+
+        # Splat each point over a disc of offsets, then resolve every pixel
+        # to its nearest candidate in one pass.
+        r = self._radius
+        offsets = [
+            (du, dv)
+            for du in range(-r, r + 1)
+            for dv in range(-r, r + 1)
+            if du * du + dv * dv <= r * r
+        ]
+        pix_parts, depth_parts, val_parts = [], [], []
+        for du, dv in offsets:
+            rr = rows + dv
+            cc = cols + du
+            ok = (rr >= 0) & (rr < height) & (cc >= 0) & (cc < width)
+            pix_parts.append(rr[ok] * width + cc[ok])
+            depth_parts.append(depth[ok])
+            val_parts.append(vals[ok])
+        pix = np.concatenate(pix_parts)
+        d = np.concatenate(depth_parts)
+        v = np.concatenate(val_parts)
+
+        order = np.lexsort((d, pix))  # by pixel, nearest first within a pixel
+        pix, v = pix[order], v[order]
+        first = np.ones(len(pix), dtype=bool)
+        first[1:] = pix[1:] != pix[:-1]
+        mask.flat[pix[first]] = v[first]
+        return mask
+
+    def _occlusion_keep(
+        self, rows: np.ndarray, cols: np.ndarray, depth: np.ndarray
+    ) -> np.ndarray:
+        """Keep points within ``occlusion_depth_margin`` of their bin's nearest
+        return, on a coarse ``occlusion_bin_px`` grid."""
+        b = self._occ_bin
+        n_bins_x = (self._width + b - 1) // b
+        n_bins_y = (self._height + b - 1) // b
+        bin_id = (rows // b) * n_bins_x + (cols // b)
+        nearest = np.full(n_bins_x * n_bins_y, np.inf)
+        np.minimum.at(nearest, bin_id, depth)
+        return depth <= nearest[bin_id] + self._occ_margin

apairo_preprocess/projection/lidar_to_camera.py

@@ -0,0 +1,177 @@
+"""Project lidar points onto a camera image plane.
+
+For each scan, every point is mapped through the lidar→camera extrinsic and
+the pinhole intrinsics to continuous pixel coordinates.  The output is
+**row-aligned** with the input scan: row *i* is ``[u, v, depth]`` for point
+*i*, or ``NaN`` when the point does not land in the image (behind the camera
+or out of bounds).  Row alignment is what makes the channel composable — any
+per-point channel (traversability labels, priors, …) carries into image space
+by indexing with the same row (see :class:`ImageMaskFromPointLabels` and
+:class:`PointFeaturesFromImage`).
+
+The camera frame must be the **optical** frame (+Z forward, +X right, +Y
+down).  Both halves of the projection come from the dataset's calibration —
+the extrinsic from the transform graph, the intrinsics from the ``cameras:``
+section::
+
+    cal = ds.calibration
+    cam = cal.get_intrinsics("camera_left")            # CameraIntrinsics
+    T   = cal.get_tf("velodyne_0", "camera_left")      # (4, 4)
+
+Typical usage::
+
+    TartanKittiDataset.run_preprocess(
+        LidarCameraProjection(
+            intrinsics=cam,              # CameraIntrinsics, or a (3, 3) K
+            extrinsics=T,                # (4, 4) T_camera_from_lidar
+            lidar_key="velodyne_0",
+            output_key="lidar_uv_left",
+        ),
+        seq_dir,
+    )
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import numpy as np
+
+from apairo.core.config import CameraIntrinsics
+from apairo.core.preprocessor import FramePreprocessor
+from apairo.core.sample import Sample
+
+
+def _distort_plumb_bob(x: np.ndarray, y: np.ndarray, d: np.ndarray) -> tuple:
+    """Apply plumb-bob (radial-tangential) distortion to normalized coords."""
+    k1, k2, p1, p2, k3 = d
+    r2 = x * x + y * y
+    radial = 1.0 + k1 * r2 + k2 * r2 * r2 + k3 * r2 * r2 * r2
+    xd = x * radial + 2.0 * p1 * x * y + p2 * (r2 + 2.0 * x * x)
+    yd = y * radial + p1 * (r2 + 2.0 * y * y) + 2.0 * p2 * x * y
+    return xd, yd
+
+
+class LidarCameraProjection(FramePreprocessor):
+    """Per-point pixel coordinates of a lidar scan in a camera image.
+
+    Output is ``(N, 3)`` float32, row-aligned with the input scan:
+    ``[u, v, depth]`` — continuous pixel coordinates (``u`` = column, ``v`` =
+    row, origin at the top-left pixel corner) and depth along the camera's
+    optical axis (metres).  Points behind the camera or projecting outside
+    ``image_size`` are ``NaN`` in all three columns, so validity is
+    ``np.isfinite(uv[:, 0])`` and ``int(u)``/``int(v)`` of a valid row always
+    index into the image.
+
+    Only the scan itself is streamed (``input_keys`` has a single channel):
+    extrinsics, intrinsics and image size are static rig properties, so the
+    preprocessor runs on asynchronous datasets as well as profiled ones.
+
+    Args:
+        intrinsics:  A :class:`~apairo.core.config.CameraIntrinsics` from
+                     ``ds.calibration.get_intrinsics(frame)``, or a bare
+                     ``(3, 3)`` pinhole camera matrix K.
+        extrinsics:  ``(4, 4)`` rigid transform ``T_camera_from_lidar`` — maps
+                     lidar-frame points into the camera *optical* frame
+                     (+Z forward).  Typically
+                     ``ds.calibration.get_tf(lidar_frame, camera_frame)``.
+        image_size:  ``(height, width)`` of the target image, in pixels.
+                     Defaults to the ``CameraIntrinsics``' recorded size;
+                     required when passing a bare K.
+        distortion:  Plumb-bob coefficients ``(k1, k2, p1, p2[, k3])``.
+                     Defaults to the ``CameraIntrinsics``' coefficients;
+                     ``None`` with a bare K means a rectified image.
+        lidar_key:   Input channel for point cloud data.
+        output_key:  Override the default output channel name ``"lidar_uv"``.
+    """
+
+    output_key: ClassVar[str] = "lidar_uv"
+    input_keys: ClassVar[list[str]] = ["lidar"]
+    timestamps_from: ClassVar[str] = "lidar"
+    sources: ClassVar[list[str]] = ["lidar"]
+
+    def __init__(
+        self,
+        intrinsics: np.ndarray | CameraIntrinsics,
+        extrinsics: np.ndarray,
+        image_size: tuple[int, int] | None = None,
+        distortion: np.ndarray | None = None,
+        lidar_key: str = "lidar",
+        output_key: str | None = None,
+    ) -> None:
+        if isinstance(intrinsics, CameraIntrinsics):
+            if intrinsics.model != "plumb_bob" and len(intrinsics.distortion):
+                raise ValueError(
+                    f"only the plumb_bob distortion model is supported, "
+                    f"got {intrinsics.model!r}"
+                )
+            if distortion is None and len(intrinsics.distortion):
+                distortion = intrinsics.distortion
+            if image_size is None:
+                if intrinsics.height is None or intrinsics.width is None:
+                    raise ValueError(
+                        "CameraIntrinsics records no image size — pass image_size="
+                    )
+                image_size = (intrinsics.height, intrinsics.width)
+            intrinsics = intrinsics.K
+        elif image_size is None:
+            raise ValueError("image_size is required when intrinsics is a bare K")
+        K = np.asarray(intrinsics, dtype=np.float64)
+        if K.shape != (3, 3):
+            raise ValueError(f"intrinsics must be (3, 3), got {K.shape}")
+        T = np.asarray(extrinsics, dtype=np.float64)
+        if T.shape != (4, 4):
+            raise ValueError(f"extrinsics must be (4, 4), got {T.shape}")
+        height, width = image_size
+        if height <= 0 or width <= 0:
+            raise ValueError(f"image_size must be positive, got {image_size}")
+        if distortion is not None:
+            d = np.asarray(distortion, dtype=np.float64).ravel()
+            if d.size == 4:
+                d = np.append(d, 0.0)  # k3 defaults to 0
+            if d.size != 5:
+                raise ValueError(
+                    f"distortion must be (k1, k2, p1, p2[, k3]), got {d.size} values"
+                )
+            self._dist = d
+        else:
+            self._dist = None
+
+        self._K = K
+        self._T = T
+        self._height = int(height)
+        self._width = int(width)
+        self._lidar_key = lidar_key
+
+        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])
+        xyz = pc[:, :3].astype(np.float64)
+        n = len(xyz)
+
+        cam = xyz @ self._T[:3, :3].T + self._T[:3, 3]
+        z = cam[:, 2]
+
+        out = np.full((n, 3), np.nan, dtype=np.float32)
+        front = z > 1e-9
+        if not front.any():
+            return out
+
+        x = cam[front, 0] / z[front]
+        y = cam[front, 1] / z[front]
+        if self._dist is not None:
+            x, y = _distort_plumb_bob(x, y, self._dist)
+
+        u = self._K[0, 0] * x + self._K[0, 1] * y + self._K[0, 2]
+        v = self._K[1, 1] * y + self._K[1, 2]
+
+        in_image = (u >= 0) & (u < self._width) & (v >= 0) & (v < self._height)
+        idx = np.flatnonzero(front)[in_image]
+        out[idx, 0] = u[in_image]
+        out[idx, 1] = v[in_image]
+        out[idx, 2] = z[idx]
+        return out

apairo_preprocess/projection/point_features.py

@@ -0,0 +1,100 @@
+"""Sample image values at projected lidar points — image → point cloud.
+
+The inverse direction of :class:`ImageMaskFromPointLabels`: given the
+per-point pixel coordinates written by :class:`LidarCameraProjection`, read
+the image at each valid point and attach the pixel's value (RGB, semantic
+class, …) to the point.  Output is row-aligned with the scan, so it loads as
+a per-point feature channel next to ``lidar``.
+
+Colouring the cloud with camera RGB and inspecting it in a 3-D viewer
+(projector, apairo_rr) is the standard visual check that extrinsics and
+intrinsics are right — miscalibration shows up immediately as colour bleeding
+across object boundaries.
+
+This preprocessor reads **two** channels on different clocks (the projection
+on the lidar clock, the image on the camera clock).  On profiled synchronous
+datasets it runs via ``run_preprocess`` directly; on asynchronous datasets,
+run it over a ``synchronize()`` view and persist with ``ChannelWriter`` (see
+``examples/traversability_image_mask.py``).
+
+Typical usage::
+
+    Rellis3DDataset.run_preprocess(
+        PointFeaturesFromImage(image_key="image_left", output_key="lidar_rgb"),
+        root,
+    )
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import numpy as np
+
+from apairo.core.preprocessor import FramePreprocessor
+from apairo.core.sample import Sample
+
+
+class PointFeaturesFromImage(FramePreprocessor):
+    """Per-point image values sampled at each point's projected pixel.
+
+    Output is ``(N, C)`` in the image's dtype, row-aligned with the scan
+    (``C`` = image channels; a single-channel ``(H, W)`` image yields
+    ``(N, 1)``).  Points without a valid projection get ``fill_value``.
+    Sampling is nearest-pixel (truncation of the continuous coordinates).
+
+    Args:
+        projection_key: Input channel with ``(N, 3)`` ``[u, v, depth]`` rows
+                        from :class:`LidarCameraProjection`.
+        image_key:      Input channel for the image, ``(H, W)`` or ``(H, W, C)``.
+        fill_value:     Value written for points with no valid projection.
+        output_key:     Override the default output channel name
+                        ``"point_features"``.
+    """
+
+    output_key: ClassVar[str] = "point_features"
+    input_keys: ClassVar[list[str]] = ["lidar_uv", "image"]
+    timestamps_from: ClassVar[str] = "lidar_uv"
+    sources: ClassVar[list[str]] = ["lidar_uv", "image"]
+
+    def __init__(
+        self,
+        projection_key: str = "lidar_uv",
+        image_key: str = "image",
+        fill_value: float = 0,
+        output_key: str | None = None,
+    ) -> None:
+        self._projection_key = projection_key
+        self._image_key = image_key
+        self._fill_value = fill_value
+
+        self.input_keys = [projection_key, image_key]
+        self.timestamps_from = projection_key
+        self.sources = [projection_key, image_key]
+        if output_key is not None:
+            self.output_key = output_key
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        uv = np.asarray(sample.data[self._projection_key])
+        img = np.asarray(sample.data[self._image_key])
+        if uv.ndim != 2 or uv.shape[1] < 2:
+            raise ValueError(f"projection must be (N, >=2) [u, v, ...], got {uv.shape}")
+        if img.ndim == 2:
+            img = img[:, :, None]
+        if img.ndim != 3:
+            raise ValueError(f"image must be (H, W) or (H, W, C), got {img.shape}")
+
+        n = len(uv)
+        height, width, channels = img.shape
+        out = np.full((n, channels), self._fill_value, dtype=img.dtype)
+
+        valid = np.isfinite(uv[:, 0])
+        cols = uv[valid, 0].astype(np.intp)
+        rows = uv[valid, 1].astype(np.intp)
+        if len(cols) and (cols.max() >= width or rows.max() >= height):
+            raise ValueError(
+                f"projection exceeds image bounds ({height}, {width}) — the "
+                f"projection was computed for a different image_size."
+            )
+        out[valid] = img[rows, cols]
+        return out

apairo_preprocess/traversability/__init__.py

@@ -0,0 +1,7 @@
+from apairo_preprocess.traversability.from_labels import TraversabilityFromLabels
+from apairo_preprocess.traversability.from_trajectory import TraversabilityFromTrajectory
+
+__all__ = [
+    "TraversabilityFromLabels",
+    "TraversabilityFromTrajectory",
+]

apairo_preprocess/traversability/from_labels.py

@@ -0,0 +1,61 @@
+"""Label-based traversability — maps semantic IDs to binary traversable/non-traversable.
+
+Each point is labelled 1 if its semantic class ID is in ``traversable_ids``, 0 otherwise.
+
+Default traversable IDs for RELLIS-3D::
+
+    {1: dirt, 3: grass, 10: asphalt, 23: concrete, 31: puddle, 33: mud}
+
+Typical usage::
+
+    Rellis3DDataset.run_preprocess(
+        TraversabilityFromLabels(),
+        "/data/rellis",
+    )
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import numpy as np
+
+from apairo.core.preprocessor import FramePreprocessor
+from apairo.core.sample import Sample
+
+_RELLIS_TRAVERSABLE_IDS: frozenset[int] = frozenset({1, 3, 10, 23, 31, 33})
+
+
+class TraversabilityFromLabels(FramePreprocessor):
+    """Label each point traversable based on its semantic class ID.
+
+    Args:
+        labels_key:      Input channel for per-point semantic labels.
+        traversable_ids: Set of semantic class IDs considered traversable.
+                         Defaults to the RELLIS-3D traversable classes.
+        output_key:      Override the default output channel name ``"trav_label"``.
+    """
+
+    output_key: ClassVar[str] = "trav_label"
+    input_keys: ClassVar[list[str]] = ["labels"]
+    timestamps_from: ClassVar[str] = "lidar"
+    sources: ClassVar[list[str]] = ["labels"]
+
+    def __init__(
+        self,
+        labels_key: str = "labels",
+        traversable_ids: frozenset[int] | None = None,
+        output_key: str | None = None,
+    ) -> None:
+        self._labels_key = labels_key
+        self._trav_ids = (
+            traversable_ids if traversable_ids is not None else _RELLIS_TRAVERSABLE_IDS
+        )
+        self.input_keys = [labels_key]
+        self.sources = [labels_key]
+        if output_key is not None:
+            self.output_key = output_key
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        labels = np.asarray(sample.data[self._labels_key])
+        return np.isin(labels, list(self._trav_ids)).astype(np.uint8)