to_filelike, update test

Author: Daniel Flores

src/torchcodec/_core/__init__.py

@@ -26,6 +26,7 @@
     encode_audio_to_file_like,
     encode_audio_to_tensor,
     encode_video_to_file,
+    encode_video_to_file_like,
     encode_video_to_tensor,
     get_ffmpeg_library_versions,
     get_frame_at_index,

src/torchcodec/_core/custom_ops.cpp

@@ -40,6 +40,8 @@ TORCH_LIBRARY(torchcodec_ns, m) {
       "encode_video_to_file(Tensor frames, int frame_rate, str filename, int? crf=None) -> ()");
   m.def(
       "encode_video_to_tensor(Tensor frames, int frame_rate, str format, int? crf=None) -> Tensor");
+  m.def(
+      "_encode_video_to_file_like(Tensor frames, int frame_rate, str format, int file_like_context, int? crf=None) -> ()");
   m.def(
       "create_from_tensor(Tensor video_tensor, str? seek_mode=None) -> Tensor");
   m.def(
@@ -606,6 +608,30 @@ at::Tensor encode_video_to_tensor(
       .encodeToTensor();
 }
 
+void _encode_video_to_file_like(
+    const at::Tensor& frames,
+    int64_t frame_rate,
+    std::string_view format,
+    int64_t file_like_context,
+    std::optional<int64_t> crf = std::nullopt) {
+  auto fileLikeContext =
+      reinterpret_cast<AVIOFileLikeContext*>(file_like_context);
+  TORCH_CHECK(
+      fileLikeContext != nullptr, "file_like_context must be a valid pointer");
+  std::unique_ptr<AVIOFileLikeContext> avioContextHolder(fileLikeContext);
+
+  VideoStreamOptions videoStreamOptions;
+  videoStreamOptions.crf = crf;
+
+  VideoEncoder encoder(
+      frames,
+      validateInt64ToInt(frame_rate, "frame_rate"),
+      format,
+      std::move(avioContextHolder),
+      videoStreamOptions);
+  encoder.encode();
+}
+
 // For testing only. We need to implement this operation as a core library
 // function because what we're testing is round-tripping pts values as
 // double-precision floating point numbers from C++ to Python and back to C++.
@@ -870,6 +896,7 @@ TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) {
   m.impl("_encode_audio_to_file_like", &_encode_audio_to_file_like);
   m.impl("encode_video_to_file", &encode_video_to_file);
   m.impl("encode_video_to_tensor", &encode_video_to_tensor);
+  m.impl("_encode_video_to_file_like", &_encode_video_to_file_like);
   m.impl("seek_to_pts", &seek_to_pts);
   m.impl("add_video_stream", &add_video_stream);
   m.impl("_add_video_stream", &_add_video_stream);

src/torchcodec/_core/ops.py

@@ -104,6 +104,9 @@ def load_torchcodec_shared_libraries():
 encode_video_to_tensor = torch._dynamo.disallow_in_graph(
     torch.ops.torchcodec_ns.encode_video_to_tensor.default
 )
+_encode_video_to_file_like = torch._dynamo.disallow_in_graph(
+    torch.ops.torchcodec_ns._encode_video_to_file_like.default
+)
 create_from_tensor = torch._dynamo.disallow_in_graph(
     torch.ops.torchcodec_ns.create_from_tensor.default
 )
@@ -203,6 +206,33 @@ def encode_audio_to_file_like(
     )
 
 
+def encode_video_to_file_like(
+    frames: torch.Tensor,
+    frame_rate: int,
+    format: str,
+    file_like: Union[io.RawIOBase, io.BufferedIOBase],
+    crf: Optional[int] = None,
+) -> None:
+    """Encode video frames to a file-like object.
+
+    Args:
+        frames: Video frames tensor
+        frame_rate: Frame rate in frames per second
+        format: Video format (e.g., "mp4", "mov", "mkv")
+        file_like: File-like object that supports write() and seek() methods
+        crf: Optional constant rate factor for encoding quality
+    """
+    assert _pybind_ops is not None
+
+    _encode_video_to_file_like(
+        frames,
+        frame_rate,
+        format,
+        _pybind_ops.create_file_like_context(file_like, True),  # True means for writing
+        crf,
+    )
+
+
 def get_frames_at_indices(
     decoder: torch.Tensor, *, frame_indices: Union[torch.Tensor, list[int]]
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -302,6 +332,17 @@ def encode_video_to_tensor_abstract(
     return torch.empty([], dtype=torch.long)
 
 
+@register_fake("torchcodec_ns::_encode_video_to_file_like")
+def _encode_video_to_file_like_abstract(
+    frames: torch.Tensor,
+    frame_rate: int,
+    format: str,
+    file_like_context: int,
+    crf: Optional[int] = None,
+) -> None:
+    return
+
+
 @register_fake("torchcodec_ns::create_from_tensor")
 def create_from_tensor_abstract(
     video_tensor: torch.Tensor, seek_mode: Optional[str]

test/test_ops.py

@@ -29,6 +29,7 @@
     create_from_tensor,
     encode_audio_to_file,
     encode_video_to_file,
+    encode_video_to_file_like,
     encode_video_to_tensor,
     get_ffmpeg_library_versions,
     get_frame_at_index,
@@ -1379,17 +1380,19 @@ def test_bad_input(self, tmp_path):
                 filename="./bad/path.mp3",
             )
 
-    def decode(self, file_path=None, tensor=None) -> torch.Tensor:
+    def decode(self, file_path=None, tensor=None, file_like=None) -> torch.Tensor:
         if file_path is not None:
             decoder = create_from_file(str(file_path), seek_mode="approximate")
         elif tensor is not None:
             decoder = create_from_tensor(tensor, seek_mode="approximate")
+        elif file_like is not None:
+            decoder = create_from_file_like(file_like, seek_mode="approximate")
         add_video_stream(decoder)
         frames, *_ = get_frames_in_range(decoder, start=0, stop=60)
         return frames
 
     @pytest.mark.parametrize("format", ("mov", "mp4", "mkv", "webm"))
-    @pytest.mark.parametrize("output_method", ("to_file", "to_tensor"))
+    @pytest.mark.parametrize("output_method", ("to_file", "to_tensor", "to_file_like"))
     def test_video_encoder_round_trip(self, tmp_path, format, output_method):
         # Test that decode(encode(decode(asset))) == decode(asset)
         ffmpeg_version = get_ffmpeg_major_version()
@@ -1416,12 +1419,24 @@ def test_video_encoder_round_trip(self, tmp_path, format, output_method):
                 crf=0,
             )
             round_trip_frames = self.decode(file_path=encoded_path).data
-        else:  # to_tensor
+        elif output_method == "to_tensor":
             format = "matroska" if format == "mkv" else format
             encoded_tensor = encode_video_to_tensor(
                 source_frames, frame_rate, format, crf=0
             )
             round_trip_frames = self.decode(tensor=encoded_tensor).data
+        else:  # to_file_like
+            format = "matroska" if format == "mkv" else format
+            file_like = io.BytesIO()
+            encode_video_to_file_like(
+                frames=source_frames,
+                frame_rate=frame_rate,
+                format=format,
+                file_like=file_like,
+                crf=0,
+            )
+            file_like.seek(0)
+            round_trip_frames = self.decode(file_like=file_like).data
 
         assert source_frames.shape == round_trip_frames.shape
         assert source_frames.dtype == round_trip_frames.dtype