| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """LeRobot frame conversion and batch writing.""" | |
| import io | |
| import logging | |
| import math | |
| import numbers | |
| from array import array | |
| from bisect import bisect_left | |
| from pathlib import Path | |
| from urllib.parse import unquote, urlparse | |
| import pyarrow as pa | |
| from pypaimon.multimodal.arrow_utils import strict_arrow_table | |
| from pypaimon.multimodal.hdf5 import _SnapshotRecorder | |
| from pypaimon.multimodal.lerobot.schema import _feature_shape | |
| from pypaimon.multimodal.table import _target_schema | |
| from pypaimon.table.row.blob import VideoFrameDescriptor | |
| _DECLARED_NUMERIC_RANGES = { | |
| "uint8": (0, 255), | |
| "uint16": (0, 65535), | |
| "uint32": (0, 4294967295), | |
| "float16": (-65504.0, 65504.0), | |
| "float32": (-3.4028234663852886e38, 3.4028234663852886e38), | |
| } | |
| _NUMERIC_DTYPES = { | |
| "int8", | |
| "int16", | |
| "int32", | |
| "int64", | |
| "uint8", | |
| "uint16", | |
| "uint32", | |
| "float16", | |
| "float32", | |
| "float64", | |
| } | |
| _BOOLEAN_DTYPES = {"bool", "boolean"} | |
| _VIDEO_TIMESTAMP_TOLERANCE = 1e-4 | |
| _LOGGER = logging.getLogger(__name__) | |
| def _strict_lerobot_table(data, target_schema, source, batch_index): | |
| return strict_arrow_table( | |
| data, | |
| target_schema, | |
| source.path, | |
| batch_index, | |
| "LeRobot", | |
| ) | |
| def _write_dataset( | |
| table, | |
| dataset, | |
| info, | |
| source, | |
| source_schema, | |
| batch_size, | |
| metadata, | |
| video_fields=()): | |
| table_commit = table.raw_table.new_batch_write_builder().new_commit() | |
| try: | |
| messages = _prepare_dataset( | |
| table.raw_table, dataset, info, source, source_schema, batch_size, | |
| metadata, video_fields) | |
| return _commit_messages(table_commit, messages) | |
| finally: | |
| table_commit.close() | |
| def _commit_messages(table_commit, messages): | |
| snapshot_recorder = _SnapshotRecorder() | |
| table_commit.add_commit_callback(snapshot_recorder) | |
| table_commit.commit(messages) | |
| if snapshot_recorder.snapshot_id is None: | |
| raise RuntimeError( | |
| "LeRobot append committed without reporting a snapshot id.") | |
| return snapshot_recorder.snapshot_id | |
| def _abort_messages(table, messages): | |
| if not messages: | |
| return | |
| commit = table.new_batch_write_builder().new_commit() | |
| try: | |
| commit.abort(messages) | |
| finally: | |
| commit.close() | |
| def _prepare_dataset( | |
| table, dataset, info, source, source_schema, batch_size, metadata, | |
| video_fields=(), episode_indices=None, transform=None): | |
| from pypaimon.common.options import CoreOptions | |
| remux_videos = table.options.options.get(CoreOptions.WRITE_LEROBOT_REMUX_EPISODES) | |
| if video_fields and remux_videos: | |
| from pypaimon.multimodal.video_remux import _require_remux_av | |
| _require_remux_av() | |
| target_schema = _target_schema(table) | |
| write_builder = table.new_batch_write_builder() | |
| table_write = None | |
| prepared = False | |
| messages = None | |
| batch_count = 0 | |
| row_count = 0 | |
| episodes = metadata["episodes"] | |
| current_episode = None | |
| expected_tasks = set() | |
| observed_tasks = set() | |
| video_sources = {} | |
| source_factory = None | |
| try: | |
| table_write = write_builder.new_write() | |
| reader_factory = getattr( | |
| dataset, "video_uri_reader_factory", None) | |
| if video_fields and remux_videos: | |
| from pypaimon.common.uri_reader import UriReaderFactory | |
| from pypaimon.multimodal.video_remux import _EpisodeVideoSources | |
| if reader_factory is None: | |
| source_factory = UriReaderFactory({}) | |
| reader_factory = source_factory | |
| video_sources = _EpisodeVideoSources(reader_factory) | |
| reader_factory = video_sources | |
| if video_fields and reader_factory is not None: | |
| table_write.with_blob_uri_reader_factory(reader_factory) | |
| for source_episode, episode_index, episode_begin, task_indices, \ | |
| begin, end in _episode_batches( | |
| info, | |
| batch_size, | |
| episodes, | |
| video_fields, | |
| episode_indices, | |
| ): | |
| if episode_index != current_episode: | |
| if current_episode is not None: | |
| _validate_episode_tasks( | |
| current_episode, expected_tasks, observed_tasks) | |
| current_episode = episode_index | |
| expected_tasks = set(task_indices) | |
| observed_tasks = set() | |
| if video_fields: | |
| table_write.roll_before_group_if_needed( | |
| int(source_episode["length"])) | |
| if remux_videos: | |
| video_sources.next_episode() | |
| batch = _read_batch( | |
| dataset, | |
| info, | |
| begin, | |
| end, | |
| source_schema, | |
| episode=source_episode, | |
| video_sources=video_sources, | |
| ) | |
| seen_tasks = _validate_frame_controls( | |
| batch, | |
| int(info["fps"]), | |
| episode_index, | |
| episode_begin, | |
| begin, | |
| task_indices, | |
| metadata["subtask_indices"], | |
| ) | |
| observed_tasks.update(seen_tasks) | |
| if transform is not None: | |
| batch = transform(batch, episode_index) | |
| batch = pa.Table.from_arrays([ | |
| batch[field.name] if field.name in batch.column_names | |
| else pa.nulls(batch.num_rows, type=field.type) | |
| for field in target_schema | |
| ], schema=target_schema) | |
| batch = _strict_lerobot_table( | |
| batch, | |
| target_schema, | |
| source, | |
| batch_count, | |
| ) | |
| table_write.write_arrow(batch) | |
| batch_count += 1 | |
| row_count += batch.num_rows | |
| if current_episode is not None: | |
| _validate_episode_tasks( | |
| current_episode, expected_tasks, observed_tasks) | |
| expected_rows = (int(info["total_frames"]) if episode_indices is None | |
| else sum(episodes[i]["length"] for i in episode_indices)) | |
| if row_count != expected_rows: | |
| raise ValueError( | |
| "LeRobot metadata reports %d frames but import produced %d." | |
| % (expected_rows, row_count)) | |
| messages = table_write.prepare_commit() | |
| prepared = True | |
| return messages | |
| except BaseException: | |
| if table_write is not None and not prepared: | |
| table_write.abort() | |
| raise | |
| finally: | |
| try: | |
| if table_write is not None: | |
| try: | |
| table_write.close() | |
| except BaseException: | |
| # prepare_commit handed off these files; writer.abort() no | |
| # longer owns them, and no caller has received the messages. | |
| if messages is not None: | |
| try: | |
| _abort_messages(table, messages) | |
| except Exception: | |
| _LOGGER.exception("Failed to abort prepared LeRobot files.") | |
| raise | |
| finally: | |
| if source_factory is not None: | |
| source_factory.close() | |
| def _episode_batches( | |
| info, batch_size, episodes, video_fields=(), episode_indices=None): | |
| episode_count = int(info.get("total_episodes", 0)) | |
| total_frames = int(info["total_frames"]) | |
| expected_begin = 0 | |
| ordinals = range(episode_count) if episode_indices is None else episode_indices | |
| for ordinal in ordinals: | |
| episode = episodes.iloc[ordinal] if hasattr(episodes, "iloc") \ | |
| else episodes[ordinal] | |
| episode_index = int(_python_scalar(episode["episode_index"])) | |
| begin = int(_python_scalar(episode["dataset_from_index"])) | |
| end = int(_python_scalar(episode["dataset_to_index"])) | |
| length = int(_python_scalar(episode["length"])) | |
| if episode_index != ordinal or (episode_indices is None and begin != expected_begin) \ | |
| or end <= begin or length != end - begin: | |
| raise ValueError( | |
| "LeRobot episode %d has an invalid index, range, or length." | |
| % ordinal) | |
| episode_begin = begin | |
| task_indices = episode.get("task_indices", ()) | |
| source_episode = episode if video_fields else None | |
| while begin < end: | |
| batch_end = min(begin + batch_size, end) | |
| yield ( | |
| source_episode, | |
| episode_index, | |
| episode_begin, | |
| task_indices, | |
| begin, | |
| batch_end, | |
| ) | |
| begin = batch_end | |
| expected_begin = end | |
| if episode_indices is None and expected_begin != total_frames: | |
| raise ValueError( | |
| "LeRobot episode ranges cover %d frames but metadata reports %d." | |
| % (expected_begin, total_frames)) | |
| def _validate_frame_controls( | |
| batch, | |
| fps, | |
| episode_index, | |
| episode_begin, | |
| begin, | |
| task_indices, | |
| subtask_indices=None): | |
| required = [ | |
| "index", "episode_index", "frame_index", "timestamp", "task_index" | |
| ] | |
| if subtask_indices is not None: | |
| required.append("subtask_index") | |
| missing = [name for name in required if name not in batch.column_names] | |
| if missing: | |
| raise ValueError( | |
| "LeRobot frame data is missing control columns: %s." | |
| % ", ".join(missing)) | |
| allowed_tasks = set(task_indices) | |
| values = { | |
| name: batch.column(name).to_pylist() | |
| for name in required | |
| } | |
| seen_tasks = set() | |
| for offset in range(batch.num_rows): | |
| index = begin + offset | |
| frame_index = index - episode_begin | |
| _require_control_integer( | |
| values["index"][offset], "index", index, index) | |
| _require_control_integer( | |
| values["episode_index"][offset], | |
| "episode_index", | |
| episode_index, | |
| index, | |
| ) | |
| _require_control_integer( | |
| values["frame_index"][offset], | |
| "frame_index", | |
| frame_index, | |
| index, | |
| ) | |
| timestamp = values["timestamp"][offset] | |
| expected_timestamp = pa.scalar( | |
| frame_index / fps, | |
| type=batch.schema.field("timestamp").type, | |
| ).as_py() | |
| if (isinstance(timestamp, bool) | |
| or not isinstance(timestamp, numbers.Real) | |
| or not math.isclose( | |
| float(timestamp), float(expected_timestamp), | |
| rel_tol=0.0, abs_tol=1e-4)): | |
| raise ValueError( | |
| "LeRobot frame %d has timestamp %r; expected %r." | |
| % (index, timestamp, expected_timestamp)) | |
| task_index = _control_integer( | |
| values["task_index"][offset], "task_index", index) | |
| if task_index not in allowed_tasks: | |
| raise ValueError( | |
| "LeRobot frame %d has task_index %d outside Episode %d " | |
| "tasks %s." | |
| % (index, task_index, episode_index, | |
| sorted(allowed_tasks))) | |
| seen_tasks.add(task_index) | |
| if subtask_indices is not None: | |
| if values["subtask_index"][offset] is None: | |
| continue | |
| subtask_index = _control_integer( | |
| values["subtask_index"][offset], "subtask_index", index) | |
| if subtask_index not in subtask_indices: | |
| raise ValueError( | |
| "LeRobot frame %d has subtask_index %d outside [0, %d)." | |
| % (index, subtask_index, len(subtask_indices))) | |
| return seen_tasks | |
| def _validate_episode_tasks(episode_index, expected, actual): | |
| if actual != expected: | |
| raise ValueError( | |
| "LeRobot Episode %d declares task indices %s but its " | |
| "frames use %s." | |
| % (episode_index, sorted(expected), sorted(actual))) | |
| def _require_control_integer(value, name, expected, frame_index): | |
| actual = _control_integer(value, name, frame_index) | |
| if actual != expected: | |
| raise ValueError( | |
| "LeRobot frame %d has %s %d; expected %d." | |
| % (frame_index, name, actual, expected)) | |
| def _control_integer(value, name, frame_index): | |
| value = _python_scalar(value) | |
| if isinstance(value, bool) or not isinstance(value, numbers.Integral): | |
| raise ValueError( | |
| "LeRobot frame %d has non-integer %s %r." | |
| % (frame_index, name, value)) | |
| return int(value) | |
| def _read_batch( | |
| dataset, info, begin, end, schema, episode=None, | |
| video_sources=None): | |
| read_batch = getattr(dataset, "read_batch", None) | |
| if callable(read_batch): | |
| raw = read_batch(begin, end) | |
| else: | |
| raw = dataset.hf_dataset.with_format("arrow")[begin:end] | |
| if isinstance(raw, pa.RecordBatch): | |
| raw = pa.Table.from_batches([raw]) | |
| elif not isinstance(raw, pa.Table): | |
| raw = pa.Table.from_pydict(raw) | |
| features = info["features"] | |
| video_rows = None | |
| video_timestamp_type = None | |
| if any(feature.get("dtype") == "video" | |
| for feature in features.values()): | |
| video_rows = _validate_video_rows( | |
| raw, info, episode, begin, end) | |
| video_timestamp_type = raw.schema.field("timestamp").type | |
| arrays = [] | |
| fields = [] | |
| for name, feature in features.items(): | |
| field = schema.field(name) | |
| dtype = feature["dtype"] | |
| if dtype == "video": | |
| values = _video_frame_descriptors( | |
| dataset, | |
| info, | |
| episode, | |
| video_rows, | |
| name, | |
| feature, | |
| begin, | |
| end, | |
| video_sources if video_sources is not None else {}, | |
| video_timestamp_type, | |
| ) | |
| elif name not in raw.column_names: | |
| raise ValueError( | |
| "LeRobot data is missing metadata feature %s." % name) | |
| else: | |
| values = raw.column(name).to_pylist() | |
| if dtype == "image": | |
| image_reader = getattr(dataset, "image_bytes", None) | |
| if callable(image_reader): | |
| values = [image_reader(value) for value in values] | |
| else: | |
| values = [_image_bytes(value, dataset.root) | |
| for value in values] | |
| else: | |
| values = [_normalize_value(value, feature, name) | |
| for value in values] | |
| arrays.append(_safe_array(values, field, name, dtype)) | |
| fields.append(field) | |
| return pa.Table.from_arrays(arrays, schema=pa.schema(fields)) | |
| def _video_frame_descriptors( | |
| dataset, info, episode, video_rows, name, feature, begin, end, cache, | |
| timestamp_type): | |
| if episode is None: | |
| raise ValueError("LeRobot video import requires Episode metadata.") | |
| episode_begin = _nonnegative_integer( | |
| episode["dataset_from_index"], "dataset_from_index") | |
| episode_end = _nonnegative_integer( | |
| episode["dataset_to_index"], "dataset_to_index") | |
| if begin < episode_begin or end > episode_end: | |
| raise ValueError( | |
| "LeRobot video batch [%d, %d) crosses Episode range [%d, %d)." | |
| % (begin, end, episode_begin, episode_end) | |
| ) | |
| fps = _video_fps(info, feature, name) | |
| prefix = "videos/%s/" % name | |
| try: | |
| chunk_index = _nonnegative_integer( | |
| episode[prefix + "chunk_index"], prefix + "chunk_index") | |
| file_index = _nonnegative_integer( | |
| episode[prefix + "file_index"], prefix + "file_index") | |
| from_timestamp = float(_python_scalar( | |
| episode[prefix + "from_timestamp"])) | |
| to_timestamp = float(_python_scalar( | |
| episode[prefix + "to_timestamp"])) | |
| except (KeyError, TypeError, ValueError) as error: | |
| raise ValueError( | |
| "LeRobot Episode metadata is missing video mapping for %s." | |
| % name | |
| ) from error | |
| episode_length = episode_end - episode_begin | |
| _validate_video_timestamp_range( | |
| from_timestamp, to_timestamp, fps, episode_length, name) | |
| source_key = (name, chunk_index, file_index) | |
| source = cache.get(source_key) | |
| if source is None: | |
| uri, length = _video_source( | |
| dataset, info, episode, name, chunk_index, file_index) | |
| source = (uri, length, _video_sample_timestamps(dataset, uri)) | |
| cache[source_key] = source | |
| uri, length, sample_timestamps = source | |
| from pypaimon.multimodal.video_remux import _EpisodeVideoSources | |
| if isinstance(cache, _EpisodeVideoSources): | |
| ordinals = [ | |
| _video_frame_ordinal(sample_timestamps, timestamp, from_timestamp, | |
| timestamp_type, name) | |
| for unused_frame_index, timestamp in video_rows | |
| ] | |
| return cache.descriptors( | |
| uri, from_timestamp, to_timestamp, | |
| [sample_timestamps[ordinal] for ordinal in ordinals]) | |
| return [ | |
| VideoFrameDescriptor( | |
| uri, | |
| 0, | |
| length, | |
| _video_frame_ordinal( | |
| sample_timestamps, | |
| timestamp, | |
| from_timestamp, | |
| timestamp_type, | |
| name, | |
| ), | |
| ).serialize() | |
| for unused_frame_index, timestamp in video_rows | |
| ] | |
| def _video_sample_timestamps(dataset, uri): | |
| resolver = getattr(dataset, "video_sample_timestamps", None) | |
| if callable(resolver): | |
| values = resolver(uri) | |
| else: | |
| try: | |
| import av | |
| except ImportError as error: | |
| raise ImportError( | |
| "LeRobot video import requires PyAV. Install it with " | |
| "`pip install 'pypaimon[lerobot]'`.") from error | |
| input_stream = None | |
| parsed = urlparse(uri) | |
| if parsed.scheme in ("", "file"): | |
| source = unquote(parsed.path) if parsed.scheme else uri | |
| else: | |
| factory = getattr(dataset, "video_uri_reader_factory", None) | |
| if factory is None: | |
| raise ValueError( | |
| "LeRobot video source %s cannot be inspected." % uri) | |
| input_stream = factory.create(uri).new_input_stream(uri) | |
| source = input_stream | |
| try: | |
| with av.open(source) as container: | |
| stream = container.streams.video[0] | |
| values = [ | |
| float(packet.pts * (packet.time_base or stream.time_base)) | |
| for packet in container.demux(stream) | |
| if packet.pts is not None and not packet.is_discard | |
| ] | |
| finally: | |
| if input_stream is not None: | |
| input_stream.close() | |
| timestamps = array("d", sorted(float(value) for value in values)) | |
| if not timestamps or any(not math.isfinite(value) for value in timestamps): | |
| raise ValueError( | |
| "LeRobot video source %s has no valid frame timestamps." % uri) | |
| return timestamps | |
| def _video_frame_ordinal( | |
| timestamps, frame_timestamp, from_timestamp, timestamp_type, name): | |
| timestamp = from_timestamp + frame_timestamp | |
| position = bisect_left(timestamps, timestamp) | |
| candidates = [] | |
| if position: | |
| candidates.append(position - 1) | |
| if position < len(timestamps): | |
| candidates.append(position) | |
| ordinal = min( | |
| candidates, | |
| key=lambda index: (abs(timestamps[index] - timestamp), index), | |
| ) | |
| expected = pa.scalar( | |
| timestamps[ordinal] - from_timestamp, type=timestamp_type).as_py() | |
| if not math.isclose( | |
| frame_timestamp, | |
| float(expected), | |
| rel_tol=0.0, | |
| abs_tol=_VIDEO_TIMESTAMP_TOLERANCE): | |
| raise ValueError( | |
| "LeRobot video feature %s has no frame matching timestamp %s." | |
| % (name, timestamp)) | |
| return ordinal | |
| def _validate_video_rows(raw, info, episode, begin, end): | |
| if episode is None: | |
| raise ValueError("LeRobot video import requires Episode metadata.") | |
| required = ("episode_index", "frame_index", "timestamp") | |
| missing = [name for name in required if name not in raw.column_names] | |
| if missing: | |
| raise ValueError( | |
| "LeRobot video import requires frame columns %s." | |
| % ", ".join(missing) | |
| ) | |
| episode_index = _nonnegative_integer( | |
| episode["episode_index"], "episode_index") | |
| episode_begin = _nonnegative_integer( | |
| episode["dataset_from_index"], "dataset_from_index") | |
| expected_frames = range(begin - episode_begin, end - episode_begin) | |
| actual_episodes = raw.column("episode_index").to_pylist() | |
| actual_frames = raw.column("frame_index").to_pylist() | |
| timestamps = raw.column("timestamp").to_pylist() | |
| timestamp_type = raw.schema.field("timestamp").type | |
| fps = _positive_fps(info.get("fps"), "dataset") | |
| result = [] | |
| for offset, expected_frame in enumerate(expected_frames): | |
| actual_episode = _nonnegative_integer( | |
| actual_episodes[offset], "episode_index") | |
| actual_frame = _nonnegative_integer( | |
| actual_frames[offset], "frame_index") | |
| if actual_episode != episode_index or actual_frame != expected_frame: | |
| raise ValueError( | |
| "LeRobot frame rows do not match Episode %d range [%d, %d)." | |
| % (episode_index, begin, end) | |
| ) | |
| timestamp = timestamps[offset] | |
| expected_timestamp = pa.scalar( | |
| expected_frame / fps, type=timestamp_type).as_py() | |
| if (isinstance(timestamp, bool) | |
| or not isinstance(timestamp, numbers.Real) | |
| or not math.isclose( | |
| float(timestamp), float(expected_timestamp), | |
| rel_tol=0.0, | |
| abs_tol=_VIDEO_TIMESTAMP_TOLERANCE)): | |
| raise ValueError( | |
| "LeRobot frame %d has timestamp %r; expected %r." | |
| % (actual_frame, timestamp, expected_timestamp) | |
| ) | |
| result.append((actual_frame, float(timestamp))) | |
| return result | |
| def _validate_video_timestamp_range( | |
| from_timestamp, to_timestamp, fps, episode_length, name): | |
| if (not math.isfinite(from_timestamp) | |
| or not math.isfinite(to_timestamp) | |
| or from_timestamp < 0 | |
| or to_timestamp <= from_timestamp): | |
| raise ValueError( | |
| "LeRobot video feature %s has invalid timestamp range [%s, %s)." | |
| % (name, from_timestamp, to_timestamp) | |
| ) | |
| duration = to_timestamp - from_timestamp | |
| expected_duration = episode_length / fps | |
| if not math.isclose( | |
| duration, | |
| expected_duration, | |
| rel_tol=0.0, | |
| abs_tol=_VIDEO_TIMESTAMP_TOLERANCE): | |
| raise ValueError( | |
| "LeRobot video feature %s has duration %s, but Episode with %d " | |
| "frames at FPS %s requires %s." | |
| % (name, duration, episode_length, fps, expected_duration) | |
| ) | |
| def _video_fps(info, feature, name): | |
| global_fps = _positive_fps(info.get("fps"), "dataset") | |
| values = [] | |
| for key in ("info", "video_info"): | |
| details = feature.get(key) | |
| if isinstance(details, dict) and "video.fps" in details: | |
| values.append(details["video.fps"]) | |
| if "fps" in feature: | |
| values.append(feature["fps"]) | |
| value = values[0] if values else global_fps | |
| fps = _positive_fps(value, "video feature %s" % name) | |
| if any(not math.isclose( | |
| fps, _positive_fps(other, "video feature %s" % name), | |
| rel_tol=1e-6, abs_tol=1e-6) for other in values[1:]) \ | |
| or not math.isclose( | |
| fps, global_fps, rel_tol=1e-6, abs_tol=1e-6): | |
| raise ValueError( | |
| "LeRobot video feature %s FPS does not match dataset FPS." | |
| % name | |
| ) | |
| return fps | |
| def _positive_fps(value, owner): | |
| try: | |
| fps = float(value) | |
| except (TypeError, ValueError) as error: | |
| raise ValueError( | |
| "LeRobot %s is missing a valid FPS." % owner | |
| ) from error | |
| if not math.isfinite(fps) or fps <= 0: | |
| raise ValueError( | |
| "LeRobot %s is missing a valid FPS." % owner) | |
| return fps | |
| def _nonnegative_integer(value, name): | |
| value = _python_scalar(value) | |
| if isinstance(value, bool) \ | |
| or not isinstance(value, numbers.Integral) or value < 0: | |
| raise ValueError( | |
| "LeRobot %s must be a non-negative integer; found %r." | |
| % (name, value)) | |
| return int(value) | |
| def _video_source( | |
| dataset, info, episode, name, chunk_index, file_index): | |
| resolver = getattr(dataset, "video_source", None) | |
| if callable(resolver): | |
| return resolver(name, episode) | |
| template = info.get("video_path") | |
| if not isinstance(template, str) or not template: | |
| raise ValueError("LeRobot v3 metadata is missing info.video_path.") | |
| relative = template.format( | |
| video_key=name, | |
| chunk_index=chunk_index, | |
| file_index=file_index, | |
| ) | |
| root = Path(dataset.root).resolve() | |
| path = (root / relative).resolve() | |
| try: | |
| path.relative_to(root) | |
| except ValueError as error: | |
| raise ValueError( | |
| "LeRobot video path must stay within the source directory: %s" | |
| % relative | |
| ) from error | |
| if not path.is_file(): | |
| raise FileNotFoundError("LeRobot video file does not exist: %s" % path) | |
| length = path.stat().st_size | |
| if length <= 0: | |
| raise ValueError("LeRobot video file is empty: %s" % path) | |
| return path.as_uri(), length | |
| def _safe_array(values, field, name, dtype): | |
| _validate_declared_range(values, field.type, name, dtype) | |
| try: | |
| return pa.array(values).cast(field.type, safe=True) | |
| except (pa.ArrowException, TypeError, ValueError, OverflowError) as error: | |
| raise ValueError( | |
| "LeRobot feature %s cannot be safely converted to %s: %s" | |
| % (name, field.type, error)) from error | |
| def _validate_declared_range(values, target_type, name, dtype): | |
| if (pa.types.is_list(target_type) | |
| or pa.types.is_large_list(target_type) | |
| or pa.types.is_fixed_size_list(target_type)): | |
| for value in values: | |
| if value is not None: | |
| _validate_declared_range( | |
| value, target_type.value_type, name, dtype) | |
| return | |
| value_range = _DECLARED_NUMERIC_RANGES.get(dtype) | |
| for value in values: | |
| if value is None: | |
| continue | |
| _validate_value_domain(value, name, dtype) | |
| if value_range is None: | |
| continue | |
| minimum, maximum = value_range | |
| if dtype.startswith("float") and not isinstance( | |
| value, numbers.Integral): | |
| try: | |
| if not math.isfinite(value): | |
| continue | |
| except (TypeError, ValueError, OverflowError) as error: | |
| raise ValueError( | |
| "LeRobot feature %s contains a value incompatible with " | |
| "dtype %s: %r" % (name, dtype, value)) from error | |
| try: | |
| out_of_range = value < minimum or value > maximum | |
| except (TypeError, ValueError, OverflowError) as error: | |
| raise ValueError( | |
| "LeRobot feature %s contains a value incompatible with " | |
| "dtype %s: %r" % (name, dtype, value)) from error | |
| if out_of_range: | |
| raise ValueError( | |
| "LeRobot feature %s contains a value outside the %s " | |
| "range [%s, %s]: %r" | |
| % (name, dtype, minimum, maximum, value)) | |
| def _validate_value_domain(value, name, dtype): | |
| if dtype in _NUMERIC_DTYPES: | |
| valid = isinstance(value, numbers.Real) \ | |
| and not isinstance(value, bool) | |
| expected = "numeric" | |
| elif dtype in _BOOLEAN_DTYPES: | |
| valid = isinstance(value, bool) | |
| expected = "boolean" | |
| elif dtype == "string": | |
| valid = isinstance(value, str) | |
| expected = "string" | |
| else: | |
| return | |
| if not valid: | |
| raise ValueError( | |
| "LeRobot feature %s declares dtype %s but contains a non-%s " | |
| "value: %r" % (name, dtype, expected, value)) | |
| def _normalize_value(value, feature, name): | |
| shape = _feature_shape(feature, name) | |
| if shape in ((), (1,)): | |
| if isinstance(value, (list, tuple)): | |
| if len(value) != 1: | |
| raise ValueError( | |
| "LeRobot feature %s expected shape %s, got %s." | |
| % (name, shape, _value_shape(value))) | |
| return value[0] | |
| return _python_scalar(value) | |
| actual_shape = _value_shape(value) | |
| if actual_shape != shape: | |
| raise ValueError( | |
| "LeRobot feature %s expected shape %s, got %s." | |
| % (name, shape, actual_shape)) | |
| return value | |
| def _value_shape(value): | |
| if hasattr(value, "shape"): | |
| return tuple(int(size) for size in value.shape) | |
| if isinstance(value, (list, tuple)): | |
| if not value: | |
| return (0,) | |
| child = _value_shape(value[0]) | |
| if any(_value_shape(item) != child for item in value[1:]): | |
| return (len(value), -1) | |
| return (len(value),) + child | |
| return () | |
| def _python_scalar(value): | |
| item = getattr(value, "item", None) | |
| if callable(item): | |
| return item() | |
| return value | |
| def _image_bytes(value, root): | |
| if value is None: | |
| raise ValueError("LeRobot image feature contains a null frame.") | |
| if isinstance(value, (bytes, bytearray, memoryview)): | |
| return bytes(value) | |
| if isinstance(value, dict): | |
| body = value.get("bytes") | |
| if body is not None: | |
| return bytes(body) | |
| image_path = value.get("path") | |
| if image_path: | |
| root = Path(root).resolve() | |
| path = Path(image_path) | |
| if not path.is_absolute(): | |
| path = root / path | |
| path = path.resolve() | |
| try: | |
| path.relative_to(root) | |
| except ValueError as error: | |
| raise ValueError( | |
| "LeRobot image path must stay within the source " | |
| "directory: %s" % image_path) from error | |
| return path.read_bytes() | |
| return _encode_media_frame(value) | |
| def _encode_media_frame(value, channel_first=None): | |
| try: | |
| import numpy as np | |
| from PIL import Image | |
| except ImportError as error: | |
| raise ImportError( | |
| "LeRobot media import requires numpy and Pillow from the " | |
| "'pypaimon[lerobot]' extra.") from error | |
| if isinstance(value, Image.Image): | |
| image = value | |
| else: | |
| detach = getattr(value, "detach", None) | |
| if callable(detach): | |
| value = detach().cpu().numpy() | |
| array = np.asarray(value) | |
| if channel_first is True or ( | |
| channel_first is None | |
| and array.ndim == 3 | |
| and array.shape[0] in (1, 3, 4)): | |
| array = np.transpose(array, (1, 2, 0)) | |
| if np.issubdtype(array.dtype, np.floating): | |
| array = np.rint(np.clip(array, 0.0, 1.0) * 255.0).astype(np.uint8) | |
| if array.ndim == 3 and array.shape[2] == 1: | |
| array = array[:, :, 0] | |
| try: | |
| image = Image.fromarray(array) | |
| except (KeyError, TypeError, ValueError) as error: | |
| raise ValueError( | |
| "Unsupported LeRobot media frame shape or dtype: %s, %s." | |
| % (array.shape, array.dtype)) from error | |
| output = io.BytesIO() | |
| image.save(output, format="PNG") | |
| return output.getvalue() |