apairo-preprocess 0.2.0__py3-none-any.whl

apairo_preprocess/__init__.py

@@ -0,0 +1,27 @@
+from apairo_preprocess.traversability import TraversabilityFromLabels, TraversabilityFromTrajectory
+from apairo_preprocess.ground import GroundSegmentationCSF, GroundSegmentationFromLabels, GroundSegmentationRANSAC
+from apairo_preprocess.priors import GroundHeightFromLabels, TrajectoryDistance
+from apairo_preprocess.odometry import GICPOdometry, KissICPOdometry
+from apairo_preprocess.pointcloud import VoxeliseCoords, VoxeliseLabels, VoxelisePointCloud, RemoveRobotPoints
+from apairo_preprocess.projection import BEVRasterisation, ImageMaskFromPointLabels, LidarCameraProjection, PointFeaturesFromImage, to_uint8_image
+
+__all__ = [
+    "TraversabilityFromLabels",
+    "TraversabilityFromTrajectory",
+    "GroundSegmentationCSF",
+    "GroundSegmentationFromLabels",
+    "GroundSegmentationRANSAC",
+    "GroundHeightFromLabels",
+    "TrajectoryDistance",
+    "GICPOdometry",
+    "KissICPOdometry",
+    "VoxelisePointCloud",
+    "VoxeliseLabels",
+    "VoxeliseCoords",
+    "RemoveRobotPoints",
+    "LidarCameraProjection",
+    "PointFeaturesFromImage",
+    "ImageMaskFromPointLabels",
+    "BEVRasterisation",
+    "to_uint8_image",
+]

apairo_preprocess/ground/__init__.py

@@ -0,0 +1,9 @@
+from apairo_preprocess.ground.csf import GroundSegmentationCSF
+from apairo_preprocess.ground.from_labels import GroundSegmentationFromLabels
+from apairo_preprocess.ground.ransac import GroundSegmentationRANSAC
+
+__all__ = [
+    "GroundSegmentationCSF",
+    "GroundSegmentationFromLabels",
+    "GroundSegmentationRANSAC",
+]

apairo_preprocess/ground/csf.py

@@ -0,0 +1,95 @@
+"""Ground segmentation via Cloth Simulation Filter (CSF)."""
+
+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 _run_csf(
+    xyz: np.ndarray,
+    cloth_resolution: float,
+    class_threshold: float,
+    rigidness: int,
+) -> tuple[list[int], list[int]]:
+    """Run CSF and return (ground_indices, non_ground_indices).
+
+    Requires ``pip install CSF``.
+    """
+    try:
+        import CSF  # type: ignore[import-untyped]
+    except ImportError as exc:
+        raise ImportError(
+            "GroundSegmentationCSF requires the CSF package: pip install CSF"
+        ) from exc
+
+    csf = CSF.CSF()
+    csf.params.bSloopSmooth = False
+    csf.params.cloth_resolution = cloth_resolution
+    csf.params.rigidness = rigidness
+    csf.params.time_step = 0.65
+    csf.params.class_threshold = class_threshold
+    csf.params.interations = 500
+
+    csf.setPointCloud(xyz.tolist())
+    ground_idx = CSF.VecInt()
+    non_ground_idx = CSF.VecInt()
+    csf.do_filtering(ground_idx, non_ground_idx)
+    return list(ground_idx), list(non_ground_idx)
+
+
+class GroundSegmentationCSF(FramePreprocessor):
+    """Per-point ground/non-ground labels via Cloth Simulation Filter.
+
+    Runs CSF on the voxelised point cloud.  More accurate than RANSAC on
+    uneven terrain; requires ``pip install CSF``.
+
+    Output labels: **0 = ground, 1 = non-ground**.
+
+    Args:
+        voxelised_key:    Input channel (voxelised point cloud, shape (N, D≥3)).
+        cloth_resolution: CSF cloth resolution (metres).  Default 0.5 m.
+        class_threshold:  CSF point-to-cloth distance threshold (metres).
+                          Default 0.5 m.
+        rigidness:        CSF cloth rigidness (1–3).  3 = flat terrain.
+                          Default 3.
+        output_key:       Override the default channel name ``"ground_csf"``.
+    """
+
+    output_key: ClassVar[str] = "ground_csf"
+    input_keys: ClassVar[list[str]] = ["voxelised"]
+    timestamps_from: ClassVar[str] = "voxelised"
+    sources: ClassVar[list[str]] = ["voxelised"]
+
+    def __init__(
+        self,
+        voxelised_key: str = "voxelised",
+        cloth_resolution: float = 0.5,
+        class_threshold: float = 0.5,
+        rigidness: int = 3,
+        output_key: str | None = None,
+    ) -> None:
+        self._voxelised_key = voxelised_key
+        self._cloth_resolution = cloth_resolution
+        self._class_threshold = class_threshold
+        self._rigidness = rigidness
+        self.input_keys = [voxelised_key]
+        self.sources = [voxelised_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]
+        ground_idx, _ = _run_csf(
+            xyz,
+            cloth_resolution=self._cloth_resolution,
+            class_threshold=self._class_threshold,
+            rigidness=self._rigidness,
+        )
+        labels = np.ones(len(xyz), dtype=np.uint8)
+        labels[ground_idx] = 0
+        return labels

