From 6bc12f9e446bbf068c0dee4c240bf7f9e2d7bdd7 Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 3 Jun 2026 23:56:32 +0200 Subject: [PATCH] add video frame count method --- docs/episode-data/frames-and-videos.md | 15 ++++++ .../pipeline/utils/cache/decoder_cache.py | 3 ++ src/refiner/video/types.py | 41 +++++++++++++++ tests/readers/test_zarr_reader.py | 3 ++ tests/test_video_decode.py | 51 +++++++++++++++++++ 5 files changed, 113 insertions(+) diff --git a/docs/episode-data/frames-and-videos.md b/docs/episode-data/frames-and-videos.md index a1d33a0c..ac9f71a0 100644 --- a/docs/episode-data/frames-and-videos.md +++ b/docs/episode-data/frames-and-videos.md @@ -53,11 +53,26 @@ sequences. They share the same core operations: | Operation | Purpose | | --- | --- | +| `get_frame_count()` | Return the exact frame count when it is known or reported by the encoded container. | | `clipped(...)` | Create a time-bounded view. | | `iter_frames()` | Decode video frames lazily. | | `iter_numpy_frames()` | Decode RGB arrays lazily. | | `write_to(...)` | Let a writer copy/remux/transcode media. | +## Video Frame Counts + +```python +async def add_frame_count(row): + video = row.videos["observation.images.top"] + return row.update(video_frame_count=await video.get_frame_count()) +``` + +`get_frame_count()` is explicit because path-backed and bytes-backed videos may +need to open the encoded container. In-memory frame arrays return their exact +length. Frame sequences return their provided count, or count the repeatable +sequence when no count was provided. Encoded videos use container metadata and +raise an error when the container does not report an exact frame count. + ## Clip A Video View ```python diff --git a/src/refiner/pipeline/utils/cache/decoder_cache.py b/src/refiner/pipeline/utils/cache/decoder_cache.py index 4b901e95..b497b1d8 100644 --- a/src/refiner/pipeline/utils/cache/decoder_cache.py +++ b/src/refiner/pipeline/utils/cache/decoder_cache.py @@ -18,6 +18,7 @@ class VideoSourceProbe: width: int height: int fps: float | None + frame_count: int | None time_base: Fraction codec: str | None pix_fmt: str | None @@ -85,11 +86,13 @@ def _probe_video_source( return None stream_fps = stream.average_rate or stream.base_rate + frame_count = int(stream.frames) if stream.frames else None codec_obj = getattr(getattr(stream, "codec_context", None), "codec", None) return VideoSourceProbe( width=int(stream.width), height=int(stream.height), fps=float(stream_fps) if stream_fps is not None else None, + frame_count=frame_count, time_base=Fraction(cast(Any, stream.time_base)), codec=str( getattr(codec_obj, "canonical_name", None) diff --git a/src/refiner/video/types.py b/src/refiner/video/types.py index 382cdce7..f7db7894 100644 --- a/src/refiner/video/types.py +++ b/src/refiner/video/types.py @@ -26,6 +26,8 @@ @runtime_checkable class VideoSource(Protocol): + async def get_frame_count(self) -> int: ... + def clipped( self, *, @@ -94,6 +96,21 @@ def uri(self) -> str: def open(self) -> IO[bytes]: return self.data_file.open(mode="rb") + async def get_frame_count(self) -> int: + from refiner.video.remux import prepare_video_source + + prepared = await prepare_video_source(video=self) + try: + if self.from_timestamp_s is not None or self.to_timestamp_s is not None: + raise ValueError( + "encoded video frame count is unavailable for clipped videos" + ) + if prepared.probe is None or prepared.probe.frame_count is None: + raise ValueError(f"Video frame count is unavailable for {self.uri!r}") + return prepared.probe.frame_count + finally: + prepared.close() + def clipped( self, *, @@ -161,6 +178,22 @@ class VideoBytes: def open(self) -> IO[bytes]: return io.BytesIO(self.data) + async def get_frame_count(self) -> int: + from refiner.video.remux import prepare_video_source + + prepared = await prepare_video_source(video=self) + try: + if self.from_timestamp_s is not None or self.to_timestamp_s is not None: + raise ValueError( + "encoded video frame count is unavailable for clipped videos" + ) + if prepared.probe is None or prepared.probe.frame_count is None: + source = self.uri or type(self).__name__ + raise ValueError(f"Video frame count is unavailable for {source!r}") + return prepared.probe.frame_count + finally: + prepared.close() + def clipped( self, *, @@ -250,6 +283,11 @@ def duration_s(self) -> float | None: return None return self.frame_count / float(self.fps) + async def get_frame_count(self) -> int: + if self.frame_count is not None: + return self.frame_count + return sum(1 for _ in self.iter_frame_arrays()) + def iter_frame_arrays(self) -> Iterator[np.ndarray]: start_idx = ( 0 @@ -409,6 +447,9 @@ def frame_count(self) -> int: def duration_s(self) -> float: return self.frame_count / float(self.fps) + async def get_frame_count(self) -> int: + return self.frame_count + def iter_frame_arrays(self) -> Iterator[np.ndarray]: yield from self._array diff --git a/tests/readers/test_zarr_reader.py b/tests/readers/test_zarr_reader.py index ac95365f..1a883e9f 100644 --- a/tests/readers/test_zarr_reader.py +++ b/tests/readers/test_zarr_reader.py @@ -34,6 +34,9 @@ def finalized_workers( class _EmptyVideoSource: + async def get_frame_count(self): + return 0 + def clipped(self, **_kwargs): return self diff --git a/tests/test_video_decode.py b/tests/test_video_decode.py index f010e625..e5af2381 100644 --- a/tests/test_video_decode.py +++ b/tests/test_video_decode.py @@ -68,6 +68,35 @@ def test_iter_frames_respects_clip_bounds(tmp_path) -> None: assert [frame.timestamp_s for frame in frames] == [0.2, 0.4, 0.6] +def test_video_file_get_frame_count_uses_container_metadata(tmp_path) -> None: + path = tmp_path / "video.mp4" + _write_video(path, num_frames=5, fps=5) + video = mdr.video.VideoFile(DataFile.resolve(path)) + + assert asyncio.run(video.get_frame_count()) == 5 + + +def test_video_bytes_get_frame_count_uses_container_metadata(tmp_path) -> None: + path = tmp_path / "video.mp4" + _write_video(path, num_frames=4, fps=5) + video = mdr.video.VideoBytes(path.read_bytes(), uri=str(path)) + + assert asyncio.run(video.get_frame_count()) == 4 + + +def test_clipped_encoded_video_get_frame_count_raises(tmp_path) -> None: + path = tmp_path / "video.mp4" + _write_video(path, num_frames=5, fps=5) + video = mdr.video.VideoFile( + DataFile.resolve(path), + from_timestamp_s=0.0, + to_timestamp_s=0.4, + ) + + with pytest.raises(ValueError, match="clipped videos"): + asyncio.run(video.get_frame_count()) + + def test_video_frame_array_clip_returns_frame_view() -> None: frames = np.stack([np.full((4, 4, 3), value, dtype=np.uint8) for value in range(6)]) video = mdr.video.VideoFrameArray(frames, fps=10) @@ -81,6 +110,13 @@ def test_video_frame_array_clip_returns_frame_view() -> None: assert [int(frame[0, 0, 0]) for frame in clipped_frames] == [2, 3, 4] +def test_video_frame_array_get_frame_count() -> None: + frames = np.stack([np.full((4, 4, 3), value, dtype=np.uint8) for value in range(6)]) + video = mdr.video.VideoFrameArray(frames, fps=10) + + assert asyncio.run(video.get_frame_count()) == 6 + + def test_video_frame_array_iter_frames() -> None: frames = np.stack([np.full((4, 4, 3), value, dtype=np.uint8) for value in range(3)]) video = mdr.video.VideoFrameArray(frames, fps=5) @@ -114,6 +150,21 @@ def frames(): assert calls == 6 +def test_video_frame_sequence_get_frame_count_counts_unknown_sequence() -> None: + calls = 0 + + def frames(): + nonlocal calls + for value in range(3): + calls += 1 + yield np.full((4, 4, 3), value, dtype=np.uint8) + + video = mdr.video.VideoFrameSequence(frames, fps=5) + + assert asyncio.run(video.get_frame_count()) == 3 + assert calls == 3 + + def test_video_frame_sequence_rejects_one_shot_iterators() -> None: frames = (np.full((4, 4, 3), value, dtype=np.uint8) for value in range(3))