feat(common): load LiDAR frames from Parquet, PCD or bin by their own layout

t4_devkit/common/io.py

@@ -1,14 +1,19 @@
 from __future__ import annotations
 
 import json
-from typing import Any
+import os
+from functools import lru_cache
+from typing import Any, Sequence
 
+import numpy as np
+import pyarrow.parquet as pq
+import pypcd4
 from typing_extensions import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from t4_devkit.typing import PathLike
 
-__all__ = ("load_json", "save_json")
+__all__ = ("load_json", "save_json", "load_pointcloud")
 
 
 def load_json(filename: PathLike) -> Any:
@@ -34,3 +39,64 @@ def save_json(data: Any, filename: PathLike) -> None:
     """
     with open(filename, "w") as f:
         json.dump(data, f, ensure_ascii=False, indent=4)
+
+
+def load_pointcloud(
+    filepath: PathLike,
+    *,
+    frame_index: int | None = None,
+    fields: Sequence[tuple[str, str]] | None = None,
+    columns: Sequence[str] | None = None,
+) -> dict[str, np.ndarray]:
+    """Load one pointcloud frame as one array per point field, in the dtype stored in the file.
+
+    The layout is read from the file itself (`.parquet`, `.pcd`), or from `fields` for `.bin`,
+    which has no header. Compression and encodings are read from the file.
+
+    Args:
+        filepath (PathLike): `.parquet` file of a channel, or `.pcd` / `.bin` file of one frame.
+        frame_index (int | None, optional): Frame to read from a `.parquet` file (one row group per frame).
+        fields (Sequence[tuple[str, str]] | None, optional): (name, dtype) of each point field
+            of a `.bin` file, in order, e.g. from `sensor.lidar.fields`.
+        columns (Sequence[str] | None, optional): Fields to return; all fields if None.
+
+    Returns:
+        Arrays of shape (N,) keyed by field name, in file order. A `.parquet` column that is
+        null on every point of the frame is left out.
+    """
+    filepath = os.fspath(filepath)
+    if filepath.endswith(".parquet"):
+        if frame_index is None:
+            raise ValueError("frame_index is required for a .parquet file")
+        mtime = os.stat(filepath).st_mtime_ns
+        frame = _open_parquet(filepath, mtime, os.getpid()).read_row_group(
+            frame_index, columns=columns
+        )
+        # a column null on every point (e.g. point labels of a frame without labels) is left out
+        return {
+            name: frame[name].to_numpy()
+            for name in frame.column_names
+            if frame.num_rows == 0 or frame[name].null_count < frame.num_rows
+        }
+
+    if filepath.endswith(".pcd"):
+        data = pypcd4.PointCloud.from_path(filepath).pc_data
+    elif filepath.endswith(".bin"):
+        if fields is None:
+            raise ValueError("fields is required for a .bin file")
+        data = np.fromfile(filepath, dtype=np.dtype(list(fields)))
+    else:
+        raise ValueError(f"Unexpected filetype: {filepath}")
+
+    names = data.dtype.names if columns is None else columns
+    return {name: data[name] for name in names}
+
+
+@lru_cache(maxsize=32)
+def _open_parquet(filepath: str, mtime: int, pid: int) -> pq.ParquetFile:
+    """Open a Parquet file once per process and file version, memory-mapped.
+
+    Reopening per frame re-reads the footer. `mtime` and `pid` are only part of the cache key,
+    so a rewritten file or a forked worker gets its own handle.
+    """
+    return pq.ParquetFile(filepath, memory_map=True)

t4_devkit/common/timestamp.py

@@ -2,13 +2,14 @@
 
 from typing import TYPE_CHECKING
 
+from attrs import define
 from typing_extensions import deprecated
 
 if TYPE_CHECKING:
     from t4_devkit.typing import ScalarLike
 
 
-__all__ = ["microseconds2seconds", "seconds2microseconds"]
+__all__ = ["microseconds2seconds", "seconds2microseconds", "Stamp"]
 
 
 @deprecated("Please use `microseconds2seconds(...)` instead.")
@@ -65,3 +66,25 @@ def seconds2microseconds(timestamp: ScalarLike) -> float:
         Timestamp in [us].
     """
     return 1e6 * timestamp
+
+
+@define
+class Stamp:
+    """A dataclass to represent timestamp.
+
+    Attributes:
+        sec (int): Seconds.
+        nanosec (int): Nanoseconds.
+    """
+
+    sec: int
+    nanosec: int
+
+    @property
+    def in_seconds(self) -> float:
+        """Convert timestamp to seconds as a float.
+
+        Returns:
+            float: Timestamp in seconds.
+        """
+        return self.sec + self.nanosec * 1e-9

t4_devkit/dataclass/pointcloud.py

@@ -2,14 +2,15 @@
 
 import os.path as osp
 import struct
+import warnings
 from abc import abstractmethod