apairo_preprocess/ground/from_labels.py

@@ -0,0 +1,55 @@
+"""Ground segmentation from semantic class IDs."""
+
+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_GROUND_IDS: frozenset[int] = frozenset({1, 3, 10, 23, 31, 33})
+
+
+class GroundSegmentationFromLabels(FramePreprocessor):
+    """Per-point ground/non-ground labels from semantic class IDs.
+
+    Each point is labelled ground (0) if its semantic class ID is in
+    ``ground_ids``, non-ground (1) otherwise.
+
+    Output labels: **0 = ground, 1 = non-ground** — same convention as
+    ``GroundSegmentationCSF`` and ``GroundSegmentationRANSAC``.
+
+    Default ground IDs for RELLIS-3D::
+
+        {1: dirt, 3: grass, 10: asphalt, 23: concrete, 31: puddle, 33: mud}
+
+    Args:
+        labels_key: Input channel for per-point semantic labels.
+        ground_ids: Set of semantic class IDs considered ground.
+                    Defaults to the RELLIS-3D ground classes.
+        output_key: Override the default channel name ``"ground_labels"``.
+    """
+
+    output_key: ClassVar[str] = "ground_labels"
+    input_keys: ClassVar[list[str]] = ["labels"]
+    timestamps_from: ClassVar[str] = "lidar"
+    sources: ClassVar[list[str]] = ["labels"]
+
+    def __init__(
+        self,
+        labels_key: str = "labels",
+        ground_ids: frozenset[int] | None = None,
+        output_key: str | None = None,
+    ) -> None:
+        self._labels_key = labels_key
+        self._ground_ids = ground_ids if ground_ids is not None else _RELLIS_GROUND_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._ground_ids))).astype(np.uint8)

apairo_preprocess/ground/ransac.py

