Encoder: properly set frame pts values (#726)

Author: NicolasHug

src/torchcodec/_core/Encoder.cpp

@@ -293,8 +293,7 @@ void AudioEncoder::encode() {
     encodeInnerLoop(autoAVPacket, convertedAVFrame);
 
     numEncodedSamples += numSamplesToEncode;
-    // TODO-ENCODING set frame pts correctly, and test against it.
-    // avFrame->pts += static_cast<int64_t>(numSamplesToEncode);
+    avFrame->pts += static_cast<int64_t>(numSamplesToEncode);
   }
   TORCH_CHECK(numEncodedSamples == numSamples, "Hmmmmmm something went wrong.");
 

src/torchcodec/_core/FFMPEGCommon.cpp

@@ -274,6 +274,7 @@ UniqueAVFrame convertAudioAVFrameSamples(
       convertedAVFrame,
       "Could not allocate frame for sample format conversion.");
 
+  convertedAVFrame->pts = srcAVFrame->pts;
   convertedAVFrame->format = static_cast<int>(outSampleFormat);
 
   convertedAVFrame->sample_rate = outSampleRate;

src/torchcodec/_frame.py

@@ -125,7 +125,7 @@ class AudioSamples(Iterable):
     pts_seconds: float
     """The :term:`pts` of the first sample, in seconds."""
     duration_seconds: float
-    """The duration of the sampleas, in seconds."""
+    """The duration of the samples, in seconds."""
     sample_rate: int
     """The sample rate of the samples, in Hz."""
 

test/test_encoders.py

@@ -1,5 +1,8 @@
+import json
+import os
 import re
 import subprocess
+from pathlib import Path
 
 import pytest
 import torch
@@ -16,12 +19,89 @@
 )
 
 
+@pytest.fixture
+def with_ffmpeg_debug_logs():
+    # Fixture that sets the ffmpeg logs to DEBUG mode
+    previous_log_level = os.environ.get("TORCHCODEC_FFMPEG_LOG_LEVEL", "QUIET")
+    os.environ["TORCHCODEC_FFMPEG_LOG_LEVEL"] = "DEBUG"
+    yield
+    os.environ["TORCHCODEC_FFMPEG_LOG_LEVEL"] = previous_log_level
+
+
+def validate_frames_properties(*, actual: Path, expected: Path):
+    # actual and expected are files containing encoded audio data.  We call
+    # `ffprobe` on both, and assert that the frame properties match (pts,
+    # duration, etc.)
+
+    frames_actual, frames_expected = (
+        json.loads(
+            subprocess.run(
+                [
+                    "ffprobe",
+                    "-v",
+                    "error",
+                    "-hide_banner",
+                    "-select_streams",
+                    "a:0",
+                    "-show_frames",
+                    "-of",
+                    "json",
+                    f"{f}",
+                ],
+                check=True,
+                capture_output=True,
+                text=True,
+            ).stdout
+        )["frames"]
+        for f in (actual, expected)
+    )
+
+    # frames_actual and frames_expected are both a list of dicts, each dict
+    # corresponds to a frame and each key-value pair corresponds to a frame
+    # property like pts, nb_samples, etc., similar to the AVFrame fields.
+    assert isinstance(frames_actual, list)
+    assert all(isinstance(d, dict) for d in frames_actual)
+
+    assert len(frames_actual) > 3  # arbitrary sanity check
+    assert len(frames_actual) == len(frames_expected)
+
+    # non-exhaustive list of the props we want to test for:
+    required_props = (
+        "pts",
+        "pts_time",
+        "sample_fmt",
+        "nb_samples",
+        "channels",
+        "duration",
+        "duration_time",
+    )
+
+    for frame_index, (d_actual, d_expected) in enumerate(
+        zip(frames_actual, frames_expected)
+    ):
+        if get_ffmpeg_major_version() >= 6:
+            assert all(required_prop in d_expected for required_prop in required_props)
+
+        for prop in d_expected:
+            if prop == "pkt_pos":
+                # pkt_pos is the position of the packet *in bytes* in its
+                # stream. We don't always match FFmpeg exactly on this,
+                # typically on compressed formats like mp3. It's probably
+                # because we are not writing the exact same headers, or
+                # something like this. In any case, this doesn't seem to be
+                # critical.
+                continue
+            assert (
+                d_actual[prop] == d_expected[prop]
+            ), f"\nComparing: {actual}\nagainst reference: {expected},\nthe {prop} property is different at frame {frame_index}:"
+
+
 class TestAudioEncoder:
 
     def decode(self, source) -> torch.Tensor:
         if isinstance(source, TestContainerFile):
             source = str(source.path)
