Update public API of multi-stream encoder (#1416)

Author: NicolasHug

src/torchcodec/encoders/__init__.py

@@ -1,2 +1,3 @@
 from ._audio_encoder import AudioEncoder  # noqa
+from ._multi_stream_encoder import Encoder  # noqa
 from ._video_encoder import VideoEncoder  # noqa

src/torchcodec/encoders/_multi_stream_encoder.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from typing import Any
 
 from torch import Tensor
@@ -16,7 +17,7 @@ def __init__(self, encoder_tensor: Tensor, stream_index: int):
         self._encoder_tensor = encoder_tensor
         self._stream_index = stream_index
 
-    def write(self, frames: Tensor) -> None:
+    def add_frames(self, frames: Tensor) -> None:
         _core.streaming_encoder_add_frames(
             self._encoder_tensor, frames, self._stream_index
         )
@@ -27,13 +28,13 @@ def __init__(self, encoder_tensor: Tensor, stream_index: int):
         self._encoder_tensor = encoder_tensor
         self._stream_index = stream_index
 
-    def write(self, samples: Tensor) -> None:
+    def add_samples(self, samples: Tensor) -> None:
         _core.streaming_encoder_add_samples(
             self._encoder_tensor, samples, self._stream_index
         )
 
 
-class StreamingEncoder:
+class Encoder:
     def __init__(self):
         self._encoder_tensor = _core.create_streaming_encoder()
 
@@ -73,33 +74,31 @@ def add_audio(
         sample_rate: int,
         num_channels: int,
         bit_rate: int | None = None,
-        # TODO MultiStreamEncoder: Decide on public API for 'output' params
-        output_num_channels: int | None = None,
-        output_sample_rate: int | None = None,
+        out_num_channels: int | None = None,
+        out_sample_rate: int | None = None,
     ) -> _AudioStream:
         stream_index = _core.streaming_encoder_add_audio_stream(
             self._encoder_tensor,
             sample_rate=sample_rate,
             num_channels=num_channels,
             bit_rate=bit_rate,
-            output_num_channels=output_num_channels,
-            output_sample_rate=output_sample_rate,
+            output_num_channels=out_num_channels,
+            output_sample_rate=out_sample_rate,
         )
         return _AudioStream(self._encoder_tensor, stream_index)
 
-    # TODO MultiStreamEncoder: Maybe there should 2 separate methods, one for
-    # file, one for file-like.
-    def open(self, dest, *, format: str | None = None) -> "StreamingEncoder":
-        if format is not None:
-            _core.streaming_encoder_open_file_like(self._encoder_tensor, format, dest)
-        else:
-            _core.streaming_encoder_open_file(self._encoder_tensor, str(dest))
+    def open_file(self, dest: str | Path) -> "Encoder":
+        _core.streaming_encoder_open_file(self._encoder_tensor, str(dest))
+        return self
+
+    def open_file_like(self, dest, *, format: str) -> "Encoder":
+        _core.streaming_encoder_open_file_like(self._encoder_tensor, format, dest)
         return self
 
     def close(self) -> None:
         _core.streaming_encoder_close(self._encoder_tensor)
 
-    def __enter__(self) -> "StreamingEncoder":
+    def __enter__(self) -> "Encoder":
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:

test/smoke_test.py

@@ -8,8 +8,7 @@
 from torchcodec import ffmpeg_major_version
 from torchcodec._frame import AudioSamples, Frame, FrameBatch
 from torchcodec.decoders import AudioDecoder, VideoDecoder
-from torchcodec.encoders import AudioEncoder, VideoEncoder
-from torchcodec.encoders._multi_stream_encoder import StreamingEncoder
+from torchcodec.encoders import AudioEncoder, Encoder, VideoEncoder
 
 
 NUM_FRAMES = 10
@@ -300,7 +299,7 @@ def test_resample_on_encode(self, tmp_path):
         assert abs(decoded.data.shape[1] - expected_num_samples) <= 1
 
 
-class TestStreamingEncoder:
+class TestEncoder:
     def test_video_and_audio_chunked(self, tmp_path):
         frames = torch.randint(
             0, 256, (NUM_FRAMES, 3, HEIGHT, WIDTH), dtype=torch.uint8
@@ -309,7 +308,7 @@ def test_video_and_audio_chunked(self, tmp_path):
         samples = torch.rand(NUM_AUDIO_CHANNELS, NUM_SAMPLES) * 2 - 1
         path = tmp_path / "av.mkv"
 
-        enc = StreamingEncoder()
+        enc = Encoder()
         video = enc.add_video(
             height=HEIGHT,
             width=WIDTH,
@@ -318,12 +317,12 @@ def test_video_and_audio_chunked(self, tmp_path):
             crf=0,
         )
         audio = enc.add_audio(sample_rate=sr, num_channels=NUM_AUDIO_CHANNELS)
-        enc.open(dest=path)
+        enc.open_file(path)
         with enc:
-            video.write(frames[:5])
-            audio.write(samples[:, : NUM_SAMPLES // 2])
-            video.write(frames[5:])
-            audio.write(samples[:, NUM_SAMPLES // 2 :])
+            video.add_frames(frames[:5])
+            audio.add_samples(samples[:, : NUM_SAMPLES // 2])
+            video.add_frames(frames[5:])
+            audio.add_samples(samples[:, NUM_SAMPLES // 2 :])
 
         video_dec = VideoDecoder(path)
         assert len(video_dec) == NUM_FRAMES
@@ -358,18 +357,18 @@ def test_cuda_encoding(self, tmp_path):
         frames = frames.cuda()
 
         cuda_path = tmp_path / "cuda.mp4"
-        enc = StreamingEncoder()
+        enc = Encoder()
         video = enc.add_video(
             height=HEIGHT,
             width=WIDTH,
             frame_rate=FRAME_RATE,
             device="cuda",
         )
-        with enc.open(dest=cuda_path):
-            video.write(frames)
+        with enc.open_file(cuda_path):
+            video.add_frames(frames)
 
         cpu_path = tmp_path / "cpu.mp4"
-        cpu_enc = StreamingEncoder()
+        cpu_enc = Encoder()
         cpu_video = cpu_enc.add_video(
             height=HEIGHT,
             width=WIDTH,