@@ -0,0 +1,95 @@
+"""Ground segmentation via RANSAC plane fitting."""
+
+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 _ransac_plane(
+    xyz: np.ndarray,
+    n_iter: int = 200,
+    inlier_threshold: float = 0.15,
+    seed: int = 0,
+) -> tuple[np.ndarray, float]:
+    """Fit a ground plane with RANSAC.
+
+    Returns (normal, d) for the plane equation n · x + d = 0,
+    with the normal oriented upward (positive z component).
+    """
+    rng = np.random.default_rng(seed)
+    n = len(xyz)
+    best_inliers = -1
+    best_normal = np.array([0.0, 0.0, 1.0])
+    best_d = float(-np.median(xyz[:, 2]))
+
+    for _ in range(n_iter):
+        idx = rng.choice(n, 3, replace=False)
+        p0, p1, p2 = xyz[idx]
+        normal = np.cross(p1 - p0, p2 - p0)
+        norm = np.linalg.norm(normal)
+        if norm < 1e-8:
+            continue
+        normal = normal / norm
+        if normal[2] < 0:
+            normal = -normal
+        d = -float(np.dot(normal, p0))
+
+        dist = np.abs(xyz @ normal + d)
+        n_inl = int((dist < inlier_threshold).sum())
+        if n_inl > best_inliers:
+            best_inliers = n_inl
+            best_normal = normal
+            best_d = d
+
+    return best_normal, best_d
+
+
+class GroundSegmentationRANSAC(FramePreprocessor):
+    """Per-point ground/non-ground labels via RANSAC plane fitting.
+
+    Fits a plane to the voxelised point cloud with RANSAC.  Points whose
+    distance to the plane is within ``inlier_threshold`` are labelled ground
+    (0); all others are labelled non-ground (1).
+
+    Output labels: **0 = ground, 1 = non-ground**.
+
+    Args:
+        voxelised_key:     Input channel (voxelised point cloud, shape (N, D≥3)).
+        n_iter:            Number of RANSAC iterations.  Default 200.
+        inlier_threshold:  Distance threshold (metres) for a point to be ground.
+                           Default 0.15 m.
+        output_key:        Override the default channel name ``"ground_ransac"``.
+    """
+
+    output_key: ClassVar[str] = "ground_ransac"
+    input_keys: ClassVar[list[str]] = ["voxelised"]
+    timestamps_from: ClassVar[str] = "voxelised"
+    sources: ClassVar[list[str]] = ["voxelised"]
+
+    def __init__(
+        self,
+        voxelised_key: str = "voxelised",
+        n_iter: int = 200,
+        inlier_threshold: float = 0.15,
+        output_key: str | None = None,
+    ) -> None:
+        self._voxelised_key = voxelised_key
+        self._n_iter = n_iter
+        self._inlier_threshold = inlier_threshold
+        self.input_keys = [voxelised_key]
+        self.sources = [voxelised_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]
+        normal, d = _ransac_plane(
+            xyz, n_iter=self._n_iter, inlier_threshold=self._inlier_threshold
+        )
+        dist = np.abs(xyz @ normal + d)
+        return (dist > self._inlier_threshold).astype(np.uint8)

apairo_preprocess/odometry/__init__.py

@@ -0,0 +1,4 @@
+from apairo_preprocess.odometry.gicp import GICPOdometry
+from apairo_preprocess.odometry.kissicp import KissICPOdometry
+
+__all__ = ["GICPOdometry", "KissICPOdometry"]

apairo_preprocess/odometry/gicp.py

