[codex] add video frame count method

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

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)

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
 

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
 

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))