diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6d4d74d8..9a004ca5 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/docs/schema/v1/index.md b/docs/schema/v1/index.md index 781ca93b..c012e60a 100644 --- a/docs/schema/v1/index.md +++ b/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("/") # 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("/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`. diff --git a/docs/schema/v1/requirement.md b/docs/schema/v1/requirement.md index 407918cb..1b2c7a23 100644 --- a/docs/schema/v1/requirement.md +++ b/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`) diff --git a/t4_devkit/common/io.py b/t4_devkit/common/io.py index 9e107d9f..866b4b06 100644 --- a/t4_devkit/common/io.py +++ b/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. diff --git a/t4_devkit/sanity/record/base.py b/t4_devkit/sanity/record/base.py index e712957f..62540fb5 100644 --- a/t4_devkit/sanity/record/base.py +++ b/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)) diff --git a/t4_devkit/sanity/record/rec023.py b/t4_devkit/sanity/record/rec023.py index 3ea679fc..061a91f5 100644 --- a/t4_devkit/sanity/record/rec023.py +++ b/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 diff --git a/t4_devkit/sanity/record/rec024.py b/t4_devkit/sanity/record/rec024.py index 2c4ca181..54791d0f 100644 --- a/t4_devkit/sanity/record/rec024.py +++ b/t4_devkit/sanity/record/rec024.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 REC024(RecordsChecker): ) tables = ( SchemaName.SENSOR, + SchemaName.SAMPLE, SchemaName.SAMPLE_DATA, SchemaName.MAP, SchemaName.ROUTE, @@ -44,8 +45,8 @@ class REC024(RecordsChecker): ) def check(self, context: SanityContext) -> list[Reason] | None: - declared = declared_capabilities(context) + declared = capability_items(context.manifest.unwrap().capabilities) return [ Reason(f"Present but not declared: {item}") - for item in sorted(present_capabilities(context) - declared) + for item in sorted(present_capabilities(context, self.tables) - declared) ] or None diff --git a/t4_devkit/schema/v1/__init__.py b/t4_devkit/schema/v1/__init__.py index 16b9f4e9..fcc37f7c 100644 --- a/t4_devkit/schema/v1/__init__.py +++ b/t4_devkit/schema/v1/__init__.py @@ -24,3 +24,4 @@ from .vehicle import * # noqa from .vehicle_state import * # noqa from .visibility import * # noqa +from .writer import * # noqa diff --git a/t4_devkit/schema/v1/manifest.py b/t4_devkit/schema/v1/manifest.py index 7a9a1667..c1e97f57 100644 --- a/t4_devkit/schema/v1/manifest.py +++ b/t4_devkit/schema/v1/manifest.py @@ -2,11 +2,14 @@ import os.path as osp from enum import Enum, unique -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Mapping, Sequence from attrs import converters, define, field, validators from t4_devkit.common.io import load_json +from t4_devkit.schema.tables.sensor import SensorModality + +from .name import SchemaName if TYPE_CHECKING: from t4_devkit.typing import PathLike @@ -120,6 +123,73 @@ class Capabilities: vehicle_state: bool = field(validator=validators.instance_of(bool)) scenario: bool = field(validator=validators.instance_of(bool)) + @classmethod + def from_tables(cls, tables: Mapping[SchemaName, Sequence[Any]]) -> Capabilities: + """Return what the tables contain: channels with frames, filled tables, tasks with labels. + + `prediction` needs every sample to have a key frame (all frames annotated); `planning` + needs scenarios. + + Args: + tables (Mapping[SchemaName, Sequence[Any]]): 1.0 records per table; a missing table + is empty. + + Returns: + The capabilities. + """ + + def records(table: SchemaName) -> Sequence[Any]: + return tables.get(table, []) + + sensors = records(SchemaName.SENSOR) + with_frames = {frame.sensor_token for frame in records(SchemaName.SAMPLE_DATA)} + channels = { + modality: list( + dict.fromkeys( + sensor.channel + for sensor in sensors + if sensor.modality == modality and sensor.token in with_frames + ) + ) + for modality in SensorModality + } + lidar_columns = { + field.name + for sensor in sensors + if sensor.lidar is not None + for field in sensor.lidar.fields + } + samples = records(SchemaName.SAMPLE) + things = {category.token for category in records(SchemaName.CATEGORY) if category.is_thing} + object_anns = records(SchemaName.OBJECT_ANN) + has_task = { + Task.DETECTION3D: bool(records(SchemaName.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(SchemaName.KEYPOINT)), + # states, or lamp instances linked to the map + Task.TRAFFIC_LIGHT: bool(records(SchemaName.TRAFFIC_LIGHT)) + or any( + instance.map_element.id is not None for instance in records(SchemaName.INSTANCE) + ), + Task.PREDICTION: bool(samples) and all(sample.has_key_frame for sample in samples), + Task.PLANNING: bool(records(SchemaName.SCENARIO)), + } + return cls( + tasks=[task for task in Task if has_task[task]], + modalities=Modalities( + camera=channels[SensorModality.CAMERA], + lidar=channels[SensorModality.LIDAR], + radar=channels[SensorModality.RADAR], + ), + map=bool(records(SchemaName.MAP)), + route=bool(records(SchemaName.ROUTE)), + traffic_light=bool(records(SchemaName.TRAFFIC_LIGHT)), + vehicle_state=bool(records(SchemaName.VEHICLE_STATE)), + scenario=bool(records(SchemaName.SCENARIO)), + ) + @define class AnnotationSpecification: diff --git a/t4_devkit/schema/v1/writer.py b/t4_devkit/schema/v1/writer.py new file mode 100644 index 00000000..7da8152a --- /dev/null +++ b/t4_devkit/schema/v1/writer.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import os +import shutil +from pathlib import Path +from typing import TYPE_CHECKING, Any, Iterable, Mapping, Sequence + +from t4_devkit.common.io import save_pointcloud, save_video +from t4_devkit.common.serialize import serialize_dataclass, serialize_dataclasses + +from .manifest import MANIFEST_FILENAME, Capabilities, Manifest +from .name import SchemaName + +if TYPE_CHECKING: + import numpy as np + + from t4_devkit.typing import PathLike + +__all__ = ["DatasetWriter"] + + +class DatasetWriter: + """Write a T4Dataset 1.0 dataset folder: sensor data, then tables and `manifest.json`. + + The writer writes files; the caller decides their content. Data files go first + (`write_lidar`, `write_camera`, `copy`), then `save` writes the tables and the manifest, with + the hash of every file in the folder. + + Examples: + >>> writer = DatasetWriter("out//") + >>> path, frame_count = writer.write_lidar("LIDAR_TOP", frames, fields) + >>> writer.copy("in/map/lanelet2_map.osm", "map/lanelet2_map.osm") + >>> manifest = writer.save(tables, manifest_info) + """ + + # Folders whose files are listed in `manifest.json` as assets. + ASSET_DIRS: tuple[str, ...] = ("data", "map", "input_bag") + + def __init__(self, output_root: PathLike) -> None: + """ + Args: + output_root (PathLike): Dataset root to create, i.e. `/`. + + Raises: + FileExistsError: If the folder already exists: a dataset is never written in place. + """ + self.output_root = Path(output_root) + if self.output_root.exists(): + raise FileExistsError(f"{self.output_root} already exists") + self.output_root.mkdir(parents=True) + + def write_lidar( + self, + channel: str, + frames: Iterable[Mapping[str, np.ndarray]], + fields: Sequence[tuple[str, str]], + **kwargs: Any, + ) -> tuple[str, int]: + """Write the frames of a LiDAR channel to `data//.parquet`. + + Args: + channel (str): Channel name. + frames (Iterable[Mapping[str, np.ndarray]]): Point fields per frame. + fields (Sequence[tuple[str, str]]): (name, dtype) of each column. + **kwargs: Options of `save_pointcloud`. + + Returns: + The path relative to the dataset root and the number of frames. + + Raises: + FileExistsError: If the channel's file is already written. + """ + path = self._data_path(channel, "parquet") + frame_count = save_pointcloud(self.output_root / path, frames, fields, **kwargs) + return path, frame_count + + def write_camera( + self, + channel: str, + frames: Iterable[np.ndarray], + frame_rate: int, + **kwargs: Any, + ) -> dict[str, Any]: + """Write the frames of a camera channel to `data//.mp4`. + + Args: + channel (str): Channel name. + frames (Iterable[np.ndarray]): RGB images. + frame_rate (int): Frames per second. + **kwargs: Options of `save_video`. + + Returns: + The values of `sensor.camera.video`, `path` included. + + Raises: + FileExistsError: If the channel's file is already written. + """ + path = self._data_path(channel, "mp4") + video = save_video(self.output_root / path, frames, frame_rate, **kwargs) + return {"path": path} | video + + def copy(self, source: PathLike, path: str) -> None: + """Copy a file or folder as it is, e.g. the map, `input_bag/` or per-frame images. + + Args: + source (PathLike): File or folder to copy. + path (str): Destination relative to the dataset root. + """ + destination = self.output_root / path + destination.parent.mkdir(parents=True, exist_ok=True) + if Path(source).is_dir(): + shutil.copytree(source, destination) + else: + shutil.copy2(source, destination) + + def save( + self, + tables: Mapping[SchemaName, Sequence[Any]], + manifest_info: Mapping[str, Any], + ) -> Manifest: + """Write the tables and `manifest.json`. + + Every required table is written, and every optional table with records. The manifest + lists them and every file under `data/`, `map/` and `input_bag/` with its hash and size. + Its `capabilities` are computed from the tables unless given, and t4-devkit is added to + its `generators`. + + Args: + tables (Mapping[SchemaName, Sequence[Any]]): 1.0 records per table. + manifest_info (Mapping[str, Any]): The other `manifest.json` fields: identity, + provenance, migration notes, ... + + Returns: + The manifest written. + """ + table_entries = {} + for table in SchemaName: + records = tables.get(table, []) + if not records and table.is_optional(): + continue + path = f"annotation/{table.filename}" + self._write_json(path, serialize_dataclasses(list(records))) + table_entries[table.value] = {"path": path, "sha256": self._sha256(path)} + + assets = [ + { + "path": path, + "sha256": self._sha256(path), + "bytes": (self.output_root / path).stat().st_size, + } + for path in self._asset_paths() + ] + generator = {"name": "t4-devkit", "version": importlib.metadata.version("t4-devkit")} + generators = list(manifest_info.get("generators", [])) + if generator not in generators: + generators.append(generator) + capabilities = manifest_info.get("capabilities") or Capabilities.from_tables(tables) + + data = dict(manifest_info) + data.update( + generators=generators, + capabilities=capabilities, + tables=table_entries, + assets=assets, + ) + manifest = Manifest(**data) + self._write_json(MANIFEST_FILENAME, serialize_dataclass(manifest)) + return manifest + + def _data_path(self, channel: str, extension: str) -> str: + path = f"data/{channel}/{channel}.{extension}" + if (self.output_root / path).exists(): + raise FileExistsError(f"{path} is already written: a channel has one file") + (self.output_root / path).parent.mkdir(parents=True, exist_ok=True) + return path + + def _asset_paths(self) -> list[str]: + """Every file under the asset folders, relative to the dataset root, sorted.""" + return sorted( + filepath.relative_to(self.output_root).as_posix() + for directory in self.ASSET_DIRS + for filepath in (self.output_root / directory).rglob("*") + if filepath.is_file() + ) + + def _write_json(self, path: str, data: Any) -> None: + filepath = self.output_root / path + filepath.parent.mkdir(parents=True, exist_ok=True) + filepath.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n") + + def _sha256(self, path: str) -> str: + digest = hashlib.sha256() + with open(os.fspath(self.output_root / path), "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/tests/common/test_io.py b/tests/common/test_io.py index 75bf2af2..e7bc77bd 100644 --- a/tests/common/test_io.py +++ b/tests/common/test_io.py @@ -8,7 +8,14 @@ import pytest from PIL import Image -from t4_devkit.common.io import load_image, load_json, load_pointcloud, save_json +from t4_devkit.common.io import ( + load_image, + load_json, + load_pointcloud, + save_json, + save_pointcloud, + save_video, +) SAMPLE_PARQUET = Path(__file__).parents[1] / "sample/t4dataset_v1/data/LIDAR_TOP/LIDAR_TOP.parquet" SAMPLE_VIDEO = Path(__file__).parents[1] / "sample/t4dataset_v1/data/CAM_FRONT/CAM_FRONT.mp4" @@ -96,3 +103,82 @@ def test_load_image_file(tmp_path: Path) -> None: Image.fromarray(pixels).save(path) assert np.array_equal(load_image(path), pixels) + + +def test_save_pointcloud(tmp_path: Path) -> None: + """Frames come back bit-exact in their dtypes, one row group each, empty frames included.""" + fields = [ + ("x", "float32"), + ("intensity", "uint8"), + ("channel", "uint16"), + ("time_stamp", "uint32"), + ] + rng = np.random.default_rng(0) + frames = [ + { + "x": rng.normal(size=n).astype(np.float32), + "intensity": rng.integers(0, 256, n).astype(np.uint8), + "channel": rng.integers(0, 128, n).astype(np.uint16), + "time_stamp": rng.integers(0, 100_000_000, n).astype(np.uint32), + } + for n in (5, 0, 3) + ] + path = tmp_path / "LIDAR_TOP.parquet" + + assert save_pointcloud(path, frames, fields) == 3 + for frame_index, expected in enumerate(frames): + frame = load_pointcloud(path, frame_index=frame_index) + assert list(frame) == [name for name, _ in fields] + for name, dtype in fields: + assert frame[name].dtype == np.dtype(dtype) + assert np.array_equal(frame[name], expected[name]) + + columns = pq.ParquetFile(path).metadata.row_group(0).to_dict()["columns"] + encodings = {column["path_in_schema"]: column["encodings"] for column in columns} + assert "BYTE_STREAM_SPLIT" in encodings["x"] + assert "RLE_DICTIONARY" in encodings["intensity"] + assert "DELTA_BINARY_PACKED" in encodings["channel"] + assert "DELTA_BINARY_PACKED" in encodings["time_stamp"] + + +@pytest.mark.parametrize( + ("dtype", "bit_depth", "max_value"), [(np.uint8, 8, 255), (np.uint16, 10, 65535)] +) +def test_save_video(tmp_path: Path, dtype, bit_depth: int, max_value: int) -> None: + """Frames read back by index, in any order, close to what was written.""" + images = [] + for i in range(3): + image = np.zeros((48, 64, 3), dtype=dtype) + image[..., 0] = np.linspace(0, max_value, 64).astype(dtype) + image[..., 1] = max_value // 4 * i + images.append(image) + path = tmp_path / "CAM_FRONT.mp4" + + video = save_video(path, iter(images), frame_rate=10) + assert video["frame_count"] == 3 and video["timescale"] == 10 + assert video["bit_depth"] == bit_depth + for frame_index in (2, 0, 1): + image = load_image(path, frame_index=frame_index) + assert image.dtype == dtype and image.shape == (48, 64, 3) + # lossy, but the green level of each frame is kept + green = image[..., 1].mean() / max_value + assert abs(green - 0.25 * frame_index) < 0.03 + + +def test_save_pointcloud_null_column(tmp_path: Path) -> None: + """A field that is None in a frame is written null and left out when read.""" + fields = [("x", "float32"), ("label", "uint8")] + frames = [ + {"x": np.array([1.0, 2.0], dtype=np.float32), "label": np.array([0, 3], dtype=np.uint8)}, + {"x": np.array([4.0], dtype=np.float32), "label": None}, + ] + path = tmp_path / "LIDAR_TOP.parquet" + + assert save_pointcloud(path, frames, fields) == 2 + assert load_pointcloud(path, frame_index=0)["label"].tolist() == [0, 3] + assert list(load_pointcloud(path, frame_index=1)) == ["x"] + + # what the reader returns can be written back as it is + again = tmp_path / "again.parquet" + save_pointcloud(again, [load_pointcloud(path, frame_index=i) for i in range(2)], fields) + assert list(load_pointcloud(again, frame_index=1)) == ["x"] diff --git a/tests/sample/t4dataset_v1/annotation/caption.json b/tests/sample/t4dataset_v1/annotation/caption.json index 5d091446..12a21b1a 100644 --- a/tests/sample/t4dataset_v1/annotation/caption.json +++ b/tests/sample/t4dataset_v1/annotation/caption.json @@ -6,7 +6,9 @@ "end_timestamp": 1704067200200000, "sample_data_token": null, "scenario_token": "scenario0", - "instance_tokens": ["instance_car"], + "instance_tokens": [ + "instance_car" + ], "caption_type": "description", "text": "The ego vehicle follows a car.", "language": "en", diff --git a/tests/sample/t4dataset_v1/annotation/ego_pose.json b/tests/sample/t4dataset_v1/annotation/ego_pose.json index b1f807ec..0d846ce6 100644 --- a/tests/sample/t4dataset_v1/annotation/ego_pose.json +++ b/tests/sample/t4dataset_v1/annotation/ego_pose.json @@ -1,8 +1,17 @@ [ { "token": "ego_cam0", - "translation": [0.0, 0.0, 0.0], - "rotation": [1.0, 0.0, 0.0, 0.0], + "translation": [ + 0.0, + 0.0, + 0.0 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ], "timestamp": 1704067200000000, "twist": null, "acceleration": null, @@ -13,8 +22,17 @@ }, { "token": "ego_cam1", - "translation": [1.0, 0.0, 0.0], - "rotation": [1.0, 0.0, 0.0, 0.0], + "translation": [ + 1.0, + 0.0, + 0.0 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ], "timestamp": 1704067200100000, "twist": null, "acceleration": null, @@ -25,8 +43,17 @@ }, { "token": "ego_lidar0", - "translation": [0.0, 0.0, 0.0], - "rotation": [1.0, 0.0, 0.0, 0.0], + "translation": [ + 0.0, + 0.0, + 0.0 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ], "timestamp": 1704067200000000, "twist": null, "acceleration": null, @@ -37,8 +64,17 @@ }, { "token": "ego_lidar1", - "translation": [1.0, 0.0, 0.0], - "rotation": [1.0, 0.0, 0.0, 0.0], + "translation": [ + 1.0, + 0.0, + 0.0 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ], "timestamp": 1704067200100000, "twist": null, "acceleration": null, @@ -49,8 +85,17 @@ }, { "token": "ego_lidar2", - "translation": [2.0, 0.0, 0.0], - "rotation": [1.0, 0.0, 0.0, 0.0], + "translation": [ + 2.0, + 0.0, + 0.0 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ], "timestamp": 1704067200200000, "twist": null, "acceleration": null, diff --git a/tests/sample/t4dataset_v1/annotation/keypoint.json b/tests/sample/t4dataset_v1/annotation/keypoint.json index a31da9c6..bd3f13c4 100644 --- a/tests/sample/t4dataset_v1/annotation/keypoint.json +++ b/tests/sample/t4dataset_v1/annotation/keypoint.json @@ -3,8 +3,15 @@ "token": "keypoint0", "sample_data_token": "cam0", "instance_token": "instance_car", - "category_tokens": ["category_car"], - "keypoints": [[12.0, 18.0]], + "category_tokens": [ + "category_car" + ], + "keypoints": [ + [ + 12.0, + 18.0 + ] + ], "num_keypoints": 1, "annotation_type": "manual", "provenance": null, diff --git a/tests/sample/t4dataset_v1/annotation/map.json b/tests/sample/t4dataset_v1/annotation/map.json index 5cdd7134..28828b93 100644 --- a/tests/sample/t4dataset_v1/annotation/map.json +++ b/tests/sample/t4dataset_v1/annotation/map.json @@ -2,7 +2,9 @@ { "token": "map", "name": null, - "log_tokens": ["log"], + "log_tokens": [ + "log" + ], "format": "lanelet2_osm", "filename": "map/lanelet2_map.osm", "sha256": "4b325cf23adf604023824cc4a3243a7aac56c690fa872b68ce58e42a4529f054", diff --git a/tests/sample/t4dataset_v1/annotation/object_ann.json b/tests/sample/t4dataset_v1/annotation/object_ann.json index 000a2bc7..d8bafa4b 100644 --- a/tests/sample/t4dataset_v1/annotation/object_ann.json +++ b/tests/sample/t4dataset_v1/annotation/object_ann.json @@ -5,7 +5,12 @@ "instance_token": "instance_car", "category_token": "category_car", "attribute_tokens": [], - "bbox": [10.0, 10.0, 30.0, 20.0], + "bbox": [ + 10.0, + 10.0, + 30.0, + 20.0 + ], "mask": null, "orientation": null, "number": null, @@ -19,9 +24,17 @@ "instance_token": null, "category_token": "category_road", "attribute_tokens": [], - "bbox": [0.0, 24.0, 64.0, 48.0], + "bbox": [ + 0.0, + 24.0, + 64.0, + 48.0 + ], "mask": { - "size": [48, 64], + "size": [ + 48, + 64 + ], "counts": "aDBoMGgwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDA=" }, "orientation": null, diff --git a/tests/sample/t4dataset_v1/annotation/route.json b/tests/sample/t4dataset_v1/annotation/route.json index 924c9de0..a1b3123c 100644 --- a/tests/sample/t4dataset_v1/annotation/route.json +++ b/tests/sample/t4dataset_v1/annotation/route.json @@ -7,17 +7,38 @@ "valid_until": null, "route_uuid": null, "start_pose": { - "translation": [0.0, 0.0, 0.0], - "rotation": [1.0, 0.0, 0.0, 0.0] + "translation": [ + 0.0, + 0.0, + 0.0 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ] }, "goal_pose": { - "translation": [100.0, 0.0, 0.0], - "rotation": [1.0, 0.0, 0.0, 0.0] + "translation": [ + 100.0, + 0.0, + 0.0 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ] }, "segments": [ { "preferred_primitive_id": 1000, - "primitive_ids": [1000, 1001] + "primitive_ids": [ + 1000, + 1001 + ] } ], "allow_modification": false, diff --git a/tests/sample/t4dataset_v1/annotation/sample_annotation.json b/tests/sample/t4dataset_v1/annotation/sample_annotation.json index b31a929a..761b4ca5 100644 --- a/tests/sample/t4dataset_v1/annotation/sample_annotation.json +++ b/tests/sample/t4dataset_v1/annotation/sample_annotation.json @@ -3,12 +3,31 @@ "token": "box0", "sample_token": "sample0", "instance_token": "instance_car", - "attribute_tokens": ["attribute_moving"], + "attribute_tokens": [ + "attribute_moving" + ], "visibility_token": "visibility_full", - "translation": [10.0, 0.0, 1.0], - "size": [1.9, 4.5, 1.6], - "rotation": [1.0, 0.0, 0.0, 0.0], - "velocity": [1.0, 0.0, 0.0], + "translation": [ + 10.0, + 0.0, + 1.0 + ], + "size": [ + 1.9, + 4.5, + 1.6 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ], + "velocity": [ + 1.0, + 0.0, + 0.0 + ], "acceleration": null, "num_lidar_pts": 0, "num_radar_pts": null, @@ -22,12 +41,31 @@ "token": "box1", "sample_token": "sample1", "instance_token": "instance_car", - "attribute_tokens": ["attribute_moving"], + "attribute_tokens": [ + "attribute_moving" + ], "visibility_token": "visibility_full", - "translation": [11.0, 0.0, 1.0], - "size": [1.9, 4.5, 1.6], - "rotation": [1.0, 0.0, 0.0, 0.0], - "velocity": [1.0, 0.0, 0.0], + "translation": [ + 11.0, + 0.0, + 1.0 + ], + "size": [ + 1.9, + 4.5, + 1.6 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ], + "velocity": [ + 1.0, + 0.0, + 0.0 + ], "acceleration": null, "num_lidar_pts": 0, "num_radar_pts": null, diff --git a/tests/sample/t4dataset_v1/annotation/scenario.json b/tests/sample/t4dataset_v1/annotation/scenario.json index c3c419fc..807ba4dd 100644 --- a/tests/sample/t4dataset_v1/annotation/scenario.json +++ b/tests/sample/t4dataset_v1/annotation/scenario.json @@ -8,7 +8,9 @@ "status": null, "lateral_decision": null, "longitudinal_decision": null, - "events": ["following"], + "events": [ + "following" + ], "scenery": null, "dynamic_entities": null, "justification": null, diff --git a/tests/sample/t4dataset_v1/annotation/sensor.json b/tests/sample/t4dataset_v1/annotation/sensor.json index f5bfcdc8..c8524b89 100644 --- a/tests/sample/t4dataset_v1/annotation/sensor.json +++ b/tests/sample/t4dataset_v1/annotation/sensor.json @@ -4,17 +4,44 @@ "channel": "CAM_FRONT", "index": 0, "modality": "camera", - "translation": [1.0, 0.0, 1.5], - "rotation": [0.5, -0.5, 0.5, -0.5], + "translation": [ + 1.0, + 0.0, + 1.5 + ], + "rotation": [ + 0.5, + -0.5, + 0.5, + -0.5 + ], "camera": { "width": 64, "height": 48, "intrinsic": [ - [50.0, 0.0, 32.0], - [0.0, 50.0, 24.0], - [0.0, 0.0, 1.0] + [ + 50.0, + 0.0, + 32.0 + ], + [ + 0.0, + 50.0, + 24.0 + ], + [ + 0.0, + 0.0, + 1.0 + ] + ], + "distortion": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 ], - "distortion": [0.0, 0.0, 0.0, 0.0, 0.0], "distortion_model": "plumb_bob", "is_rectified": false, "video": { @@ -41,8 +68,17 @@ "channel": "LIDAR_TOP", "index": 1, "modality": "lidar", - "translation": [0.0, 0.0, 2.0], - "rotation": [1.0, 0.0, 0.0, 0.0], + "translation": [ + 0.0, + 0.0, + 2.0 + ], + "rotation": [ + 1.0, + 0.0, + 0.0, + 0.0 + ], "camera": null, "lidar": { "container": { diff --git a/tests/sample/t4dataset_v1/data/CAM_FRONT/CAM_FRONT.mp4 b/tests/sample/t4dataset_v1/data/CAM_FRONT/CAM_FRONT.mp4 index c14d7f4a..3e49e0c8 100644 Binary files a/tests/sample/t4dataset_v1/data/CAM_FRONT/CAM_FRONT.mp4 and b/tests/sample/t4dataset_v1/data/CAM_FRONT/CAM_FRONT.mp4 differ diff --git a/tests/sample/t4dataset_v1/data/LIDAR_TOP/LIDAR_TOP.parquet b/tests/sample/t4dataset_v1/data/LIDAR_TOP/LIDAR_TOP.parquet index 9193998b..37928b6f 100644 Binary files a/tests/sample/t4dataset_v1/data/LIDAR_TOP/LIDAR_TOP.parquet and b/tests/sample/t4dataset_v1/data/LIDAR_TOP/LIDAR_TOP.parquet differ diff --git a/tests/sample/t4dataset_v1/make_sample_v1.py b/tests/sample/t4dataset_v1/make_sample_v1.py index b385d703..ebb97aff 100644 --- a/tests/sample/t4dataset_v1/make_sample_v1.py +++ b/tests/sample/t4dataset_v1/make_sample_v1.py @@ -4,29 +4,25 @@ (LIDAR_TOP, Parquet) frame each, then a non-key frame with a LiDAR frame only (a channel can be missing). One record in every optional table. All data is synthetic and deterministic. -Requires an ffmpeg build with libsvtav1, and npx to format JSON with the same prettier version -as pre-commit (so committing does not change the hashes in manifest.json): +Written with `DatasetWriter`, like a real dataset (no FFmpeg needed): - FFMPEG=/path/to/ffmpeg python tests/sample/t4dataset_v1/make_sample_v1.py + python tests/sample/t4dataset_v1/make_sample_v1.py """ from __future__ import annotations import base64 import hashlib -import json -import os -import subprocess +import shutil +import tempfile from pathlib import Path import numpy as np -import pyarrow as pa -import pyarrow.parquet as pq from pycocotools import mask as cocomask +from t4_devkit.schema.v1 import SCHEMAS, DatasetWriter, SchemaName + ROOT = Path(__file__).parent -FFMPEG = os.environ.get("FFMPEG", "ffmpeg") -PRETTIER = ["npx", "--yes", "prettier@4.0.0-alpha.8", "--write", "--log-level", "warn"] NUM_FRAMES = 2 # key frames, with a camera and a LiDAR frame each NUM_SAMPLES = NUM_FRAMES + 1 # the last sample is a non-key frame with a LiDAR frame only @@ -36,21 +32,16 @@ T0 = 1704067200000000 # 2024-01-01T00:00:00Z in microseconds DT = 1_000_000 // FPS -CAMERA_PATH = "data/CAM_FRONT/CAM_FRONT.mp4" -LIDAR_PATH = "data/LIDAR_TOP/LIDAR_TOP.parquet" +LIDAR_FIELDS = [ + ("x", "float32"), + ("y", "float32"), + ("z", "float32"), + ("intensity", "float32"), + ("ring", "uint16"), +] MAP_PATH = "map/lanelet2_map.osm" -def write_json(path: Path, data: object) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(data, indent=2) + "\n") - subprocess.run([*PRETTIER, str(path)], check=True) - - -def sha256(path: Path) -> str: - return hashlib.sha256(path.read_bytes()).hexdigest() - - def road_mask() -> dict[str, object]: """RLE mask of the road: the bottom half of the image (rows 24-47), as pycocotools writes it.""" mask = np.zeros((HEIGHT, WIDTH), dtype=np.uint8, order="F") @@ -59,53 +50,27 @@ def road_mask() -> dict[str, object]: return {"size": [HEIGHT, WIDTH], "counts": base64.b64encode(rle["counts"]).decode()} -def write_camera() -> None: - """Encode NUM_FRAMES synthetic RGB frames as intra-only AV1 with pts == frame_index.""" - frames = np.zeros((NUM_FRAMES, HEIGHT, WIDTH, 3), dtype=np.uint8) +def camera_frames() -> list[np.ndarray]: + """NUM_FRAMES synthetic RGB frames.""" + frames = [] for i in range(NUM_FRAMES): - frames[i, :, :, 0] = np.linspace(0, 255, WIDTH, dtype=np.uint8) - frames[i, :, :, 1] = 128 * i - path = ROOT / CAMERA_PATH - path.parent.mkdir(parents=True, exist_ok=True) - # fmt: off - cmd = [ - FFMPEG, "-y", "-loglevel", "error", - "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{WIDTH}x{HEIGHT}", "-framerate", str(FPS), - "-i", "-", - "-c:v", "libsvtav1", "-preset", "6", "-crf", "26", "-g", "1", - "-svtav1-params", "pred-struct=0:scd=0:film-grain=0:fast-decode=2:lp=0:log-level=1", - "-pix_fmt", "yuv420p", "-color_range", "pc", "-colorspace", "bt470bg", - "-color_primaries", "bt709", "-color_trc", "iec61966-2-1", - "-video_track_timescale", str(FPS), "-movflags", "+faststart", - "-map_metadata", "-1", "-fflags", "+bitexact", - str(path), - ] - # fmt: on - subprocess.run(cmd, input=frames.tobytes(), check=True) - - -def write_lidar() -> None: - """Write NUM_SAMPLES frames of NUM_POINTS points, one row group per frame.""" + frame = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8) + frame[:, :, 0] = np.linspace(0, 255, WIDTH, dtype=np.uint8) + frame[:, :, 1] = 128 * i + frames.append(frame) + return frames + + +def lidar_frames() -> list[dict[str, np.ndarray]]: + """NUM_SAMPLES frames of NUM_POINTS points.""" rng = np.random.default_rng(0) - path = ROOT / LIDAR_PATH - path.parent.mkdir(parents=True, exist_ok=True) - schema = pa.schema( - [ - ("x", pa.float32()), - ("y", pa.float32()), - ("z", pa.float32()), - ("intensity", pa.float32()), - ("ring", pa.uint16()), - ] - ) - with pq.ParquetWriter(path, schema, compression="zstd") as writer: - for _ in range(NUM_SAMPLES): - xyzi = rng.uniform(-10, 10, (NUM_POINTS, 4)).astype(np.float32) - ring = np.arange(NUM_POINTS, dtype=np.uint16) - table = pa.table( - [*(pa.array(xyzi[:, k]) for k in range(4)), pa.array(ring)], schema=schema - ) - writer.write_table(table, row_group_size=NUM_POINTS) + frames = [] + for _ in range(NUM_SAMPLES): + xyzi = rng.uniform(-10, 10, (NUM_POINTS, 4)).astype(np.float32) + frame = {name: xyzi[:, k] for k, (name, _) in enumerate(LIDAR_FIELDS[:4])} + frame["ring"] = np.arange(NUM_POINTS, dtype=np.uint16) + frames.append(frame) + return frames # Added to the 0.x sample's map: a traffic light (way) and its regulatory element, controlling @@ -131,9 +96,8 @@ def write_lidar() -> None: """ -def write_map() -> None: - path = ROOT / MAP_PATH - path.parent.mkdir(parents=True, exist_ok=True) +def write_map(path: Path) -> None: + """Write the 0.x sample's map with the traffic light added.""" osm = (ROOT.parent / "t4dataset" / MAP_PATH).read_text() path.write_text(osm.replace("\n\n", MAP_ADDITIONS)) @@ -146,7 +110,7 @@ def label( return {"annotation_type": annotation_type, "provenance": provenance, "confidence": confidence} -def tables() -> dict[str, list[dict]]: +def tables(camera_video: dict, lidar_path: str, map_sha256: str) -> dict[str, list[dict]]: identity = [1.0, 0.0, 0.0, 0.0] zero = [0.0, 0.0, 0.0] samples = [f"sample{i}" for i in range(NUM_SAMPLES)] @@ -161,8 +125,8 @@ def link(tokens: list[str], i: int) -> dict: sample_data = [] ego_pose = [] for channel, sensor, path, fileformat, num_points, num_frames in ( - ("cam", "sensor_cam", CAMERA_PATH, "av1_mp4", None, NUM_FRAMES), - ("lidar", "sensor_lidar", LIDAR_PATH, "parquet", NUM_POINTS, NUM_SAMPLES), + ("cam", "sensor_cam", camera_video["path"], "av1_mp4", None, NUM_FRAMES), + ("lidar", "sensor_lidar", lidar_path, "parquet", NUM_POINTS, NUM_SAMPLES), ): tokens = [f"{channel}{i}" for i in range(num_frames)] for i, token in enumerate(tokens): @@ -226,15 +190,7 @@ def link(tokens: list[str], i: int) -> dict: "distortion": [0.0, 0.0, 0.0, 0.0, 0.0], "distortion_model": "plumb_bob", "is_rectified": False, - "video": { - "path": CAMERA_PATH, - "format": "av1_mp4", - "codec_profile": "svt-av1 preset6 crf26 intra", - "pixel_format": "yuv420p full-range bt601", - "bit_depth": 8, - "frame_count": NUM_FRAMES, - "timescale": FPS, - }, + "video": camera_video, }, "lidar": None, "radar": None, @@ -256,17 +212,11 @@ def link(tokens: list[str], i: int) -> dict: "lidar": { "container": { "format": "parquet", - "path": LIDAR_PATH, + "path": lidar_path, "frame_count": NUM_SAMPLES, "compression": "zstd", }, - "fields": [ - {"name": "x", "dtype": "float32"}, - {"name": "y", "dtype": "float32"}, - {"name": "z", "dtype": "float32"}, - {"name": "intensity", "dtype": "float32"}, - {"name": "ring", "dtype": "uint16"}, - ], + "fields": [{"name": name, "dtype": dtype} for name, dtype in LIDAR_FIELDS], "source_channels": [], "frame_id": "base_link", "is_undistorted": True, @@ -321,7 +271,7 @@ def link(tokens: list[str], i: int) -> dict: "log_tokens": ["log"], "format": "lanelet2_osm", "filename": MAP_PATH, - "sha256": sha256(ROOT / MAP_PATH), + "sha256": map_sha256, "map_id": "sample_map", "version_id": "sample_map_version", "projector": { @@ -512,62 +462,61 @@ def link(tokens: list[str], i: int) -> dict: } -def main() -> None: - write_camera() - write_lidar() - write_map() - - table_entries = {} - for name, records in tables().items(): - path = ROOT / "annotation" / f"{name}.json" - write_json(path, records) - table_entries[name] = {"path": f"annotation/{name}.json", "sha256": sha256(path)} - - assets = [ - {"path": p, "sha256": sha256(ROOT / p), "bytes": (ROOT / p).stat().st_size} - for p in (CAMERA_PATH, LIDAR_PATH, MAP_PATH) - ] - - manifest = { - "schema_version": "1.0.0", - "dataset_id": "sample", - "version": 0, - "created_at": "2024-01-01T00:00:00Z", - "generators": [{"name": "t4-devkit", "version": "sample"}], - "capabilities": { - "tasks": ["detection3d", "detection2d", "segmentation2d", "keypoint", "traffic_light"], - "modalities": {"camera": ["CAM_FRONT"], "lidar": ["LIDAR_TOP"], "radar": []}, - "map": True, - "route": True, - "traffic_light": True, - "vehicle_state": True, - "scenario": True, +# Manifest fields the writer does not compute (it adds the files, capabilities and t4-devkit). +MANIFEST_INFO = { + "schema_version": "1.0.0", + "dataset_id": "sample", + "version": 0, + "created_at": "2024-01-01T00:00:00Z", + "generators": [], + "annotation_specifications": {}, + "provenance": { + "rosbag": { + "log_file_ids": [], + "storage": None, + "record_start": "2024-01-01T00:00:00Z", + "record_end": "2024-01-01T00:00:01Z", }, - "annotation_specifications": {}, - "tables": table_entries, - "assets": assets, - "provenance": { - "rosbag": { - "log_file_ids": [], - "storage": None, - "record_start": "2024-01-01T00:00:00Z", - "record_end": "2024-01-01T00:00:01Z", - }, - "topics": { - "CAM_FRONT": "/sensing/camera/camera0/image_rect_color/compressed", - "LIDAR_TOP": "/sensing/lidar/concatenated/pointcloud", - }, - "process_config_id": None, - "process_config_version_id": None, + "topics": { + "CAM_FRONT": "/sensing/camera/camera0/image_rect_color/compressed", + "LIDAR_TOP": "/sensing/lidar/concatenated/pointcloud", }, - "project_id": None, - "source_project_id": None, - "is_synthetic": True, - "is_discontinuous": None, - "migrated_from_schema_version": None, - "migration_notes": None, - } - write_json(ROOT / "manifest.json", manifest) + "process_config_id": None, + "process_config_version_id": None, + }, + "project_id": None, + "source_project_id": None, + "is_synthetic": True, + "is_discontinuous": None, + "migrated_from_schema_version": None, + "migration_notes": None, +} + + +def main() -> None: + with tempfile.TemporaryDirectory() as tmp: + output_root = Path(tmp) / "sample" + writer = DatasetWriter(output_root) + camera_video = writer.write_camera("CAM_FRONT", camera_frames(), FPS) + lidar_path, _ = writer.write_lidar("LIDAR_TOP", lidar_frames(), LIDAR_FIELDS) + map_file = Path(tmp) / "lanelet2_map.osm" + write_map(map_file) + writer.copy(map_file, MAP_PATH) + + map_sha256 = hashlib.sha256(map_file.read_bytes()).hexdigest() + records = { + SchemaName(name): [SCHEMAS.get(SchemaName(name)).from_dict(data) for data in rows] + for name, rows in tables(camera_video, lidar_path, map_sha256).items() + } + writer.save(records, MANIFEST_INFO) + + for name in ("annotation", "data", "map", "manifest.json"): + target = ROOT / name + if target.is_dir(): + shutil.rmtree(target) + elif target.exists(): + target.unlink() + shutil.move(output_root / name, target) if __name__ == "__main__": diff --git a/tests/sample/t4dataset_v1/manifest.json b/tests/sample/t4dataset_v1/manifest.json index 109a399c..95820f7d 100644 --- a/tests/sample/t4dataset_v1/manifest.json +++ b/tests/sample/t4dataset_v1/manifest.json @@ -6,7 +6,7 @@ "generators": [ { "name": "t4-devkit", - "version": "sample" + "version": "0.8.0" } ], "capabilities": { @@ -15,11 +15,16 @@ "detection2d", "segmentation2d", "keypoint", - "traffic_light" + "traffic_light", + "planning" ], "modalities": { - "camera": ["CAM_FRONT"], - "lidar": ["LIDAR_TOP"], + "camera": [ + "CAM_FRONT" + ], + "lidar": [ + "LIDAR_TOP" + ], "radar": [] }, "map": true, @@ -40,11 +45,11 @@ }, "ego_pose": { "path": "annotation/ego_pose.json", - "sha256": "5743a2253f46acfec636b04375939ab9ec494ddc9643cf529e8ff1c5658939cc" + "sha256": "41c24c1c135773c2037e63f8aafe9080e2345eedd8a6ec71872d5ef7cabf6ce2" }, "sensor": { "path": "annotation/sensor.json", - "sha256": "195694efc286e2d61d9e062d64e59a525becfacc54c5aac72bf877060e0481f6" + "sha256": "dd66fc0cadb8cb1e54831067a7f13c9f159de8d686f5f9ba85976706e3ecc80e" }, "scene": { "path": "annotation/scene.json", @@ -60,7 +65,7 @@ }, "map": { "path": "annotation/map.json", - "sha256": "d869de9d95d4ed8560ab8463c9e9aac442e2d07f03f8ce4d798cdd51828bd3b8" + "sha256": "5cd6ff58a1c1d942de187f2902239091e2f662d21d1d564975b300a7c9d4fbab" }, "vehicle_state": { "path": "annotation/vehicle_state.json", @@ -84,15 +89,15 @@ }, "sample_annotation": { "path": "annotation/sample_annotation.json", - "sha256": "30a0ed3f5ba5ced1cfd4e994859da8502ff44b7e6490059a09141c3a3a86f88c" + "sha256": "85afcc5f772d398bafc56b77df9c726d8c0a0e3e1e4421c666bf527cb4adedf8" }, "object_ann": { "path": "annotation/object_ann.json", - "sha256": "d3a84be002e1b8db8ed98d4e547e8baf6a43b1aa101ba6fcd3773a70ddfc7c9e" + "sha256": "5434655b3ce0246de3418e79760c97b428ed5ad2c89964abcc67e9a52ec2873b" }, "keypoint": { "path": "annotation/keypoint.json", - "sha256": "00b4271d8cf39b1b16fa86673a7b7246fbeaeca189ffe64fe6625a74c7adfdaa" + "sha256": "1aff568f5509d4b44a297c80b56aea77c6d6264d3e7992b8dff741a1b59ecac0" }, "traffic_light": { "path": "annotation/traffic_light.json", @@ -100,27 +105,27 @@ }, "route": { "path": "annotation/route.json", - "sha256": "6badf9c9fc00722b62763bc483f24e04a53df9cc5a717130008581b41693cecb" + "sha256": "80b6d0d982fc7d7a972a76e14f348421d7586ba84590f007836857eaa071abc7" }, "scenario": { "path": "annotation/scenario.json", - "sha256": "7d751a7b093765efda5bdde70dbfc61d20bc58de00e7f17a7e6dc68e92a4d068" + "sha256": "8c5b20607c45a4e2cf92ff3b291a036550ed1c0b9d1534f4fa2f4c8dc978c8ad" }, "caption": { "path": "annotation/caption.json", - "sha256": "6871b347a5f918a007ede1b3719512764c7ecda05eb8238959dcf5bd52279f03" + "sha256": "377049baba19c58c7dbb376d67dc6c122fc7aef262be7e81eb04e7c4b1e05e00" } }, "assets": [ { "path": "data/CAM_FRONT/CAM_FRONT.mp4", - "sha256": "a09be0d80dad958cd0a4eac632c3eb7af0fb57de48659b56a58acd6eefbb5682", - "bytes": 894 + "sha256": "2c98ef16ed7269de658b46d5ece3698a9ba3f76449c6487e1c737c1e7fd00a93", + "bytes": 939 }, { "path": "data/LIDAR_TOP/LIDAR_TOP.parquet", - "sha256": "837f92226dabd8dbdcc27141dbae5e69ae0e487a55a2d1117fc2997e24e3fa8b", - "bytes": 3701 + "sha256": "ea286e166f6a4e5266e2765ff60e84f27eef1d32517621ea7bf4f56ee2a2b2f8", + "bytes": 3045 }, { "path": "map/lanelet2_map.osm", diff --git a/tests/schema/v1/test_manifest.py b/tests/schema/v1/test_manifest.py index da986627..19a68588 100644 --- a/tests/schema/v1/test_manifest.py +++ b/tests/schema/v1/test_manifest.py @@ -4,9 +4,17 @@ from pathlib import Path import pytest +from attrs import evolve from t4_devkit.common.serialize import serialize_dataclass -from t4_devkit.schema.v1 import Manifest, load_manifest +from t4_devkit.schema.v1 import ( + Capabilities, + Manifest, + SchemaName, + Task, + load_dataset, + load_manifest, +) SAMPLE_ROOT = Path(__file__).parents[2] / "sample" @@ -49,3 +57,17 @@ def test_invalid_sha256() -> None: data["assets"][0]["sha256"] = "not-a-hash" with pytest.raises(ValueError): Manifest(**data) + + +def test_capabilities_prediction_and_planning() -> None: + """Prediction needs every sample to have a key frame; planning needs scenarios.""" + tables = dict(load_dataset(SAMPLE_ROOT / "t4dataset_v1").tables) + tasks = Capabilities.from_tables(tables).tasks + assert Task.PLANNING in tasks + assert Task.PREDICTION not in tasks # sample2 has no key frame + + tables[SchemaName.SAMPLE] = [evolve(s, has_key_frame=True) for s in tables[SchemaName.SAMPLE]] + tables[SchemaName.SCENARIO] = [] + tasks = Capabilities.from_tables(tables).tasks + assert Task.PREDICTION in tasks + assert Task.PLANNING not in tasks diff --git a/tests/schema/v1/test_writer.py b/tests/schema/v1/test_writer.py new file mode 100644 index 00000000..142a1efb --- /dev/null +++ b/tests/schema/v1/test_writer.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from t4_devkit.common.io import load_image, load_pointcloud +from t4_devkit.common.serialize import serialize_dataclass, serialize_dataclasses +from t4_devkit.sanity import sanity_check +from t4_devkit.schema.v1 import DatasetWriter, SchemaName, load_dataset + +SAMPLE_V1 = Path(__file__).parents[2] / "sample" / "t4dataset_v1" + + +def _manifest_info(dataset) -> dict: + """The manifest fields the caller gives: everything but the files and computed parts.""" + info = serialize_dataclass(dataset.manifest) + for name in ("tables", "assets", "capabilities"): + del info[name] + return info + + +def _sensor(dataset, channel: str): + return next(sensor for sensor in dataset.tables[SchemaName.SENSOR] if sensor.channel == channel) + + +def test_round_trip(tmp_path: Path) -> None: + """A dataset written from loaded records loads back to the same records, and passes sanity.""" + dataset = load_dataset(SAMPLE_V1) + data_root = tmp_path / "sample" # named as the sample's `dataset_id` + writer = DatasetWriter(data_root) + for directory in ("data", "map"): + writer.copy(SAMPLE_V1 / directory, directory) + manifest = writer.save(dataset.tables, _manifest_info(dataset)) + + assert manifest.capabilities == dataset.manifest.capabilities + assert any(generator.name == "t4-devkit" for generator in manifest.generators) + loaded = load_dataset(data_root) + for table in SchemaName: + assert serialize_dataclasses(loaded.tables[table]) == serialize_dataclasses( + dataset.tables[table] + ), table + + result = sanity_check(data_root.as_posix(), checksum=True) + assert not [report.id for report in result.reports if report.is_failed()] + assert {report.id for report in result.reports if not report.is_passed(strict=True)} == { + "STR001", # no version directory + "STR005", # no input_bag + "STR009", # no pointcloud map + } + + +def test_write_data(tmp_path: Path) -> None: + """LiDAR and camera frames written by the writer match the sample's sensor rows.""" + dataset = load_dataset(SAMPLE_V1) + writer = DatasetWriter(tmp_path / "sample") + + lidar = _sensor(dataset, "LIDAR_TOP") + fields = [(field.name, field.dtype) for field in lidar.lidar.fields] + frames = [ + load_pointcloud(SAMPLE_V1 / lidar.lidar.container.path, frame_index=i) + for i in range(lidar.lidar.container.frame_count) + ] + path, frame_count = writer.write_lidar("LIDAR_TOP", frames, fields) + assert (path, frame_count) == (lidar.lidar.container.path, lidar.lidar.container.frame_count) + for i, frame in enumerate(frames): + written = load_pointcloud(tmp_path / "sample" / path, frame_index=i) + assert all((written[name] == frame[name]).all() for name, _ in fields) + + camera = _sensor(dataset, "CAM_FRONT") + images = ( + load_image(SAMPLE_V1 / camera.camera.video.path, frame_index=i) + for i in range(camera.camera.video.frame_count) + ) + video = writer.write_camera("CAM_FRONT", images, frame_rate=camera.camera.video.timescale) + assert video == serialize_dataclass(camera.camera.video) + + +def test_existing_folder(tmp_path: Path) -> None: + with pytest.raises(FileExistsError): + DatasetWriter(tmp_path) + + +def test_channel_written_once(tmp_path: Path) -> None: + """A channel has one file: writing it again is refused, not overwritten.""" + writer = DatasetWriter(tmp_path / "sample") + fields = [("x", "float32")] + writer.write_lidar("LIDAR_TOP", [{"x": np.zeros(2, dtype=np.float32)}], fields) + with pytest.raises(FileExistsError, match="LIDAR_TOP"): + writer.write_lidar("LIDAR_TOP", [{"x": np.ones(2, dtype=np.float32)}], fields)