Add tests

Author: NicolasHug

src/torchcodec/decoders/_audio_decoder.py

@@ -25,10 +25,13 @@ def __init__(
         source: Union[str, Path, bytes, Tensor],
         *,
         stream_index: Optional[int] = None,
+        sample_rate: Optional[int] = None,
     ):
         self._decoder = create_decoder(source=source, seek_mode="approximate")
 
-        core.add_audio_stream(self._decoder, stream_index=stream_index)
+        core.add_audio_stream(
+            self._decoder, stream_index=stream_index, sample_rate=sample_rate
+        )
 
         (
             self.metadata,
@@ -39,6 +42,10 @@ def __init__(
             decoder=self._decoder, stream_index=stream_index, media_type="audio"
         )
         assert isinstance(self.metadata, core.AudioStreamMetadata)  # mypy
+        self._source_sample_rate = self.metadata.sample_rate
+        self._desired_sample_rate = (
+            sample_rate if sample_rate is not None else self._source_sample_rate
+        )
 
     def get_samples_played_in_range(
         self, start_seconds: float, stop_seconds: Optional[float] = None
@@ -75,11 +82,7 @@ def get_samples_played_in_range(
         # So we do some basic math to figure out the position of the view that
         # we'll return.
 
-        # TODO: sample_rate is either the original one from metadata, or the
-        # user-specified one (NIY)
-        assert isinstance(self.metadata, core.AudioStreamMetadata)  # mypy
-        sample_rate = self.metadata.sample_rate
-
+        sample_rate = self._desired_sample_rate
         # TODO: metadata's sample_rate should probably not be Optional
         assert sample_rate is not None  # mypy.
 
@@ -94,7 +97,7 @@ def get_samples_played_in_range(
             output_pts_seconds = first_pts
 
         num_samples = frames.shape[1]
-        last_pts = first_pts + num_samples / self.metadata.sample_rate
+        last_pts = first_pts + num_samples / sample_rate
         if stop_seconds is not None and stop_seconds < last_pts:
             offset_end = num_samples - round((last_pts - stop_seconds) * sample_rate)
         else:

test/decoders/test_decoders.py

@@ -26,6 +26,8 @@
     NASA_AUDIO_MP3,
     NASA_VIDEO,
     SINE_MONO_S32,
+    SINE_MONO_S32_44100,
+    SINE_MONO_S32_8000,
 )
 
 
@@ -1088,3 +1090,65 @@ def test_format_conversion(self):
 
         reference_frames = asset.get_frame_data_by_range(start=0, stop=asset.num_frames)
         torch.testing.assert_close(all_samples.data, reference_frames)
+
+    @pytest.mark.parametrize(
+        "start_seconds, stop_seconds",
+        (
+            (0, None),
+            (0, 4),
+            (0, 3),
+            (2, None),
+            (2, 3),
+        ),
+    )
+    def test_sample_rate_conversion(self, start_seconds, stop_seconds):
+        # When start_seconds is not exactly 0, we have to increase the tolerance
+        # a bit. This is because sample_rate conversion relies on a sliding
+        # window of samples: if we start a stream in the middle, the first few
+        # samples aren't able to take advantage of the preceeding samples.
+        atol = 1e-4 if start_seconds == 0 else 1e-2
+        rtol = 1e-6
+
+        # Upsample
+        decoder = AudioDecoder(SINE_MONO_S32_44100.path)
+        assert decoder.metadata.sample_rate == 44_100
+        frames_44100_native = decoder.get_samples_played_in_range(
+            start_seconds=start_seconds, stop_seconds=stop_seconds
+        )
+        assert frames_44100_native.sample_rate == 44_100
+
+        decoder = AudioDecoder(SINE_MONO_S32.path, sample_rate=44_100)
+        frames_upsampled_to_44100 = decoder.get_samples_played_in_range(
+            start_seconds=start_seconds, stop_seconds=stop_seconds
+        )
+        assert decoder.metadata.sample_rate == 16_000
+        assert frames_upsampled_to_44100.sample_rate == 44_100
+
+        torch.testing.assert_close(
+            frames_upsampled_to_44100.data,
+            frames_44100_native.data,
+            atol=atol,
+            rtol=rtol,
+        )
+
+        # Downsample
+        decoder = AudioDecoder(SINE_MONO_S32_8000.path)
+        assert decoder.metadata.sample_rate == 8000
+        frames_8000_native = decoder.get_samples_played_in_range(
+            start_seconds=start_seconds, stop_seconds=stop_seconds
+        )
+        assert frames_8000_native.sample_rate == 8000
+
+        decoder = AudioDecoder(SINE_MONO_S32.path, sample_rate=8000)
+        frames_downsampled_to_8000 = decoder.get_samples_played_in_range(
+            start_seconds=start_seconds, stop_seconds=stop_seconds
+        )
+        assert decoder.metadata.sample_rate == 16_000
+        assert frames_downsampled_to_8000.sample_rate == 8000
+
+        torch.testing.assert_close(
+            frames_downsampled_to_8000.data,
+            frames_8000_native.data,
+            atol=atol,
+            rtol=rtol,
+        )

test/decoders/test_ops.py

@@ -884,11 +884,11 @@ def test_decode_before_frame_start(self):
         torch.testing.assert_close(frames, all_frames)
 
     def test_sample_rate_conversion(self):
-        def get_all_frames(asset, sample_rate=None):
+        def get_all_frames(asset, sample_rate=None, stop_seconds=None):
             decoder = create_from_file(str(asset.path), seek_mode="approximate")
             add_audio_stream(decoder, sample_rate=sample_rate)
             frames, *_ = get_frames_by_pts_in_range_audio(
-                decoder, start_seconds=0, stop_seconds=None
+                decoder, start_seconds=0, stop_seconds=stop_seconds
             )
             return frames