| """Build a self-contained v3 dataset from completed annotations, one episode at a time.""" |
|
|
| import json |
| from concurrent.futures import ThreadPoolExecutor |
| from copy import deepcopy |
| from fractions import Fraction |
| from pathlib import Path |
|
|
| import av |
| import numpy as np |
| import pandas as pd |
| import pyarrow as pa |
| import pyarrow.compute as pc |
| import pyarrow.parquet as pq |
| from studio import relative_path, validate_subset_features |
|
|
| from lerobot.datasets.compute_stats import aggregate_stats, get_feature_stats |
| from lerobot.datasets.io_utils import load_tasks, write_stats, write_tasks |
|
|
| DATA_PATH = "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet" |
| VIDEO_PATH = "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4" |
| EPISODES_PATH = "meta/episodes/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet" |
|
|
|
|
| def trim_video(source, destination, offset, length, fps, *, lossless=False): |
| """Seek to an episode and encode exactly its frames, with timestamps starting at zero.""" |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| rate = Fraction(str(fps)).limit_denominator(100000) |
| with av.open(str(source)) as src, av.open(str(destination), "w") as dst: |
| stream = src.streams.video[0] |
| stream.codec_context.thread_count = 2 |
| output = dst.add_stream( |
| "libx264rgb" if lossless else "libx264", |
| rate=rate, |
| options={"crf": "0" if lossless else "18", "preset": "fast", "g": "30"}, |
| ) |
| output.width, output.height = stream.codec_context.width, stream.codec_context.height |
| output.pix_fmt = "rgb24" if lossless else "yuv420p" |
| output.codec_context.thread_count = 2 |
| output.time_base = 1 / rate |
| src.seek(int(offset / stream.time_base), stream=stream, backward=True) |
| count = 0 |
| for frame in src.decode(stream): |
| if frame.pts is None: |
| continue |
| timestamp = float(frame.pts * stream.time_base) |
| if timestamp < offset - 0.5 / fps: |
| continue |
| if abs(timestamp - (offset + count / fps)) > 0.51 / fps: |
| raise ValueError("Source video does not match the episode's frame timestamps.") |
| encoded = frame.reformat(format=output.pix_fmt) |
| encoded.pts, encoded.time_base = count, 1 / rate |
| for packet in output.encode(encoded): |
| dst.mux(packet) |
| count += 1 |
| if count == length: |
| break |
| if count != length: |
| raise ValueError(f"Source video ended early: expected {length} frames, found {count}.") |
| for packet in output.encode(): |
| dst.mux(packet) |
| |
| with av.open(str(destination)) as check: |
| if sum(1 for _ in check.decode(video=0)) != length: |
| raise ValueError("Trimmed video frame count does not match the episode.") |
|
|
|
|
| class EpisodeSubset: |
| """Only completed episodes enter this writer. Upload each returned file set atomically.""" |
|
|
| def __init__(self, source_root, root, source_ids, camera_workers=1): |
| if camera_workers not in (1, 2, 3): |
| raise ValueError("Export supports one to three camera workers") |
| self.camera_workers = camera_workers |
| self.source_root, self.root = Path(source_root), Path(root) |
| self.root.mkdir(parents=True, exist_ok=True) |
| self.info = json.loads((self.source_root / "meta/info.json").read_text()) |
| validate_subset_features(self.info) |
| self.source_info = deepcopy(self.info) |
| self.cameras = [k for k, f in self.info["features"].items() if f["dtype"] == "video"] |
| self.source_tasks = load_tasks(self.source_root) |
| metadata = pq.read_table( |
| self.source_root / "meta/episodes", filters=[("episode_index", "in", list(source_ids))] |
| ) |
| self.source_rows = {int(r["episode_index"]): r for r in metadata.to_pylist()} |
| if set(self.source_rows) != set(source_ids): |
| raise ValueError("Episode subset metadata does not match the selected IDs.") |
| self.mapping, self.tasks, self.task_mapping = [], [], {} |
| self.total_frames, self.stats = 0, {} |
| self.info.update(data_path=DATA_PATH, video_path=VIDEO_PATH, chunks_size=1000) |
| for camera in self.cameras: |
| video_info = self.info["features"][camera].setdefault("info", {}) |
| video_info.update( |
| { |
| "video.codec": "h264", |
| "video.pix_fmt": "yuv420p", |
| "video.fps": self.info["fps"], |
| "video.g": 30, |
| "video.crf": 18, |
| "video.preset": "fast", |
| "video.extra_options": {}, |
| "has_audio": False, |
| } |
| ) |
|
|
| def add(self, source_id, table, *, hand_features=None): |
| """Materialize an annotated episode, retaining all its cameras and existing other styles.""" |
| if any(m["source_episode_index"] == source_id for m in self.mapping): |
| raise ValueError("Episode was already added to the subset.") |
| row = self.source_rows[source_id] |
| table = table.filter(pc.equal(table["episode_index"], source_id)) |
| length = int(row["length"]) |
| if len(table) != length or table["frame_index"].to_pylist() != list(range(length)): |
| raise ValueError("Episode data length or frame order does not match its metadata.") |
| if hand_features: |
| description = { |
| "coordinates": "fixed segment-start FLU (X forward, Y left, Z up)", |
| "wrist_position_units": "inferred metres; monocular metric scale is model-dependent", |
| "rotation_units": "axis-angle radians", |
| "validity_feature": "annotation.hand_action_valid", |
| "observed_training_feature": "annotation.hand_action_observed_valid", |
| "motion_infilling": False, |
| "reference_feature": "annotation.hand_reference_frame", |
| "invalid_action": "zero placeholder; must be masked during training", |
| "action_statistics": "computed over valid predictions only; absent when none are available", |
| "minimum_segment_frames": 30, |
| } |
| for name, feature in hand_features.items(): |
| self.info["features"][name] = deepcopy(feature) |
| self.info["features"]["action"]["info"] = description |
| if "action" in self.source_info["features"]: |
| self.info["features"]["annotation.original_action"] = deepcopy( |
| self.source_info["features"]["action"] |
| ) |
| new_id = len(self.mapping) |
| chunk, file = divmod(new_id, 1000) |
| changes = { |
| "episode_index": [new_id] * length, |
| "index": list(range(self.total_frames, self.total_frames + length)), |
| } |
| source_task_names = {int(v): str(k) for k, v in self.source_tasks["task_index"].items()} |
| task_indices = table["task_index"].to_pylist() |
| for old in sorted(set(task_indices)): |
| name = source_task_names[int(old)] |
| if name not in self.tasks: |
| self.tasks.append(name) |
| self.task_mapping[int(old)] = self.tasks.index(name) |
| changes["task_index"] = [self.task_mapping[int(old)] for old in task_indices] |
| for name, values in changes.items(): |
| table = table.set_column( |
| table.schema.get_field_index(name), |
| table.schema.field(name), |
| pa.array(values, type=table[name].type), |
| ) |
| paths = [] |
| episode = { |
| "episode_index": new_id, |
| "tasks": list(dict.fromkeys(source_task_names[int(i)] for i in task_indices)), |
| "length": length, |
| "dataset_from_index": self.total_frames, |
| "dataset_to_index": self.total_frames + length, |
| "data/chunk_index": chunk, |
| "data/file_index": file, |
| "meta/episodes/chunk_index": chunk, |
| "meta/episodes/file_index": file, |
| } |
|
|
| def trim_camera(camera): |
| source_video = self.source_root / relative_path( |
| self.source_info["video_path"].format( |
| video_key=camera, |
| chunk_index=int(row[f"videos/{camera}/chunk_index"]), |
| file_index=int(row[f"videos/{camera}/file_index"]), |
| episode_index=row.get("episode_index"), |
| ) |
| ) |
| target = self.root / relative_path( |
| VIDEO_PATH.format(video_key=camera, chunk_index=chunk, file_index=file) |
| ) |
| trim_video( |
| source_video, |
| target, |
| float(row[f"videos/{camera}/from_timestamp"]), |
| length, |
| float(self.info["fps"]), |
| ) |
| return target |
|
|
| if self.camera_workers > 1: |
| with ThreadPoolExecutor(max_workers=min(self.camera_workers, len(self.cameras))) as pool: |
| targets = list(pool.map(trim_camera, self.cameras)) |
| else: |
| targets = [trim_camera(camera) for camera in self.cameras] |
| for camera, target in zip(self.cameras, targets, strict=True): |
| paths.append(target) |
| episode.update( |
| { |
| f"videos/{camera}/chunk_index": chunk, |
| f"videos/{camera}/file_index": file, |
| f"videos/{camera}/from_timestamp": 0.0, |
| f"videos/{camera}/to_timestamp": length / self.info["fps"], |
| } |
| ) |
| data_path = self.root / DATA_PATH.format(chunk_index=chunk, file_index=file) |
| data_path.parent.mkdir(parents=True, exist_ok=True) |
| pq.write_table(table, data_path, row_group_size=length) |
| paths.append(data_path) |
| |
| episode_stats = {} |
| for key, value in row.items(): |
| if key.startswith("stats/") and value is not None: |
| feature, stat = key[len("stats/") :].rsplit("/", 1) |
| if feature in self.info["features"]: |
| episode_stats.setdefault(feature, {})[stat] = np.atleast_1d(np.asarray(value)) |
| for feature in changes: |
| episode_stats[feature] = get_feature_stats(np.asarray(changes[feature]), axis=0, keepdims=True) |
| if hand_features: |
| for feature in [*hand_features, "annotation.original_action"]: |
| if feature in table.column_names: |
| values = np.asarray(table[feature].to_pylist()) |
| if feature == "action": |
| values = values[ |
| np.asarray(table["annotation.hand_action_valid"].to_pylist(), dtype=bool) |
| ] |
| if not len(values): |
| episode_stats.pop(feature, None) |
| continue |
| episode_stats[feature] = get_feature_stats(values, axis=0, keepdims=values.ndim == 1) |
| for feature, stats in episode_stats.items(): |
| for stat, value in stats.items(): |
| episode[f"stats/{feature}/{stat}"] = value.tolist() |
| episode_path = self.root / EPISODES_PATH.format(chunk_index=chunk, file_index=file) |
| episode_path.parent.mkdir(parents=True, exist_ok=True) |
| episode_table = pa.Table.from_pylist([episode]) |
| if hand_features: |
| |
| |
| |
| for stat in ("min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"): |
| name = f"stats/action/{stat}" |
| values = episode.get(name, [0] if stat == "count" else None) |
| dtype = pa.list_(pa.int64() if stat == "count" else pa.float64()) |
| array = pa.array([values], type=dtype) |
| index = episode_table.schema.get_field_index(name) |
| episode_table = ( |
| episode_table.set_column(index, name, array) |
| if index >= 0 |
| else episode_table.append_column(name, array) |
| ) |
| pq.write_table(episode_table, episode_path) |
| paths.append(episode_path) |
| self.stats = aggregate_stats([self.stats, episode_stats]) |
| self.total_frames += length |
| self.mapping.append({"episode_index": new_id, "source_episode_index": source_id}) |
| self.info.update( |
| total_episodes=len(self.mapping), |
| total_frames=self.total_frames, |
| total_tasks=len(self.tasks), |
| splits={"train": f"0:{len(self.mapping)}"}, |
| ) |
| if "total_videos" in self.info: |
| self.info["total_videos"] = len(self.mapping) * len(self.cameras) |
| if "total_chunks" in self.info: |
| self.info["total_chunks"] = chunk + 1 |
| |
| annotated_info = json.loads((self.source_root / "meta/info.json").read_text()) |
| for name, feature in annotated_info["features"].items(): |
| if feature["dtype"] == "language": |
| self.info["features"][name] = feature |
| if "tools" in annotated_info: |
| self.info["tools"] = annotated_info["tools"] |
| info_path = self.root / "meta/info.json" |
| info_path.write_text(json.dumps(self.info, indent=2) + "\n") |
| write_stats(self.stats, self.root) |
| write_tasks( |
| pd.DataFrame({"task_index": range(len(self.tasks))}, index=pd.Index(self.tasks, name="task")), |
| self.root, |
| ) |
| map_path = self.root / "annotation_studio/episode_mapping.json" |
| map_path.parent.mkdir(parents=True, exist_ok=True) |
| map_path.write_text(json.dumps(self.mapping, indent=2) + "\n") |
| paths.extend([info_path, self.root / "meta/stats.json", self.root / "meta/tasks.parquet", map_path]) |
| return new_id, paths |
|
|
