feat(write): write T4Dataset 1.0 datasets

.pre-commit-config.yaml

@@ -32,6 +32,8 @@ repos:
     rev: v4.0.0-alpha.8
     hooks:
       - id: prettier
+        # written by DatasetWriter: manifest.json holds the hashes of these files
+        exclude: ^tests/sample/t4dataset_v1/
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.6.8

docs/schema/v1/index.md

@@ -132,6 +132,21 @@ Provenance {
 }
 ```
 
+## Writing 1.0 Datasets
+
+`DatasetWriter` writes a 1.0 dataset folder. Data files go first, then the tables and `manifest.json`:
+
+```python
+>>> from t4_devkit.schema.v1 import DatasetWriter
+>>> writer = DatasetWriter("<DATASET_ID>/<DATASET_VERSION>")  # must not exist yet
+>>> path, frame_count = writer.write_lidar("LIDAR_TOP", frames, fields)  # one Parquet file
+>>> video = writer.write_camera("CAM_FRONT", images, frame_rate=10)  # one AV1 video
+>>> writer.copy("<OLD_DATASET>/map", "map")  # files kept as they are
+>>> manifest = writer.save(tables, manifest_info)  # tables, then manifest.json
+```
+
+`save` lists every table and every file under `data/`, `map/` and `input_bag/` in `manifest.json`, with its hash. It computes `capabilities` from the tables, and adds t4-devkit to `generators`. `manifest_info` gives the other manifest fields (identity, provenance, migration notes).
+
 ## Loading 0.x Datasets
 
 `load_dataset()` reads a 0.x dataset (no `manifest.json`) by converting its tables to 1.0 records when loading, on a best-effort basis. Sensor data files stay as they are: one file per frame, and `sample_data.frame_index` is `null`.

docs/schema/v1/requirement.md

@@ -100,8 +100,8 @@ Tasks and the labels `REC023` / `REC024` expect:
 | `segmentation3d` | `label` / `instance` columns in a LiDAR file                   |
 | `keypoint`       | `Keypoint` records                                             |
 | `traffic_light`  | `TrafficLight` records, or instances with a lamp `map_element` |
-| `prediction`     | Not checked                                                    |
-| `planning`       | Not checked                                                    |
+| `prediction`     | Every `Sample` has a key frame (all frames annotated)          |
+| `planning`       | `Scenario` records                                             |
 
 ## Reference (`REF`)
 

t4_devkit/common/io.py

@@ -3,11 +3,13 @@
 import json
 import os
 import threading
+from fractions import Fraction
 from functools import lru_cache
-from typing import Any, Sequence
+from typing import Any, Iterable, Mapping, Sequence
 
 import av
 import numpy as np
+import pyarrow as pa
 import pyarrow.parquet as pq
 import pypcd4
 from PIL import Image
@@ -16,7 +18,15 @@
 if TYPE_CHECKING:
     from t4_devkit.typing import PathLike
 
-__all__ = ("load_json", "save_json", "load_pointcloud", "load_image", "V0_BIN_FIELDS")
+__all__ = (
+    "load_json",
+    "save_json",
+    "load_pointcloud",
+    "save_pointcloud",
+    "load_image",
+    "save_video",
+    "V0_BIN_FIELDS",
+)
 
 # 0.x `.bin` layout: this list cut to `num_pts_feats`, all float32 (as written by the converter).
 # `.bin` stores no names; these are the 1.0 names (0.x calls `channel` "ring"), all float32.
@@ -99,6 +109,68 @@ def load_pointcloud(
     return {name: data[name] for name in names}
 
 
+def save_pointcloud(
+    filepath: PathLike,
+    frames: Iterable[Mapping[str, np.ndarray]],
+    fields: Sequence[tuple[str, str]],
+    *,
+    compression: str = "zstd",
+    compression_level: int | None = 1,
+) -> int:
+    """Save the frames of a LiDAR channel as one Parquet file, one row group per frame.
+
+    Columns are `fields`, in order. Encodings are chosen by dtype, the smallest in the 1.0
+    benchmark that every tested reader (pyarrow, Polars, DuckDB) reads: floats BYTE_STREAM_SPLIT,
+    8-bit integers dictionary, wider integers (e.g. timestamps) DELTA_BINARY_PACKED. Readers do
+    not depend on them.
+
+    Args:
+        filepath (PathLike): `.parquet` file to write.
+        frames (Iterable[Mapping[str, np.ndarray]]): Arrays of shape (N,) keyed by field name,
+            one mapping per frame. A frame may have no points. A field that is missing or None
+            is null on every point of the frame, e.g. point labels of a frame without labels.
+        fields (Sequence[tuple[str, str]]): (name, dtype) of each column, e.g. from
+            `sensor.lidar.fields`.
+        compression (str, optional): Parquet compression codec.
+        compression_level (int | None, optional): Codec level; None for codecs without levels
+            (e.g. "snappy", "none").
+
+    Returns:
+        The number of frames written.
+    """
+    schema = pa.schema([(name, pa.from_numpy_dtype(np.dtype(dtype))) for name, dtype in fields])
+    encodings = {}
+    for field in schema:
+        if pa.types.is_floating(field.type):
+            encodings[field.name] = "BYTE_STREAM_SPLIT"
+        elif field.type.bit_width > 8:
+            encodings[field.name] = "DELTA_BINARY_PACKED"
+    num_frames = 0
+    with pq.ParquetWriter(
+        os.fspath(filepath),
+        schema,
+        compression=compression,
+        compression_level=compression_level,
+        use_dictionary=[field.name for field in schema if field.name not in encodings],
+        column_encoding=encodings,
+    ) as writer:
+        for frame in frames:
+            values = [frame.get(name) for name, _ in fields]
+            num_points = next((len(v) for v in values if v is not None), None)
+            if num_points is None:
+                raise ValueError(f"Frame {num_frames} has no field with values")
+            columns = [
+                pa.nulls(num_points, field.type)
+                if value is None
+                else pa.array(np.asarray(value, dtype=dtype), field.type)
+                for field, (_, dtype), value in zip(schema, fields, values)
+            ]
+            # one row group per frame, also for a frame without points
+            writer.write_table(pa.table(columns, schema=schema), row_group_size=max(num_points, 1))
+            num_frames += 1
+    return num_frames
+
+
 @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.
@@ -152,6 +224,96 @@ def load_image(filepath: PathLike, *, frame_index: int | None = None) -> np.ndar
     raise IndexError(f"Frame {frame_index} not found in {filepath}")
 
 
+def save_video(
+    filepath: PathLike,
+    frames: Iterable[np.ndarray],
+    frame_rate: int,
+    *,
+    preset: int = 6,
+    crf: int = 26,
+) -> dict[str, Any]:
+    """Save the frames of a camera channel as one AV1 MP4 video, with PyAV's SVT-AV1.
+
+    Every frame is a keyframe. The track timescale is the frame rate, so a frame's pts is its
+    `frame_index`. Colors are YUV 4:2:0, full range, BT.601.
+
+    Args:
+        filepath (PathLike): `.mp4` file to write.
+        frames (Iterable[np.ndarray]): RGB images of shape (H, W, 3), all the same size: uint8 for
+            8-bit video, uint16 for 10-bit video.
+        frame_rate (int): Frames per second.
+        preset (int, optional): SVT-AV1 preset (speed vs size).
+        crf (int, optional): SVT-AV1 quality (lower is better).
+
+    Returns:
+        The values of `sensor.camera.video`, except `path`.
+    """
+    os.environ.setdefault("SVT_LOG", "1")
+    container = av.open(
+        os.fspath(filepath),
+        mode="w",
+        options={"movflags": "+faststart", "video_track_timescale": str(frame_rate)},
+    )
+    stream = None
+    num_frames = 0
+    bit_depth, pixel_format = 8, "yuv420p"
+    with container:
+        for image in frames:
+            if stream is None:
+                # uint16 frames: 10-bit video
+                if image.dtype == np.uint16:
+                    bit_depth, pixel_format = 10, "yuv420p10le"
+                stream = _add_video_stream(container, image, frame_rate, pixel_format, preset, crf)
+            rgb = av.VideoFrame.from_ndarray(image, format="rgb48le" if bit_depth > 8 else "rgb24")
+            frame = rgb.reformat(
+                format=pixel_format,
+                dst_colorspace=av.video.reformatter.Colorspace.ITU601,
+                dst_color_range=av.video.reformatter.ColorRange.JPEG,
+            )
+            frame.pts = num_frames
+            container.mux(stream.encode(frame))
+            num_frames += 1
+        if stream is not None:
+            container.mux(stream.encode())
+    return {
+        "format": "av1_mp4",
+        "codec_profile": f"svt-av1 preset{preset} crf{crf} intra",
+        "pixel_format": f"{pixel_format} full-range bt601",
+        "bit_depth": bit_depth,
+        "frame_count": num_frames,
+        "timescale": frame_rate,
+    }
+
+
+def _add_video_stream(
+    container: av.container.OutputContainer,
+    image: np.ndarray,
+    frame_rate: int,
+    pixel_format: str,
+    preset: int,
+    crf: int,
+) -> av.video.stream.VideoStream:
+    """The AV1 stream of a camera video, sized by its first image."""
+    stream = container.add_stream("libsvtav1", rate=frame_rate)
+    stream.height, stream.width = image.shape[:2]
+    stream.pix_fmt = pixel_format
+    stream.time_base = Fraction(1, frame_rate)
+    context = stream.codec_context
+    context.gop_size = 1  # every frame a keyframe
+    context.options = {
+        "preset": str(preset),
+        "crf": str(crf),
+        # all intra, no scene detection or film grain, quiet
+        "svtav1-params": "pred-struct=0:scd=0:film-grain=0:fast-decode=2:lp=0:log-level=1",
+    }
+    # full range, BT.601 matrix, BT.709 primaries, sRGB transfer (FFmpeg enum values)
+    context.color_range = 2
+    context.colorspace = 5
+    context.color_primaries = 1
+    context.color_trc = 13
+    return stream
+
+
 @lru_cache(maxsize=32)
 def _open_video(filepath: str, mtime: int, pid: int, thread: int) -> av.container.InputContainer:
     """Open a video file once per process, thread and file version.

t4_devkit/sanity/record/base.py

@@ -6,8 +6,8 @@
 from returns.maybe import Maybe, Nothing, Some
 from returns.pipeline import is_successful
 
-from t4_devkit.schema import SchemaName, SensorModality
-from t4_devkit.schema.v1 import FileFormat, LidarContainerFormat, Task
+from t4_devkit.schema import SchemaName
+from t4_devkit.schema.v1 import Capabilities, FileFormat, LidarContainerFormat
 from t4_devkit.schema.v1 import SchemaName as V1SchemaName
 
 from ..checker import Checker
@@ -73,14 +73,6 @@ def channel_file(sensor: Sensor) -> tuple[str, int] | None:
     return None
 
 
-# `sample_data.fileformat` values for each LiDAR container format.
-LIDAR_FRAME_FORMATS = {
-    LidarContainerFormat.PARQUET: {FileFormat.PARQUET.value},
-    LidarContainerFormat.PCD: {FileFormat.PCD.value},
-    LidarContainerFormat.BIN: {FileFormat.BIN.value, FileFormat.PCDBIN.value},
-}
-
-
 def frame_formats(sensor: Sensor) -> set[str]:
     """Return the `sample_data.fileformat` values a sensor's storage allows (any for radar,
     undefined in 1.0)."""
@@ -89,25 +81,17 @@ def frame_formats(sensor: Sensor) -> set[str]:
             return {FileFormat.AV1_MP4.value}
         return {FileFormat.JPG.value, FileFormat.PNG.value}
     if sensor.lidar is not None:
-        return LIDAR_FRAME_FORMATS[sensor.lidar.container.format]
+        lidar_formats = {
+            LidarContainerFormat.PARQUET: {FileFormat.PARQUET.value},
+            LidarContainerFormat.PCD: {FileFormat.PCD.value},
+            LidarContainerFormat.BIN: {FileFormat.BIN.value, FileFormat.PCDBIN.value},
+        }
+        return lidar_formats[sensor.lidar.container.format]
     return {fileformat.value for fileformat in FileFormat}
 
 
-# Tasks whose labels are checked against `capabilities.tasks`.
-CHECKED_TASKS = (
-    Task.DETECTION3D,
-    Task.DETECTION2D,
-    Task.SEGMENTATION2D,
-    Task.SEGMENTATION3D,
-    Task.KEYPOINT,
-    Task.TRAFFIC_LIGHT,
-)
-
-
-def declared_capabilities(context: SanityContext) -> set[str]:
-    """Return what `manifest.json` `capabilities` declares, as items like "camera:CAM_FRONT", "map",
-    "task:detection3d"."""
-    capabilities = context.manifest.unwrap().capabilities
+def capability_items(capabilities: Capabilities) -> set[str]:
+    """Return capabilities as comparable items, e.g. "camera:CAM_FRONT", "map", "task:keypoint"."""
     items = {
         f"{modality}:{channel}"
         for modality in ("camera", "lidar", "radar")
@@ -118,48 +102,11 @@ def declared_capabilities(context: SanityContext) -> set[str]:
         for name in ("map", "route", "traffic_light", "vehicle_state", "scenario")
         if getattr(capabilities, name)
     }
-    items |= {f"task:{task.value}" for task in capabilities.tasks if task in CHECKED_TASKS}
+    items |= {f"task:{task.value}" for task in capabilities.tasks}
     return items
 
 
-def present_capabilities(context: SanityContext) -> set[str]:
-    """Return what the dataset contains, as the same items as `declared_capabilities`."""
-
-    def records(table: V1SchemaName) -> list:
-        return context.load_records(table).unwrap()
-
-    sensors = records(V1SchemaName.SENSOR)
-    with_frames = {frame.sensor_token for frame in records(V1SchemaName.SAMPLE_DATA)}
-    items = {
-        f"{sensor.modality.value}:{sensor.channel}"
-        for sensor in sensors
-        if sensor.token in with_frames
-    }
-    tables = {
-        "map": V1SchemaName.MAP,
-        "route": V1SchemaName.ROUTE,
-        "traffic_light": V1SchemaName.TRAFFIC_LIGHT,
-        "vehicle_state": V1SchemaName.VEHICLE_STATE,
-        "scenario": V1SchemaName.SCENARIO,
-    }
-    items |= {name for name, table in tables.items() if records(table)}
-
-    things = {category.token for category in records(V1SchemaName.CATEGORY) if category.is_thing}
-    object_anns = records(V1SchemaName.OBJECT_ANN)
-    lidar_columns = {
-        field.name
-        for sensor in sensors
-        if sensor.modality == SensorModality.LIDAR and sensor.lidar is not None
-        for field in sensor.lidar.fields
-    }
-    tasks = {
-        Task.DETECTION3D: bool(records(V1SchemaName.SAMPLE_ANNOTATION)),
-        Task.DETECTION2D: any(ann.category_token in things for ann in object_anns),
-        Task.SEGMENTATION2D: any(ann.mask is not None for ann in object_anns),
-        Task.SEGMENTATION3D: "label" in lidar_columns,
-        Task.KEYPOINT: bool(records(V1SchemaName.KEYPOINT)),
-        Task.TRAFFIC_LIGHT: bool(records(V1SchemaName.TRAFFIC_LIGHT))
-        or any(instance.map_element.id is not None for instance in records(V1SchemaName.INSTANCE)),
-    }
-    items |= {f"task:{task.value}" for task, present in tasks.items() if present}
-    return items
+def present_capabilities(context: SanityContext, tables: tuple[V1SchemaName, ...]) -> set[str]:
+    """Return what the dataset's tables contain, as `capability_items`."""
+    records = {table: context.load_records(table).unwrap() for table in tables}
+    return capability_items(Capabilities.from_tables(records))

t4_devkit/sanity/record/rec023.py

@@ -4,7 +4,7 @@
 
 from t4_devkit.schema.v1 import SchemaName
 
-from .base import declared_capabilities, present_capabilities
+from .base import capability_items, present_capabilities
 
 from ..checker import RecordsChecker, RuleID, RuleName, Severity
 from ..registry import CHECKERS
@@ -30,6 +30,7 @@ class REC023(RecordsChecker):
     )
     tables = (
         SchemaName.SENSOR,
+        SchemaName.SAMPLE,
         SchemaName.SAMPLE_DATA,
         SchemaName.MAP,
         SchemaName.ROUTE,
@@ -44,8 +45,8 @@ class REC023(RecordsChecker):
     )
 
     def check(self, context: SanityContext) -> list[Reason] | None:
-        present = present_capabilities(context)
+        present = present_capabilities(context, self.tables)
         return [
             Reason(f"Declared but not present: {item}")
-            for item in sorted(declared_capabilities(context) - present)
+            for item in sorted(capability_items(context.manifest.unwrap().capabilities) - present)
         ] or None