-from typing import TYPE_CHECKING, ClassVar, TypeVar
+from typing import TYPE_CHECKING, ClassVar, Sequence, TypeVar
 
 import numpy as np
-import pypcd4
 from attrs import define, field
 
-from t4_devkit.common.io import load_json
+from t4_devkit.common.io import load_json, load_pointcloud
+from t4_devkit.common.timestamp import Stamp
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -27,27 +28,11 @@
     "Stamp",
 ]
 
+# Point fields of `LidarPointCloud.points`, in row order.
+XYZI = ("x", "y", "z", "intensity")
 
-@define
-class Stamp:
-    """A dataclass to represent timestamp.
-
-    Attributes:
-        sec (int): Seconds.
-        nanosec (int): Nanoseconds.
-    """
-
-    sec: int
-    nanosec: int
-
-    @property
-    def in_seconds(self) -> float:
-        """Convert timestamp to seconds as a float.
-
-        Returns:
-            float: Timestamp in seconds.
-        """
-        return self.sec + self.nanosec * 1e-9
+# 0.x `.bin` layout: this list cut to `num_pts_feats`, all float32 (as written by the converter).
+V0_BIN_FIELDS = ("x", "y", "z", "intensity", "ring", "return_type", "time_stamp")
 
 
 @define
@@ -261,30 +246,52 @@ def num_dims() -> int:
         return 4
 
     @classmethod
-    def from_file(cls, filepath: str, metainfo_filepath: str | None = None) -> Self:
-        assert filepath.endswith((".bin", ".pcd")), f"Unexpected filetype: {filepath}"
+    def from_file(
+        cls,
+        filepath: str,
+        metainfo_filepath: str | None = None,
+        frame_index: int | None = None,
+        fields: Sequence[tuple[str, str]] | None = None,
+    ) -> Self:
+        """Create an object from a pointcloud file.
+
+        Args:
+            filepath (str): `.parquet` file of a channel, or `.pcd` / `.bin` file of one frame.
+            metainfo_filepath (str | None, optional): Metainfo file of a 0.x frame.
+            frame_index (int | None, optional): Frame to read from a `.parquet` file.
+            fields (Sequence[tuple[str, str]] | None, optional): (name, dtype) of each point field
+                of a `.bin` file, e.g. from `sensor.lidar.fields`.
 
+        Returns:
+            Self instance.
+        """
         metainfo = (
             PointCloudMetainfo.from_file(metainfo_filepath)
             if metainfo_filepath is not None and osp.exists(metainfo_filepath)
             else None
         )
 
-        if filepath.endswith(".bin"):
+        if fields is None and filepath.endswith(".bin"):
             num_pts_feats = getattr(metainfo, "num_pts_feats", 5)
-            scan = np.fromfile(filepath, dtype=np.float32)
-            points = scan.reshape((-1, num_pts_feats))[:, : cls.num_dims()]
-        else:
-            pcd = pypcd4.PointCloud.from_path(filepath)
-            points = pcd.numpy(("x", "y", "z", "intensity")).astype(np.float32)
-
-        return cls(points.T, metainfo=metainfo)
+            if num_pts_feats > len(V0_BIN_FIELDS):
+                raise ValueError(f"Unknown 0.x point layout with {num_pts_feats} fields")
+            fields = [(name, "float32") for name in V0_BIN_FIELDS[:num_pts_feats]]
+
+        columns = load_pointcloud(filepath, frame_index=frame_index, fields=fields, columns=XYZI)
+        points = np.empty((len(XYZI), len(columns["x"])), dtype=np.float32)
+        for i, name in enumerate(XYZI):
+            points[i] = columns[name]
+        return cls(points, metainfo=metainfo)
 
 
 @define
 class RadarPointCloud(PointCloud):
     """A dataclass to represent radar pointcloud.
 
+    Placeholder carried over from nuScenes: `from_file` reads the nuScenes radar PCD layout
+    (18 fields). tier4_perception_dataset stores radar as `RadarTracks` JSON, which it does
+    not read, and radar is not part of 1.0 yet.
+
     Attributes:
         points (NDArrayFloat): Points matrix in the shape of (18, N).
     """
@@ -307,6 +314,10 @@ def from_file(
         ambig_states: list[int] | None = None,
         metainfo_filepath: str | None = None,
     ) -> Self:
+        warnings.warn(
+            "T4 radar is not implemented: RadarPointCloud reads the nuScenes radar PCD layout only.",
+            stacklevel=2,
+        )
         assert filepath.endswith(".pcd"), f"Unexpected filetype: {filepath}"
 
         metadata: list[str] = []
