diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 580a6a2a78a8..90233a70a4dd 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -302,6 +302,25 @@ jobs: pypaimon/tests/multimodal_lerobot_writer_test.py::PaimonLeRobotWriterTest::test_multiframe_boolean_statistics_preserve_extrema \ pypaimon/tests/multimodal_lerobot_writer_test.py::PaimonLeRobotWriterTest::test_multiframe_int64_statistics_preserve_extrema + video-minimum: + name: Video / PyAV 14.2 + timeout-minutes: 15 + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 + with: + python-version: "3.10" + - name: Install minimum video dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -e './paimon-python[video]' 'av==14.2.0' --group ./paimon-python/pyproject.toml:test + python -m pip check + python -c 'import av; assert av.__version__ == "14.2.0"' + - name: Test Episode remux + working-directory: paimon-python + run: python -m pytest -q pypaimon/tests/video_remux_test.py + torch: name: PyTorch / Python 3.10 timeout-minutes: 60 diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 82c81590c733..4e312f354e9f 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -129,6 +129,11 @@ Video features map to `BLOB`. Each MP4 is copied once per aligned file group and reads. Imports require a bucket-unaware table and check rolling before each Episode. Use `VideoFrameCollator` for scans or `PaimonLeRobotDataset` for training. +By default, imports copy the original videos. Set +`options={"write.lerobot.remux-episodes": "true"}` to split them by Episode without +re-encoding (requires PyAV 14.2+), retaining decoding dependencies. Each clip is limited +to 64 MiB and 65,536 packets; unsupported timestamps fail the import. + ## Capture LeRobot frames directly into Paimon `PaimonLeRobotWriter` implements the write-side surface used by LeRobot's diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index d462943b93f9..93e300df8dd7 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -1162,6 +1162,15 @@ class CoreOptions: .with_description("Read batch size for any file format if it supports.") ) + WRITE_LEROBOT_REMUX_EPISODES: ConfigOption[bool] = ( + ConfigOptions.key("write.lerobot.remux-episodes") + .boolean_type() + .default_value(False) + .with_description( + "Copy each LeRobot Episode's video packets into a separate MP4 " + "payload without re-encoding. Requires PyAV.") + ) + READ_VIDEO_MAX_OPEN_DECODERS: ConfigOption[int] = ( ConfigOptions.key("read.video.max-open-decoders") .int_type() diff --git a/paimon-python/pypaimon/multimodal/lerobot/loader.py b/paimon-python/pypaimon/multimodal/lerobot/loader.py index b664f5c4df8e..1f68255e96be 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/loader.py +++ b/paimon-python/pypaimon/multimodal/lerobot/loader.py @@ -110,6 +110,11 @@ def _abort_messages(table, messages): def _prepare_dataset( table, dataset, info, source, source_schema, batch_size, metadata, video_fields=(), episode_indices=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 @@ -122,11 +127,20 @@ def _prepare_dataset( 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, \ @@ -147,6 +161,8 @@ def _prepare_dataset( 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, @@ -194,18 +210,22 @@ def _prepare_dataset( table_write.abort() raise finally: - 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 + 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( @@ -450,6 +470,16 @@ def _video_frame_descriptors( 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, diff --git a/paimon-python/pypaimon/multimodal/video_remux.py b/paimon-python/pypaimon/multimodal/video_remux.py new file mode 100644 index 000000000000..0bc3a364eae5 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/video_remux.py @@ -0,0 +1,190 @@ +# 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. + +"""Copy Episode video packets into bounded MP4 payloads without re-encoding.""" + +import io +from bisect import bisect_left +from fractions import Fraction +from itertools import chain + + +class _BoundedBuffer(io.BytesIO): + def __init__(self, limit): + super().__init__() + self.limit = limit + + def write(self, data): + if self.tell() + len(data) > self.limit: + raise ValueError("Remuxed video exceeds the memory limit.") + return super().write(data) + + +def _require_remux_av(): + import av + from packaging.version import Version + + if Version(av.__version__) < Version("14.2"): + raise ImportError("Episode video remux requires PyAV 14.2 or newer.") + return av + + +def remux_video_interval(source, start, end, max_bytes=64 * 1024 * 1024): + """Return MP4 bytes and (source timestamp, output ordinal) pairs. + + The half-open interval is in seconds. Preserve a preceding keyframe and + one following GOP for decoder dependencies. Copy encoded packets; decode + only the resulting short clip to verify its presentation timestamps. + The mapping includes dependency frames; callers select resolved source PTS. + source is a path or seekable binary stream, owned by the caller. + """ + av = _require_remux_av() + + start = Fraction(str(start)) + end = Fraction(str(end)) + if start < 0 or end <= start or max_bytes <= 0: + raise ValueError("Invalid video interval or memory limit.") + output = _BoundedBuffer(max_bytes) + selected = [] + seen_pts = set() + with av.open(source) as container: + if not container.streams.video: + raise ValueError("Source has no video stream.") + stream = container.streams.video[0] + if stream.time_base is None: + raise ValueError("Video has no time base.") + seek_time = start + backoff = Fraction(1) + while True: + container.seek(int(seek_time / stream.time_base), stream=stream, + backward=True, any_frame=False) + packets = container.demux(stream) + first = next((p for p in packets if p.size), None) + if first is None or first.pts is None: + raise ValueError("Cannot find a preceding keyframe.") + first_time = first.pts * (first.time_base or stream.time_base) + if first.is_keyframe and first_time <= start: + break + # Some demuxers seek by DTS; with B-frames this may land after + # the requested presentation time. Retry an earlier seek point. + if seek_time == 0: + raise ValueError("Cannot find a preceding keyframe.") + seek_time = max(Fraction(0), start - backoff) + backoff *= 2 + with av.open(output, mode="w", format="mp4") as muxer: + target = muxer.add_stream_from_template(stream, opaque=True) + started = False + origin = None + following_gops = 0 + for packet in chain([first], packets): + if not packet.size: + continue + if packet.pts is None or packet.dts is None or packet.is_discard: + raise ValueError("Video packet has ambiguous timestamps.") + timestamp = packet.pts * (packet.time_base or stream.time_base) + if not started: + if not packet.is_keyframe or timestamp > start: + raise ValueError("Cannot find a preceding keyframe.") + started = True + origin = timestamp + if packet.is_keyframe and timestamp >= end: + following_gops += 1 + if following_gops == 2: + break + if timestamp in seen_pts: + raise ValueError("Video has duplicate presentation timestamps.") + if len(selected) >= 65536: + raise ValueError("Remuxed video exceeds the frame limit.") + seen_pts.add(timestamp) + selected.append(timestamp) + shift = origin / (packet.time_base or stream.time_base) + if shift.denominator != 1: + raise ValueError("Cannot normalize packet time base.") + packet.pts -= int(shift) + packet.dts -= int(shift) + packet.stream = target + muxer.mux(packet) + payload = output.getvalue() + # Leading B-frames of an open GOP may precede its keyframe in + # presentation order. They are outside the requested interval and the + # normalized MP4 discards them; do not assign output ordinals to them. + expected = sorted(timestamp for timestamp in selected if timestamp >= origin) + actual = [] + with av.open(io.BytesIO(payload)) as container: + for frame in container.decode(video=0): + if frame.pts is None: + raise ValueError("Remuxed frame has no presentation timestamp.") + actual.append(frame.pts * frame.time_base + origin) + if len(actual) > len(expected): + raise ValueError("Remux changed the number of frames.") + if actual != expected: + raise ValueError("Remux changed presentation timestamps or frame count.") + mapping = [(timestamp, ordinal) for ordinal, timestamp in enumerate(actual)] + if not mapping: + raise ValueError("Requested interval contains no video frames.") + return payload, mapping + + +class _EpisodeVideoSources(dict): + """URI reader adapter retaining only the active Episode payloads.""" + + def __init__(self, delegate): + super().__init__() + self.delegate = delegate + self.payloads = {} + self.clips = {} + self.sequence = 0 + + def create(self, uri): + return self if uri.startswith("paimon-remux:") else self.delegate.create(uri) + + def new_input_stream(self, uri): + return io.BytesIO(self.payloads[uri]) + + def next_episode(self): + self.clips.clear() + self.payloads.clear() + + def descriptors(self, uri, start, end, timestamps): + from pypaimon.table.row.blob import VideoFrameDescriptor + from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex + + key = (uri, start, end) + clip = self.clips.get(key) + if clip is None: + with self.delegate.create(uri).new_input_stream(uri) as source: + payload, mapping = remux_video_interval(source, start, end) + index = VideoKeyframeIndex.inspect(io.BytesIO(payload), len(payload)).serialize() + name = "paimon-remux:episode-%d" % self.sequence + self.sequence += 1 + self.payloads[name] = payload + index + clip = (name, len(payload), len(index), + [float(t) for t, unused in mapping], + [ordinal for unused, ordinal in mapping]) + self.clips[key] = clip + name, length, index_length, times, ordinals = clip + result = [] + for timestamp in timestamps: + # Source timestamps have already been validated by the importer. + position = bisect_left(times, timestamp) + candidates = [i for i in (position - 1, position) + if 0 <= i < len(times) and abs(times[i] - timestamp) < 1e-7] + if len(candidates) != 1: + raise ValueError("Remux cannot map source timestamp %r in [%r, %r); clip range %r." + % (timestamp, start, end, (times[:1], times[-1:]))) + result.append(VideoFrameDescriptor( + name, 0, length, ordinals[candidates[0]], length, index_length).serialize()) + return result diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 56966d550348..799bda3b029e 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -3070,7 +3070,27 @@ def test_imported_video_payload_can_be_decoded(self): def test_imported_video_payload_supports_training_reads(self): self._assert_imported_video_payload_can_be_decoded(True) - def _assert_imported_video_payload_can_be_decoded(self, training_reads): + @unittest.skipUnless(av is not None, "PyAV is required for MP4 decoding") + def test_remuxed_episode_payload_can_be_decoded(self): + self._assert_imported_video_payload_can_be_decoded(False, remux=True) + + @unittest.skipUnless( + av is not None + and sys.version_info >= (3, 10) + and importlib.util.find_spec("datasets") is not None + and importlib.util.find_spec("torch") is not None, + "Video training reads require Python 3.10+, PyAV, datasets, and Torch", + ) + def test_remuxed_episode_payload_supports_training_reads(self): + self._assert_imported_video_payload_can_be_decoded(True, remux=True) + + @unittest.skipUnless(av is not None, "PyAV is required for MP4 decoding") + def test_remux_local_source_does_not_use_destination_file_io(self): + self._assert_imported_video_payload_can_be_decoded( + False, remux=True, remote_destination=True) + + def _assert_imported_video_payload_can_be_decoded( + self, training_reads, remux=False, remote_destination=False): import pandas as pd temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_mp4_")) @@ -3234,6 +3254,21 @@ def read_batch(self, begin, end): "warehouse": str(temp_dir / "warehouse"), "file-io.read-coalesce.max-gap": "0 b", }) + from pypaimon.multimodal.lerobot.loader import _prepare_dataset + + def prepare(table, *args, **kwargs): + if not remote_destination: + return _prepare_dataset(table, *args, **kwargs) + open_destination = table.file_io.new_input_stream + + def open_file(path): + if str(path).endswith(".mp4"): + raise IOError("Destination FileIO cannot open local source videos") + return open_destination(path) + + with patch.object(table.file_io, "new_input_stream", side_effect=open_file): + return _prepare_dataset(table, *args, **kwargs) + with patch( "pypaimon.multimodal.lerobot.api." "_import_lerobot_dataset", @@ -3242,9 +3277,13 @@ def read_batch(self, begin, end): "pypaimon.multimodal.lerobot.api." "_open_resolved_dataset", return_value=Dataset(), + ), patch( + "pypaimon.multimodal.lerobot.loader._prepare_dataset", + side_effect=prepare, ): connection.load_from_lerobot( - "frames", temp_dir, batch_size=1) + "frames", temp_dir, batch_size=1, + options={"write.lerobot.remux-episodes": str(remux).lower()}) table = connection.get_table("frames") rows = table.scan().select([ @@ -3262,10 +3301,14 @@ def read_batch(self, begin, end): descriptor.keyframe_index_descriptor is not None for descriptor in descriptors + camera_b_descriptors )) - self.assertEqual( - [55, 56, 61, 62, 63], - [descriptor.frame_index for descriptor in descriptors], - ) + if not remux: + self.assertEqual( + [55, 56, 61, 62, 63], + [descriptor.frame_index for descriptor in descriptors], + ) + else: + self.assertEqual(2, len(set( + d.payload_descriptor for d in descriptors))) class Decoder: @@ -3292,7 +3335,7 @@ def close(self): ) try: decoded_rows = collator(rows) - self.assertEqual(1, len(collator._decoders)) + self.assertEqual(2 if remux else 1, len(collator._decoders)) finally: collator.close() np.testing.assert_allclose( @@ -3300,11 +3343,12 @@ def close(self): expected_frame_values, atol=1e-6, ) - np.testing.assert_allclose( - [row["decoded"][1] for row in decoded_rows], - [5.5, 5.6, 6.1, 6.2, 6.3], - atol=1e-6, - ) + if not remux: + np.testing.assert_allclose( + [row["decoded"][1] for row in decoded_rows], + [5.5, 5.6, 6.1, 6.2, 6.3], + atol=1e-6, + ) if not training_reads: return @@ -3373,10 +3417,12 @@ def track_file_range(path, offset, length): self.assertEqual( 1, len(dataset._video_collators[0]._decoders)) self.assertGreater(sum(transferred), 0) - self.assertLess( - sum(transferred), - sum(descriptor.length for descriptor in selected), - ) + # A short remuxed clip may be read in full. Compare with the + # original video rather than requiring sparse IO within it. + read_limit = (video_path.stat().st_size * len(selected) + if remux else + sum(d.length for d in selected)) + self.assertLess(sum(transferred), read_limit) from torch.utils.data import DataLoader worker_indices = [] diff --git a/paimon-python/pypaimon/tests/video_remux_test.py b/paimon-python/pypaimon/tests/video_remux_test.py new file mode 100644 index 000000000000..b9b8d3ee8c30 --- /dev/null +++ b/paimon-python/pypaimon/tests/video_remux_test.py @@ -0,0 +1,145 @@ +# 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. + +import io +from fractions import Fraction + +import numpy as np +import pytest + +from pypaimon.multimodal.video_remux import remux_video_interval + +av = pytest.importorskip("av") + + +@pytest.fixture(params=[0, 2]) +def video(tmp_path, request): + path = tmp_path / "input.mp4" + with av.open(str(path), "w") as container: + stream = container.add_stream("libx264", rate=10) + stream.width = 64 + stream.height = 48 + stream.pix_fmt = "yuv420p" + stream.codec_context.gop_size = 10 + stream.codec_context.max_b_frames = request.param + for i in range(50): + pixels = np.full((48, 64, 3), (i * 5) % 255, dtype=np.uint8) + frame = av.VideoFrame.from_ndarray(pixels, format="rgb24") + frame.pts = i + frame.time_base = Fraction(1, 10) + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + return path + + +@pytest.mark.parametrize("start,end", [(0, 0.4), (0.7, 2.3), (4.4, 5)]) +def test_remux_preserves_requested_frames(video, start, end): + with av.open(str(video)) as source: + expected = {f.pts * f.time_base: f.to_ndarray(format="rgb24") + for f in source.decode(video=0) + if start <= float(f.pts * f.time_base) < end} + payload, mapping = remux_video_interval(str(video), start, end) + with av.open(io.BytesIO(payload)) as source: + frames = list(source.decode(video=0)) + mapping = [(t, ordinal) for t, ordinal in mapping if start <= float(t) < end] + assert set(t for t, unused in mapping) == set(expected) + for timestamp, ordinal in mapping: + np.testing.assert_array_equal(frames[ordinal].to_ndarray(format="rgb24"), + expected[timestamp]) + + +def test_remux_bounds_output(video): + with pytest.raises(ValueError, match="memory limit"): + remux_video_interval(str(video), 0, 5, max_bytes=32) + + +def test_remux_rejects_empty_interval(video): + with pytest.raises(ValueError, match="Invalid video interval"): + remux_video_interval(str(video), 2, 1) + + +def test_remux_accepts_episode_boundary_roundoff(video): + payload, mapping = remux_video_interval(str(video), 0.70000000000002, 2.30000000000002) + assert set(Fraction(i, 10) for i in range(7, 23)).issubset(t for t, unused in mapping) + + +@pytest.mark.parametrize("version", ["12.0.0", "13.1.0", "14.1.0"]) +def test_remux_requires_supported_av_before_open(monkeypatch, version): + from unittest.mock import Mock + from pypaimon.multimodal.lerobot.loader import _prepare_dataset + + monkeypatch.setattr(av, "__version__", version) + opened = Mock(side_effect=AssertionError("Source must not be opened")) + monkeypatch.setattr(av, "open", opened) + with pytest.raises(ImportError, match="PyAV 14.2"): + remux_video_interval("unused.mp4", 0, 1) + table = Mock() + table.options.options.get.return_value = True + with pytest.raises(ImportError, match="PyAV 14.2"): + _prepare_dataset(table, Mock(), {}, None, None, 1, {}, video_fields=("camera",)) + table.new_batch_write_builder.assert_not_called() + opened.assert_not_called() + + +@pytest.mark.parametrize("b_frames", [0, 2]) +def test_remux_preserves_float32_episode_bounds(tmp_path, b_frames): + import pyarrow as pa + from pypaimon.multimodal.lerobot.loader import ( + _video_frame_ordinal, _validate_video_timestamp_range) + from pypaimon.multimodal.video_remux import _EpisodeVideoSources + from pypaimon.table.row.blob import VideoFrameDescriptor + from pypaimon.common.uri_reader import UriReaderFactory + + path = tmp_path / "float32.mp4" + with av.open(str(path), "w") as container: + stream = container.add_stream("libx264", rate=30) + stream.width, stream.height = 64, 48 + stream.pix_fmt = "yuv420p" + stream.codec_context.gop_size = 12 + stream.codec_context.max_b_frames = b_frames + for i in range(156): + frame = av.VideoFrame.from_ndarray( + np.full((48, 64, 3), i, dtype=np.uint8), format="rgb24") + frame.pts, frame.time_base = i, Fraction(1, 30) + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + with av.open(str(path)) as container: + original = list(container.decode(video=0)) + times = [float(frame.pts * frame.time_base) for frame in original] + start, end = float(np.float32(4.8)), float(np.float32(4.9)) + _validate_video_timestamp_range(start, end, 30, 3, "camera") + ordinals = [_video_frame_ordinal(times, float(np.float32(i / 30)), start, + pa.float32(), "camera") for i in range(3)] + assert ordinals == [144, 145, 146] + factory = UriReaderFactory({}) + sources = _EpisodeVideoSources(factory) + try: + # Split calls also exercise reuse of the Episode mapping across batches. + rows = [sources.descriptors(path.as_uri(), start, end, [times[i]])[0] for i in ordinals] + descriptors = [VideoFrameDescriptor.deserialize(row) for row in rows] + payload = sources.payloads[descriptors[0].uri][:descriptors[0].length] + with av.open(io.BytesIO(payload)) as container: + decoded = list(container.decode(video=0)) + for ordinal, descriptor in zip(ordinals, descriptors): + np.testing.assert_array_equal( + decoded[descriptor.frame_index].to_ndarray(format="rgb24"), + original[ordinal].to_ndarray(format="rgb24")) + finally: + factory.close()