-        return AudioDecoder(source).get_all_samples().data
+        return AudioDecoder(source).get_all_samples()
 
     def test_bad_input(self):
         with pytest.raises(ValueError, match="Expected samples to be a Tensor"):
@@ -63,12 +143,12 @@ def test_bad_input_parametrized(self, method, tmp_path):
             else dict(format="mp3")
         )
 
-        decoder = AudioEncoder(self.decode(NASA_AUDIO_MP3), sample_rate=10)
+        decoder = AudioEncoder(self.decode(NASA_AUDIO_MP3).data, sample_rate=10)
         with pytest.raises(RuntimeError, match="invalid sample rate=10"):
             getattr(decoder, method)(**valid_params)
 
         decoder = AudioEncoder(
-            self.decode(NASA_AUDIO_MP3), sample_rate=NASA_AUDIO_MP3.sample_rate
+            self.decode(NASA_AUDIO_MP3).data, sample_rate=NASA_AUDIO_MP3.sample_rate
         )
         with pytest.raises(RuntimeError, match="bit_rate=-1 must be >= 0"):
             getattr(decoder, method)(**valid_params, bit_rate=-1)
@@ -81,7 +161,7 @@ def test_bad_input_parametrized(self, method, tmp_path):
             getattr(decoder, method)(**valid_params)
 
         decoder = AudioEncoder(
-            self.decode(NASA_AUDIO_MP3), sample_rate=NASA_AUDIO_MP3.sample_rate
+            self.decode(NASA_AUDIO_MP3).data, sample_rate=NASA_AUDIO_MP3.sample_rate
         )
         for num_channels in (0, 3):
             with pytest.raises(
@@ -101,7 +181,7 @@ def test_round_trip(self, method, format, tmp_path):
             pytest.skip("Swresample with FFmpeg 4 doesn't work on wav files")
 
         asset = NASA_AUDIO_MP3
-        source_samples = self.decode(asset)
+        source_samples = self.decode(asset).data
 
         encoder = AudioEncoder(source_samples, sample_rate=asset.sample_rate)
 
@@ -116,7 +196,7 @@ def test_round_trip(self, method, format, tmp_path):
 
         rtol, atol = (0, 1e-4) if format == "wav" else (None, None)
         torch.testing.assert_close(
-            self.decode(encoded_source), source_samples, rtol=rtol, atol=atol
+            self.decode(encoded_source).data, source_samples, rtol=rtol, atol=atol
         )
 
     @pytest.mark.skipif(in_fbcode(), reason="TODO: enable ffmpeg CLI")
@@ -125,7 +205,17 @@ def test_round_trip(self, method, format, tmp_path):
     @pytest.mark.parametrize("num_channels", (None, 1, 2))
     @pytest.mark.parametrize("format", ("mp3", "wav", "flac"))
     @pytest.mark.parametrize("method", ("to_file", "to_tensor"))
-    def test_against_cli(self, asset, bit_rate, num_channels, format, method, tmp_path):
+    def test_against_cli(
+        self,
+        asset,
+        bit_rate,
+        num_channels,
+        format,
+        method,
+        tmp_path,
+        capfd,
+        with_ffmpeg_debug_logs,
+    ):
         # Encodes samples with our encoder and with the FFmpeg CLI, and checks
         # that both decoded outputs are equal
 
@@ -144,14 +234,25 @@ def test_against_cli(self, asset, bit_rate, num_channels, format, method, tmp_pa
             check=True,
         )
 
-        encoder = AudioEncoder(self.decode(asset), sample_rate=asset.sample_rate)
+        encoder = AudioEncoder(self.decode(asset).data, sample_rate=asset.sample_rate)
+
         params = dict(bit_rate=bit_rate, num_channels=num_channels)
         if method == "to_file":
             encoded_by_us = tmp_path / f"output.{format}"
             encoder.to_file(dest=str(encoded_by_us), **params)
         else:
             encoded_by_us = encoder.to_tensor(format=format, **params)
 
+        captured = capfd.readouterr()
+        if format == "wav":
+            assert "Timestamps are unset in a packet" not in captured.err
+        if format == "mp3":
+            assert "Queue input is backward in time" not in captured.err
+        if format in ("flac", "wav"):
+            assert "Encoder did not produce proper pts" not in captured.err
+        if format in ("flac", "mp3"):
+            assert "Application provided invalid" not in captured.err
+
         if format == "wav":
             rtol, atol = 0, 1e-4
         elif format == "mp3" and asset is SINE_MONO_S32 and num_channels == 2:
@@ -162,12 +263,22 @@ def test_against_cli(self, asset, bit_rate, num_channels, format, method, tmp_pa
             rtol, atol = 0, 1e-3
         else:
             rtol, atol = None, None
+        samples_by_us = self.decode(encoded_by_us)
+        samples_by_ffmpeg = self.decode(encoded_by_ffmpeg)
         torch.testing.assert_close(
-            self.decode(encoded_by_ffmpeg),
-            self.decode(encoded_by_us),
+            samples_by_us.data,
+            samples_by_ffmpeg.data,
             rtol=rtol,
             atol=atol,
         )
+        assert samples_by_us.pts_seconds == samples_by_ffmpeg.pts_seconds
+        assert samples_by_us.duration_seconds == samples_by_ffmpeg.duration_seconds
+        assert samples_by_us.sample_rate == samples_by_ffmpeg.sample_rate
+
+        if method == "to_file":
+            validate_frames_properties(actual=encoded_by_us, expected=encoded_by_ffmpeg)
+        else:
+            assert method == "to_tensor", "wrong test parametrization!"
 
     @pytest.mark.parametrize("asset", (NASA_AUDIO_MP3, SINE_MONO_S32))
     @pytest.mark.parametrize("bit_rate", (None, 0, 44_100, 999_999_999))
@@ -179,7 +290,7 @@ def test_to_tensor_against_to_file(
         if get_ffmpeg_major_version() == 4 and format == "wav":
             pytest.skip("Swresample with FFmpeg 4 doesn't work on wav files")
 
-        encoder = AudioEncoder(self.decode(asset), sample_rate=asset.sample_rate)
+        encoder = AudioEncoder(self.decode(asset).data, sample_rate=asset.sample_rate)
 
         params = dict(bit_rate=bit_rate, num_channels=num_channels)
         encoded_file = tmp_path / f"output.{format}"
@@ -189,7 +300,7 @@ def test_to_tensor_against_to_file(
         )
 
         torch.testing.assert_close(
-            self.decode(encoded_file), self.decode(encoded_tensor)
+            self.decode(encoded_file).data, self.decode(encoded_tensor).data
         )
 
     def test_encode_to_tensor_long_output(self):
@@ -205,7 +316,7 @@ def test_encode_to_tensor_long_output(self):
         INITIAL_TENSOR_SIZE = 10_000_000
         assert encoded_tensor.numel() > INITIAL_TENSOR_SIZE
 
-        torch.testing.assert_close(self.decode(encoded_tensor), samples)
+        torch.testing.assert_close(self.decode(encoded_tensor).data, samples)
 
     def test_contiguity(self):
         # Ensure that 2 waveforms with the same values are encoded in the same
@@ -262,4 +373,4 @@ def test_num_channels(
 
         if num_channels_output is None:
             num_channels_output = num_channels_input
-        assert self.decode(encoded_source).shape[0] == num_channels_output
+        assert self.decode(encoded_source).data.shape[0] == num_channels_output