@@ -399,6 +410,11 @@ def from_file(
 class SegmentationPointCloud(PointCloud):
     """A dataclass to represent segmentation pointcloud.
 
+    Deprecated:
+        Kept only for 0.x datasets, whose labels are a separate `.bin` file. In 1.0, labels are
+        columns of the LiDAR frame and come with `load_pointcloud()`. To be removed once
+        `T4Devkit` moves to 1.0.
+
     Attributes:
         points (NDArrayFloat): Points matrix in the shape of (4, N).
         labels (NDArrayU8): Label matrix.
@@ -414,24 +430,9 @@ def num_dims() -> int:
     def from_file(
         cls, point_filepath: str, label_filepath: str, metainfo_filepath: str | None = None
     ) -> Self:
-        assert point_filepath.endswith((".bin", ".pcd")), f"Unexpected filetype: {point_filepath}"
-
-        metainfo = (
-            PointCloudMetainfo.from_file(metainfo_filepath)
-            if metainfo_filepath is not None and osp.exists(metainfo_filepath)
-            else None
-        )
-
-        if point_filepath.endswith(".bin"):
-            num_pts_feats = getattr(metainfo, "num_pts_feats", 5)
-            scan = np.fromfile(point_filepath, dtype=np.float32)
-            points = scan.reshape((-1, num_pts_feats))[:, : cls.num_dims()]
-        else:
-            pcd = pypcd4.PointCloud.from_path(point_filepath)
-            points = pcd.numpy(("x", "y", "z", "intensity")).astype(np.float32)
-
+        pointcloud = LidarPointCloud.from_file(point_filepath, metainfo_filepath)
         labels = np.fromfile(label_filepath, dtype=np.uint8)
-        return cls(points.T, labels=labels, metainfo=metainfo)
+        return cls(pointcloud.points, labels=labels, metainfo=pointcloud.metainfo)
 
     def split_by_sensor(self) -> dict[str, PointCloud]:
         if self.metainfo is None:

tests/common/test_io.py

@@ -1,6 +1,15 @@
 import tempfile
+from pathlib import Path
 
-from t4_devkit.common.io import load_json, save_json
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pypcd4
+import pytest
+
+from t4_devkit.common.io import load_json, load_pointcloud, save_json
+
+SAMPLE_PARQUET = Path(__file__).parents[1] / "sample/t4dataset_v1/data/LIDAR_TOP/LIDAR_TOP.parquet"
 
 
 def test_json_io() -> None:
@@ -13,3 +22,54 @@ def test_json_io() -> None:
         ret = load_json(f.name)
 
         assert ret == data
+
+
+def test_load_pointcloud_parquet() -> None:
+    """All fields come back in the dtype stored in the file; `columns` selects fields."""
+    frame = load_pointcloud(SAMPLE_PARQUET, frame_index=1)
+    assert list(frame) == ["x", "y", "z", "intensity", "ring"]
+    assert frame["ring"].dtype == np.uint16
+    assert len(frame["x"]) == 8
+
+    assert list(load_pointcloud(SAMPLE_PARQUET, frame_index=0, columns=["x", "z"])) == ["x", "z"]
+
+    with pytest.raises(ValueError, match="frame_index"):
+        load_pointcloud(SAMPLE_PARQUET)
+
+
+def test_load_pointcloud_parquet_null_column(tmp_path: Path) -> None:
+    """A column null on every point of a frame (a frame without labels) is left out."""
+    schema = pa.schema([("x", pa.float32()), ("label", pa.uint8())])
+    path = tmp_path / "LIDAR_TOP.parquet"
+    with pq.ParquetWriter(path, schema) as writer:
+        writer.write_table(pa.table([pa.array([1.0, 2.0]), pa.array([0, 3])], schema=schema))
+        writer.write_table(pa.table([pa.array([1.0, 2.0]), pa.nulls(2, pa.uint8())], schema=schema))
+
+    labeled = load_pointcloud(path, frame_index=0)
+    assert labeled["label"].tolist() == [0, 3]
+    assert labeled["label"].dtype == np.uint8
+    assert list(load_pointcloud(path, frame_index=1)) == ["x"]
+
+
+def test_load_pointcloud_bin(tmp_path: Path) -> None:
+    """A .bin file is read with the layout given in `fields`."""
+    fields = [("x", "float32"), ("y", "float32"), ("z", "float32"), ("ring", "uint16")]
+    data = np.array([(1, 2, 3, 7), (4, 5, 6, 8)], dtype=fields)
+    data.tofile(tmp_path / "frame.bin")
+
+    frame = load_pointcloud(tmp_path / "frame.bin", fields=fields)
+    assert frame["ring"].tolist() == [7, 8]
+    assert frame["z"].dtype == np.float32
+
+    with pytest.raises(ValueError, match="fields"):
+        load_pointcloud(tmp_path / "frame.bin")
+
+
+def test_load_pointcloud_pcd(tmp_path: Path) -> None:
+    """A .pcd file is read with the layout in its header."""
+    points = np.arange(10, dtype=np.float32).reshape(2, 5)
+    pypcd4.PointCloud.from_xyzir_points(points).save(tmp_path / "frame.pcd")
+
+    frame = load_pointcloud(tmp_path / "frame.pcd")
+    assert list(frame) == ["x", "y", "z", "intensity", "ring"]
+    assert frame["ring"].dtype == np.uint16