Add file-like to Python and tests

Author: NicolasHug

src/torchcodec/encoders/_multi_stream_encoder.py

@@ -1,4 +1,3 @@
-from pathlib import Path
 from typing import Any
 
 from torch import Tensor
@@ -23,8 +22,13 @@ def write(self, samples: Tensor) -> None:
 
 
 class StreamingEncoder:
-    def __init__(self, dest: str | Path):
-        self._encoder_tensor = _core.create_streaming_encoder_to_file(str(dest))
+    def __init__(self, dest, *, format: str | None = None):
+        if format is not None:
+            self._encoder_tensor = _core.create_streaming_encoder_to_file_like(
+                format, dest
+            )
+        else:
+            self._encoder_tensor = _core.create_streaming_encoder_to_file(str(dest))
 
     def add_video(
         self,

test/test_encoders.py

@@ -1566,12 +1566,31 @@ def test_fragmented_mp4(
 
 
 class TestStreamingEncoder:
-    def test_double_close(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    @staticmethod
+    def _create_encoder(method, tmp_path, format):
+        if method == "to_file":
+            encoder_output = tmp_path / f"test.{format}"
+            return StreamingEncoder(encoder_output), encoder_output
+        elif method == "to_file_like":
+            encoder_output = io.BytesIO()
+            return StreamingEncoder(encoder_output, format=format), encoder_output
+        else:
+            raise ValueError(f"Unknown method: {method}")
+
+    @staticmethod
+    def _get_decoder_source(encoder_output):
+        if isinstance(encoder_output, io.BytesIO):
+            return encoder_output.getvalue()
+        return str(encoder_output)
+
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_double_close(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         enc.close()
         enc.close()  # double close is a no-op
 
     @pytest.mark.parametrize("format", ["mp4", "mov", "mkv"])
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
     @pytest.mark.parametrize(
         "device",
         (
@@ -1587,16 +1606,15 @@ def test_double_close(self, tmp_path):
             ),
         ),
     )
-    def test_add_video_and_encode_frames(self, tmp_path, format, device):
+    def test_add_video_and_encode_frames(self, tmp_path, format, method, device):
         source_decoder = VideoDecoder(str(TEST_SRC_2_720P.path))
         source_frames = source_decoder.get_frames_in_range(start=0, stop=10).data.to(
             device
         )
         frame_rate = source_decoder.metadata.average_fps
         percentage, atol = (96, 2) if device == "cuda" else (99, 2)
 
-        output_path = tmp_path / f"test.{format}"
-        enc = StreamingEncoder(output_path)
+        enc, encoder_output = self._create_encoder(method, tmp_path, format)
         add_video_kwargs = {
             "height": source_frames.shape[2],
             "width": source_frames.shape[3],
@@ -1615,7 +1633,9 @@ def test_add_video_and_encode_frames(self, tmp_path, format, device):
         enc.close()
 
         decoded_frames = (
-            VideoDecoder(str(output_path)).get_frames_in_range(start=0, stop=10).data
+            VideoDecoder(self._get_decoder_source(encoder_output))
+            .get_frames_in_range(start=0, stop=10)
+            .data
         )
         assert_tensor_close_on_at_least(
             decoded_frames, source_frames.cpu(), percentage=percentage, atol=atol
@@ -1625,11 +1645,20 @@ def test_create_invalid_path(self):
         with pytest.raises(RuntimeError, match="make sure it's a valid path"):
             StreamingEncoder("/nonexistent/dir/test.mp4")
 
-    def test_create_invalid_format(self, tmp_path):
-        with pytest.raises(RuntimeError, match="check the desired extension"):
-            StreamingEncoder(tmp_path / "test.bad_extension")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_create_invalid_format(self, tmp_path, method):
+        if method == "to_file":
+            with pytest.raises(RuntimeError, match="check the desired extension"):
+                StreamingEncoder(tmp_path / "test.bad_extension")
+        elif method == "to_file_like":
+            with pytest.raises(
+                RuntimeError,
+                match=r"Check the desired format\? Got format=bad_extension",
+            ):
+                StreamingEncoder(io.BytesIO(), format="bad_extension")
 
     @pytest.mark.parametrize("format", ["mp4", "mov"])
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
     @pytest.mark.parametrize(
         "device",
         (
@@ -1645,16 +1674,20 @@ def test_create_invalid_format(self, tmp_path):
             ),
         ),
     )
-    def test_fragmented_mp4(self, format, tmp_path, device):
+    def test_fragmented_mp4(self, format, tmp_path, method, device):
         source_decoder = VideoDecoder(str(TEST_SRC_2_720P.path))
         source_frames = source_decoder.get_frames_in_range(start=0, stop=10).data.to(
             device
         )
         frame_rate = source_decoder.metadata.average_fps
         percentage, atol = (96, 2) if device == "cuda" else (99, 2)
 
-        output_path = tmp_path / f"test.{format}"
-        enc = StreamingEncoder(output_path)
+        enc, encoder_output = self._create_encoder(method, tmp_path, format)
+        # In addition to the fragmentation flag, "flush_packets" and "threads"
+        # are necessary to decode frames before close().
+        # See frag flags: https://ffmpeg.org/ffmpeg-formats.html#Fragmentation
+        # TODO MultiStreamEncoder: Get a better understanding of which options
+        # are necessary for reading fragmented mp4s
         extra_options = {
             "movflags": "+frag_every_frame+empty_moov",
             "flush_packets": "1",
@@ -1676,9 +1709,10 @@ def test_fragmented_mp4(self, format, tmp_path, device):
             extra_options=extra_options,
         )
         enc.open()
+        # Here, we decode the available fragmented mp4 frames before calling close()
         for batch in [source_frames[:5], source_frames[5:]]:
             video.write(batch)
-            mid_decoder = VideoDecoder(str(output_path))
+            mid_decoder = VideoDecoder(self._get_decoder_source(encoder_output))
             num_available = len(mid_decoder)
             assert num_available > 0
             assert_tensor_close_on_at_least(
@@ -1689,25 +1723,31 @@ def test_fragmented_mp4(self, format, tmp_path, device):
             )
 
         enc.close()
+        # After close, all frames must be decodable
         assert_tensor_close_on_at_least(
-            VideoDecoder(str(output_path)).get_frames_in_range(start=0, stop=10).data,
+            VideoDecoder(self._get_decoder_source(encoder_output))
+            .get_frames_in_range(start=0, stop=10)
+            .data,
             source_frames.cpu(),
             percentage=percentage,
             atol=atol,
         )
 
-    def test_add_video_twice_errors(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_add_video_twice_errors(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         enc.add_video(height=64, width=64, frame_rate=30.0)
         with pytest.raises(RuntimeError, match="already been added"):
             enc.add_video(height=64, width=64, frame_rate=24.0)
 
-    def test_add_audio_twice_errors(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_add_audio_twice_errors(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         enc.add_audio(sample_rate=44100, num_channels=2)
         with pytest.raises(RuntimeError, match="already been added"):
             enc.add_audio(sample_rate=16000, num_channels=1)
 
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
     @pytest.mark.parametrize(
         "device",
         (
@@ -1723,15 +1763,17 @@ def test_add_audio_twice_errors(self, tmp_path):
             ),
         ),
     )
-    def test_write_frames_mismatched_dimensions_errors(self, tmp_path, device):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    def test_write_frames_mismatched_dimensions_errors(self, tmp_path, method, device):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         video = enc.add_video(height=256, width=256, frame_rate=30.0, device=device)
         enc.open()
+        # write with wrong size errors
         frames_128 = torch.randint(
             0, 256, (2, 3, 128, 128), dtype=torch.uint8, device=device
         )
         with pytest.raises(RuntimeError, match="same dimensions"):
             video.write(frames_128)
+        # write with different size than first also errors
         frames_256 = torch.randint(
             0, 256, (2, 3, 256, 256), dtype=torch.uint8, device=device
         )
@@ -1743,8 +1785,9 @@ def test_write_frames_mismatched_dimensions_errors(self, tmp_path, device):
             video.write(frames_512)
 
     @pytest.mark.needs_cuda
-    def test_write_frames_different_devices_errors(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_write_frames_different_devices_errors(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         video = enc.add_video(height=64, width=64, frame_rate=30.0)
         enc.open()
         cpu_frames = torch.randint(0, 256, (2, 3, 64, 64), dtype=torch.uint8)
@@ -1754,72 +1797,79 @@ def test_write_frames_different_devices_errors(self, tmp_path):
             video.write(cuda_frames)
 
     @pytest.mark.needs_cuda
-    def test_write_samples_on_cuda_errors(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.wav")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_write_samples_on_cuda_errors(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "wav")
         audio = enc.add_audio(sample_rate=44100, num_channels=1)
         enc.open()
         cuda_samples = torch.randn(1, 1000, device="cuda")
         with pytest.raises(RuntimeError, match="samples must be on CPU, got cuda"):
             audio.write(cuda_samples)
 
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
     @pytest.mark.parametrize(
         "device", ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda))
     )
-    def test_write_frames_without_open_errors(self, tmp_path, device):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    def test_write_frames_without_open_errors(self, tmp_path, method, device):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         video = enc.add_video(height=64, width=64, frame_rate=30.0, device=device)
         frames = torch.randint(0, 256, (5, 3, 64, 64), dtype=torch.uint8, device=device)
         with pytest.raises(
             RuntimeError, match="Call open\\(\\) before addFrames\\(\\)"
         ):
             video.write(frames)
 
-    def test_write_samples_without_open_errors(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_write_samples_without_open_errors(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         audio = enc.add_audio(sample_rate=44100, num_channels=1)
         samples = torch.randn(1, 1000)
         with pytest.raises(
             RuntimeError, match="Call open\\(\\) before addSamples\\(\\)"
         ):
             audio.write(samples)
 
-    def test_write_samples_mismatched_channels_errors(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.wav")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_write_samples_mismatched_channels_errors(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "wav")
         audio = enc.add_audio(sample_rate=44100, num_channels=1)
         enc.open()
         samples = torch.randn(2, 1000)  # 2 channels but stream expects 1
         with pytest.raises(RuntimeError, match="Expected 1 channels, got 2"):
             audio.write(samples)
 
-    def test_open_without_stream_errors(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_open_without_stream_errors(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         with pytest.raises(
             RuntimeError,
             match="Call addVideoStream\\(\\) or addAudioStream\\(\\) before open\\(\\)",
         ):
             enc.open()
 
-    def test_open_twice_errors(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_open_twice_errors(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         enc.add_video(height=64, width=64, frame_rate=30.0)
         enc.open()
         with pytest.raises(RuntimeError, match="open\\(\\) was already called"):
             enc.open()
 
-    def test_add_audio_invalid_bit_rate_errors(self, tmp_path):
-        enc = StreamingEncoder(tmp_path / "test.mp4")
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_add_audio_invalid_bit_rate_errors(self, tmp_path, method):
+        enc, _ = self._create_encoder(method, tmp_path, "mp4")
         enc.add_audio(sample_rate=44100, num_channels=2, bit_rate=-1)
         with pytest.raises(RuntimeError, match="bit_rate=-1 must be >= 0"):
             enc.open()
 
-    def test_add_audio_and_encode_samples(self, tmp_path):
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_add_audio_and_encode_samples(self, tmp_path, method):
         source_audio = AudioDecoder(str(SINE_MONO_S32.path)).get_all_samples()
         samples = source_audio.data
         sample_rate = source_audio.sample_rate
         num_channels = samples.shape[0]
 
-        output_path = tmp_path / "test.wav"
-        enc = StreamingEncoder(output_path)
+        enc, encoder_output = self._create_encoder(method, tmp_path, "wav")
         audio = enc.add_audio(sample_rate=sample_rate, num_channels=num_channels)
         enc.open()
         chunk_lengths = [1, 50, 1000, 0, 25]
@@ -1830,7 +1880,9 @@ def test_add_audio_and_encode_samples(self, tmp_path):
         audio.write(samples[:, offset:])
         enc.close()
 
-        decoded = AudioDecoder(str(output_path)).get_all_samples()
+        decoded = AudioDecoder(
+            self._get_decoder_source(encoder_output)
+        ).get_all_samples()
         assert decoded.data.shape[0] == num_channels
         assert decoded.sample_rate == sample_rate
         torch.testing.assert_close(decoded.data, samples, atol=1e-4, rtol=0)
@@ -1849,7 +1901,8 @@ def test_add_audio_and_encode_samples(self, tmp_path):
             "mov",
         ),
     )
-    def test_add_audio_and_video_and_encode(self, tmp_path, format):
+    @pytest.mark.parametrize("method", ("to_file", "to_file_like"))
+    def test_add_audio_and_video_and_encode(self, tmp_path, format, method):
         source_video_decoder = VideoDecoder(str(NASA_VIDEO.path))
         source_frames = source_video_decoder.get_frames_in_range(
             start=0, stop=len(source_video_decoder)
@@ -1859,8 +1912,7 @@ def test_add_audio_and_video_and_encode(self, tmp_path, format):
         source_samples = source_audio.data
         sample_rate = source_audio.sample_rate
 
-        output_path = tmp_path / f"test.{format}"
-        enc = StreamingEncoder(output_path)
+        enc, encoder_output = self._create_encoder(method, tmp_path, format)
         video = enc.add_video(
             height=source_frames.shape[2],
             width=source_frames.shape[3],
@@ -1881,18 +1933,24 @@ def test_add_audio_and_video_and_encode(self, tmp_path, format):
         audio.write(source_samples[:, half_samples:])
         enc.close()
 
-        decoded_video_decoder = VideoDecoder(str(output_path))
+        source = self._get_decoder_source(encoder_output)
+
+        decoded_video_decoder = VideoDecoder(source)
         decoded_frames = decoded_video_decoder.get_frames_in_range(
             start=0, stop=len(decoded_video_decoder)
         ).data
         assert_tensor_close_on_at_least(
             decoded_frames, source_frames, percentage=99, atol=2
         )
 
-        audio_decoder = AudioDecoder(str(output_path))
+        audio_decoder = AudioDecoder(source)
         decoded_audio = audio_decoder.get_all_samples()
         assert decoded_audio.sample_rate == sample_rate
         assert decoded_audio.data.shape[0] == source_samples.shape[0]
+        # Codecs for lossy audio formats (not WAV or FLAC) can add padding which causes
+        # sample count to differ, so we only compare the smaller sample count.
+        # TODO MultiStreamEncoder: The previous AudioEncoder didn't need
+        # padding after introducing a FIFO. Investigate why this is needed.
         num_samples_to_compare = min(
             decoded_audio.data.shape[1], source_samples.shape[1]
         )