Calculate num_frames, add mock tests

Author: Dan-Flores

src/torchcodec/_core/_metadata.py

@@ -129,8 +129,15 @@ def num_frames(self) -> Optional[int]:
         """
         if self.num_frames_from_content is not None:
             return self.num_frames_from_content
-        else:
+        elif self.num_frames_from_header is not None:
             return self.num_frames_from_header
+        elif (
+            self.average_fps_from_header is not None
+            and self.duration_seconds_from_header is not None
+        ):
+            return int(self.average_fps_from_header * self.duration_seconds_from_header)
+        else:
+            return None
 
     @property
     def average_fps(self) -> Optional[float]:

test/test_decoders.py

@@ -6,6 +6,8 @@
 
 import contextlib
 import gc
+import json
+from unittest.mock import patch
 
 import numpy
 import pytest
@@ -738,6 +740,76 @@ def test_get_frames_in_range(self, stream_index, device, seek_mode):
             empty_frames.duration_seconds, NASA_VIDEO.empty_duration_seconds
         )
 
+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
+    @patch("torchcodec._core._metadata._get_stream_json_metadata")
+    def test_get_frames_with_missing_num_frames_metadata(
+        self, mock_get_stream_json_metadata, device, seek_mode
+    ):
+        # Create a mock stream_dict to test that initializing VideoDecoder without
+        # num_frames_from_header and num_frames_from_content calculates num_frames
+        # using the average_fps and duration_seconds metadata.
+        mock_stream_dict = {
+            "averageFpsFromHeader": 29.97003,
+            "beginStreamSecondsFromContent": 0.0,
+            "beginStreamSecondsFromHeader": 0.0,
+            "bitRate": 128783.0,
+            "codec": "h264",
+            "durationSecondsFromHeader": 13.013,
+            "endStreamSecondsFromContent": 13.013,
+            "width": 480,
+            "height": 270,
+            "mediaType": "video",
+            "numFramesFromHeader": None,
+            "numFramesFromContent": None,
+        }
+        # Set the return value of the mock to be the mock_stream_dict
+        mock_get_stream_json_metadata.return_value = json.dumps(mock_stream_dict)
+
+        decoder = VideoDecoder(
+            NASA_VIDEO.path,
+            stream_index=3,
+            device=device,
+            seek_mode=seek_mode,
+        )
+
+        assert decoder.metadata.num_frames_from_header is None
+        assert decoder.metadata.num_frames_from_content is None
+        assert decoder.metadata.duration_seconds is not None
+        assert decoder.metadata.average_fps is not None
+        assert decoder.metadata.num_frames == int(
+            decoder.metadata.duration_seconds * decoder.metadata.average_fps
+        )
+
+        # Test get_frames_in_range
+        ref_frames9 = NASA_VIDEO.get_frame_data_by_range(
+            start=9, stop=10, stream_index=3
+        ).to(device)
+        frames9 = decoder.get_frames_in_range(start=9, stop=10)
+        assert_frames_equal(ref_frames9, frames9.data)
+
+        # Test get_frame_at
+        ref_frame9 = NASA_VIDEO.get_frame_data_by_index(9, stream_index=3).to(device)
+        frame9 = decoder.get_frame_at(9)
+        torch.testing.assert_close(ref_frame9, frame9.data)
+
+        # Test get_frames_at
+        indices = [0, 1, 25, 35]
+        ref_frames = [
+            NASA_VIDEO.get_frame_data_by_index(i, stream_index=3).to(device)
+            for i in indices
+        ]
+        frames = decoder.get_frames_at(indices)
+        for ref, frame in zip(ref_frames, frames.data):
+            torch.testing.assert_close(ref, frame)
+
+        # Test get_frames_played_in_range to get all frames
+        assert decoder.metadata.end_stream_seconds is not None
+        all_frames = decoder.get_frames_played_in_range(
+            decoder.metadata.begin_stream_seconds, decoder.metadata.end_stream_seconds
+        )
+        assert_frames_equal(all_frames.data, decoder[:])
+
     @pytest.mark.parametrize("dimension_order", ["NCHW", "NHWC"])
     @pytest.mark.parametrize(
         "frame_getter",

test/test_metadata.py

@@ -119,7 +119,7 @@ def test_get_metadata_audio_file(metadata_getter):
 
 @pytest.mark.parametrize(
     "num_frames_from_header, num_frames_from_content, expected_num_frames",
-    [(None, 10, 10), (10, None, 10), (None, None, None)],
+    [(10, 20, 20), (None, 10, 10), (10, None, 10)],
 )
 def test_num_frames_fallback(
     num_frames_from_header, num_frames_from_content, expected_num_frames
@@ -143,6 +143,34 @@ def test_num_frames_fallback(
     assert metadata.num_frames == expected_num_frames
 
 
+@pytest.mark.parametrize(
+    "average_fps_from_header, duration_seconds_from_header, expected_num_frames",
+    [(60, 10, 600), (60, None, None), (None, 10, None), (None, None, None)],
+)
+def test_calculate_num_frames_using_fps_and_duration(
+    average_fps_from_header, duration_seconds_from_header, expected_num_frames
+):
+    """Check that if num_frames_from_content and num_frames_from_header are missing,
+    `.num_frames` is calculated using average_fps_from_header and duration_seconds_from_header
+    """
+    metadata = VideoStreamMetadata(
+        duration_seconds_from_header=duration_seconds_from_header,
+        bit_rate=123,
+        num_frames_from_header=None,  # None to test calculating num_frames
+        num_frames_from_content=None,  # None to test calculating num_frames
+        begin_stream_seconds_from_header=0,
+        begin_stream_seconds_from_content=0,
+        end_stream_seconds_from_content=4,
+        codec="whatever",
+        width=123,
+        height=321,
+        average_fps_from_header=average_fps_from_header,
+        stream_index=0,
+    )
+
+    assert metadata.num_frames == expected_num_frames
+
+
 def test_repr():
     # Test for calls to print(), str(), etc. Useful to make sure we don't forget
     # to add additional @properties to __repr__