@@ -0,0 +1,96 @@
+"""Incremental scan-to-scan GICP pose estimation.
+
+Produces one (4, 4) float64 world pose per scan, accumulated via
+point-to-plane ICP.  The first scan is the world origin (identity).
+
+Cross-sequence boundaries (when the dataset mixes multiple recording
+sessions) the accumulated drift will jump, but for single-session
+splits the poses are consistent.
+
+Output channel: ``gicp_poses``  (npys — one .npy per scan)
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+import numpy as np
+
+from apairo.core.preprocessor import FramePreprocessor
+from apairo.core.sample import Sample
+
+try:
+    import open3d as o3d
+    _O3D_OK = True
+except ImportError:
+    _O3D_OK = False
+
+
+class GICPOdometry(FramePreprocessor):
+    """Stateful scan-to-scan GICP → one (4, 4) world pose per frame.
+
+    Uses Open3D Generalized ICP (point-to-plane).  Each call to
+    :meth:`process` registers the current scan against the previous one
+    and returns the accumulated world pose T_world_sensor as a (4, 4)
+    float64 array.
+
+    Args:
+        lidar_key:  Input channel for point cloud data.
+        voxel_size: Down-sampling voxel size (metres).
+        max_range:  Range filter applied before registration (metres).
+        max_corr:   Maximum correspondence distance for ICP (metres).
+        output_key: Override the default output channel name ``"gicp_poses"``.
+    """
+
+    output_key: ClassVar[str] = "gicp_poses"
+    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.3,
+        max_range: float = 50.0,
+        max_corr: float = 1.0,
+        output_key: str | None = None,
+    ) -> None:
+        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
+        if not _O3D_OK:
+            raise ImportError("open3d is required for GICPOdometry.")
+        self._voxel_size = voxel_size
+        self._max_range = max_range
+        self._max_corr = max_corr
+        self._prev_pcd = None
+        self._T_accum = np.eye(4, dtype=np.float64)
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        pc = np.asarray(sample.data[self._lidar_key])
+        xyz = pc[:, :3].astype(np.float64)
+        xyz = xyz[np.linalg.norm(xyz, axis=1) < self._max_range]
+
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(xyz)
+        pcd = pcd.voxel_down_sample(self._voxel_size)
+        pcd.estimate_normals(
+            o3d.geometry.KDTreeSearchParamHybrid(radius=self._voxel_size * 3, max_nn=30)
+        )
+
+        if self._prev_pcd is not None:
+            result = o3d.pipelines.registration.registration_generalized_icp(
+                pcd,
+                self._prev_pcd,
+                self._max_corr,
+                np.eye(4),
+                o3d.pipelines.registration.TransformationEstimationForGeneralizedICP(),
+            )
+            # result.transformation: source → target  (current → prev)
+            # T_world_cur = T_world_prev @ T_cur_to_prev
+            self._T_accum = self._T_accum @ result.transformation
+
+        self._prev_pcd = pcd
+        return self._T_accum.copy()

apairo_preprocess/odometry/kissicp.py

@@ -0,0 +1,84 @@
+"""LiDAR odometry via KISS-ICP.
+
+Produces one (4, 4) float64 world pose per scan using KISS-ICP
+(https://github.com/PRBonn/kiss-icp).  Much faster than Open3D GICP (~100x)
+and handles motion deskewing natively when per-point timestamps are available.
+
+Output channel: ``kissicp_poses``  (npys — one .npy per scan)
+
+Typical usage::
+
+    dataset.run_preprocess(
+        KissICPOdometry(voxel_size=1.0),
+        split_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
+
+try:
+    from kiss_icp.config import KISSConfig
+    from kiss_icp.kiss_icp import KissICP
+    _KISS_OK = True
+except ImportError:
+    _KISS_OK = False
+
+
+class KissICPOdometry(FramePreprocessor):
+    """Stateful KISS-ICP odometry → one (4, 4) world pose per frame.
+
+    Each call to :meth:`process` registers the current scan and returns
+    the accumulated T_world_lidar pose.  The first scan is the world origin.
+
+    Args:
+        lidar_key:  Input channel for point cloud data.
+        max_range:  Maximum point range kept before registration (metres).
+        min_range:  Minimum point range kept (metres).
+        voxel_size: Map voxel size (metres).  Controls both resolution and speed.
+        deskew:     Enable motion deskewing.  Requires per-point timestamps in
+                    the lidar channel (4th column).
+        output_key: Override the default output channel name ``"kissicp_poses"``.
+    """
+
+    output_key: ClassVar[str] = "kissicp_poses"
+    input_keys: ClassVar[list[str]] = ["lidar"]
+    timestamps_from: ClassVar[str] = "lidar"
+    sources: ClassVar[list[str]] = ["lidar"]
+
+    def __init__(
+        self,
+        lidar_key: str = "lidar",
+        max_range: float = 50.0,
+        min_range: float = 1.0,
+        voxel_size: float = 1.0,
+        deskew: bool = False,
+        output_key: str | None = None,
+    ) -> None:
+        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
+        if not _KISS_OK:
+            raise ImportError("kiss-icp is required for KissICPOdometry.  pip install kiss-icp")
+        cfg = KISSConfig()
+        cfg.data.max_range = max_range
+        cfg.data.min_range = min_range
+        cfg.data.deskew = deskew
+        cfg.mapping.voxel_size = float(voxel_size)
+        self._kiss = KissICP(config=cfg)
+        self._deskew = deskew
+
+    def __call__(self, sample: Sample) -> np.ndarray:
+        pc = np.asarray(sample.data[self._lidar_key], dtype=np.float64)
+        xyz = pc[:, :3]
+        timestamps = pc[:, 3] if self._deskew and pc.shape[1] > 3 else np.zeros(len(xyz))
+        self._kiss.register_frame(xyz, timestamps)
+        return self._kiss.last_pose.copy()  # (4, 4) float64

apairo_preprocess/pointcloud/__init__.py

@@ -0,0 +1,4 @@
+from apairo_preprocess.pointcloud.voxelise import VoxeliseCoords, VoxeliseLabels, VoxelisePointCloud
+from apairo_preprocess.pointcloud.robot_filter import RemoveRobotPoints
+
+__all__ = ["VoxelisePointCloud", "VoxeliseLabels", "VoxeliseCoords", "RemoveRobotPoints"]

apairo_preprocess/pointcloud/robot_filter.py

@@ -0,0 +1,121 @@
+"""Filtering of LiDAR points that originate from the robot's own body.
+
+When the sensor is mounted on the robot, some returns will hit the robot
+chassis.  This preprocessor removes those self-hits by masking out points
+that fall inside a geometric approximation of the robot shape (axis-aligned
+box or upright cylinder), expressed in the sensor frame.
+
+Typical usage::
+
+    dataset.run_preprocess(
+        RemoveRobotPoints(
+            shape="box",
+            x_range=(-0.6, 0.6),
+            y_range=(-0.4, 0.4),
+            z_range=(-0.5, 0.2),
+        ),
+        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
+
+_Shape = Literal["box", "cylinder"]
+
+
+class RemoveRobotPoints(FramePreprocessor):
+    """Remove scan points that lie inside the robot's body volume.
+
+    Points whose coordinates (in the sensor frame) fall inside the given
+    geometric shape are discarded.  All other points and their extra channels
+    (intensity, ring, …) are preserved unchanged.
+
+    Args:
+        lidar_key:  Input channel for point cloud data.
+        shape:      Approximation geometry for the robot body:
+
+                    * ``"box"``      — axis-aligned bounding box defined by
+                      ``x_range``, ``y_range``, and ``z_range``.
+                    * ``"cylinder"`` — upright cylinder defined by ``radius``
+                      and ``z_range``.
+
+        x_range:    ``(min, max)`` extents along X (metres). Box only.
+        y_range:    ``(min, max)`` extents along Y (metres). Box only.
+        z_range:    ``(min, max)`` extents along Z (metres). Both shapes.
+        radius:     Radius of the cylinder in the XY plane (metres).
+                    Cylinder only.
+        output_key: Override the default output channel name
+                    ``"lidar_no_robot"``.
+    """
+
+    output_key: ClassVar[str] = "lidar_no_robot"
+    input_keys: ClassVar[list[str]] = ["lidar"]
+    timestamps_from: ClassVar[str] = "lidar"
+    sources: ClassVar[list[str]] = ["lidar"]
+
+    def __init__(
+        self,
+        lidar_key: str = "lidar",
+        shape: _Shape = "box",
+        x_range: tuple[float, float] = (-1.0, 1.0),
+        y_range: tuple[float, float] = (-1.0, 1.0),
+        z_range: tuple[float, float] = (-1.0, 1.0),
+        radius: float = 1.0,
+        output_key: str | None = None,
+    ) -> None:
+        if shape not in ("box", "cylinder"):
+            raise ValueError(f"shape must be 'box' or 'cylinder', got {shape!r}")
+        if shape == "box":
+            if x_range[0] >= x_range[1]:
+                raise ValueError(f"x_range must satisfy min < max, got {x_range}")
+            if y_range[0] >= y_range[1]:
+                raise ValueError(f"y_range must satisfy min < max, got {y_range}")
+        if z_range[0] >= z_range[1]:
+            raise ValueError(f"z_range must satisfy min < max, got {z_range}")
+        if shape == "cylinder" and radius <= 0:
+            raise ValueError(f"radius must be positive, got {radius}")
+
+        self._lidar_key = lidar_key
+        self._shape = shape
+        self._x_range = x_range
+        self._y_range = y_range
+        self._z_range = z_range
+        self._radius = float(radius)
+
+        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}")
+
+        xyz = pc[:, :3]
+
+        if self._shape == "box":
+            inside = (
+                (xyz[:, 0] >= self._x_range[0])
+                & (xyz[:, 0] <= self._x_range[1])
+                & (xyz[:, 1] >= self._y_range[0])
+                & (xyz[:, 1] <= self._y_range[1])
+                & (xyz[:, 2] >= self._z_range[0])
+                & (xyz[:, 2] <= self._z_range[1])
+            )
+        else:
+            dist_xy = np.linalg.norm(xyz[:, :2], axis=1)
+            inside = (
+                (dist_xy <= self._radius)
+                & (xyz[:, 2] >= self._z_range[0])
+                & (xyz[:, 2] <= self._z_range[1])
+            )
+
+        return pc[~inside]