| """Build a self-contained v3 dataset from completed annotations, one episode at a time.""" |
|
|
| import json |
| 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): |
| """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("libx264", rate=rate, options={"crf": "18", "preset": "fast", "g": "30"}) |
| output.width, output.height = stream.codec_context.width, stream.codec_context.height |
| output.pix_fmt = "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="yuv420p") |
| 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): |
| 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): |
| """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.") |
| 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, |
| } |
| for camera in self.cameras: |
| 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"]), |
| ) |
| ) |
| 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"]), |
| ) |
| 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) |
| 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) |
| pq.write_table(pa.Table.from_pylist([episode]), 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 |
|
|
