Merge branch 'main' into codec_select_encode_option

Dan-Flores · web-flow · commit b9435048b524 · 2025-11-13T23:38:57.000-05:00
diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml
@@ -62,7 +62,7 @@ jobs:
         run: python -m pip install --upgrade pip
       - name: Install dependencies and FFmpeg
         run: |
-          python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
+          python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
           conda install "ffmpeg=7.0.1" pkg-config pybind11 -c conda-forge
           ffmpeg -version
       - name: Build and install torchcodec
diff --git a/docs/source/api_ref_transforms.rst b/docs/source/api_ref_transforms.rst
@@ -0,0 +1,17 @@
+.. _transforms:
+
+=====================
+torchcodec.transforms
+=====================
+
+.. currentmodule:: torchcodec.transforms
+
+For a tutorial, see: TODO_DECODER_TRANSFORMS_TUTORIAL.
+
+.. autosummary::
+    :toctree: generated/
+    :nosignatures:
+    :template: dataclass.rst
+
+    DecoderTransform
+    Resize
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -209,6 +209,7 @@ def __call__(self, filename):
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3/", None),
     "torch": ("https://pytorch.org/docs/stable/", None),
+    "torchvision": ("https://docs.pytorch.org/vision/stable/", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
     "PIL": ("https://pillow.readthedocs.io/en/stable/", None),
     "matplotlib": ("https://matplotlib.org/stable/", None),
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -125,3 +125,4 @@ Encoding
    api_ref_decoders
    api_ref_encoders
    api_ref_samplers
+   api_ref_transforms
diff --git a/mypy.ini b/mypy.ini
@@ -4,3 +4,4 @@ files = src/torchcodec
 show_error_codes = True
 pretty = True
 allow_redefinition = True
+follow_untyped_imports = True
diff --git a/src/torchcodec/__init__.py b/src/torchcodec/__init__.py
@@ -9,7 +9,7 @@
 # Note: usort wants to put Frame and FrameBatch after decoders and samplers,
 # but that results in circular import.
 from ._frame import AudioSamples, Frame, FrameBatch  # usort:skip # noqa
-from . import decoders, encoders, samplers  # noqa
+from . import decoders, encoders, samplers, transforms  # noqa
 
 try:
     # Note that version.py is generated during install.
diff --git a/src/torchcodec/_core/Encoder.cpp b/src/torchcodec/_core/Encoder.cpp
@@ -766,6 +766,10 @@ void VideoEncoder::initializeEncoder(
         std::to_string(videoStreamOptions.crf.value()).c_str(),
         0);
   }
+  if (videoStreamOptions.preset.has_value()) {
+    av_dict_set(
+        &options, "preset", videoStreamOptions.preset.value().c_str(), 0);
+  }
   int status = avcodec_open2(avCodecContext_.get(), avCodec, &options);
   av_dict_free(&options);
 
diff --git a/src/torchcodec/_core/StreamOptions.h b/src/torchcodec/_core/StreamOptions.h
@@ -46,13 +46,11 @@ struct VideoStreamOptions {
 
   // Encoding options
   std::optional<std::string> codec;
-  // TODO-VideoEncoder: Consider adding other optional fields here
-  // (bit rate, gop size, max b frames, preset)
-  std::optional<double> crf;
-
   // Optional pixel format for video encoding (e.g., "yuv420p", "yuv444p")
   // If not specified, uses codec's default format.
   std::optional<std::string> pixelFormat;
+  std::optional<double> crf;
+  std::optional<std::string> preset;
 };
 
 struct AudioStreamOptions {
diff --git a/src/torchcodec/_core/custom_ops.cpp b/src/torchcodec/_core/custom_ops.cpp
@@ -37,11 +37,11 @@ TORCH_LIBRARY(torchcodec_ns, m) {
   m.def(
       "_encode_audio_to_file_like(Tensor samples, int sample_rate, str format, int file_like_context, int? bit_rate=None, int? num_channels=None, int? desired_sample_rate=None) -> ()");
   m.def(
-      "encode_video_to_file(Tensor frames, int frame_rate, str filename, str? codec=None, str? pixel_format=None, float? crf=None) -> ()");
+      "encode_video_to_file(Tensor frames, int frame_rate, str filename, str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None) -> ()");
   m.def(
-      "encode_video_to_tensor(Tensor frames, int frame_rate, str format, str? codec=None, str? pixel_format=None, float? crf=None) -> Tensor");
+      "encode_video_to_tensor(Tensor frames, int frame_rate, str format, str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None) -> Tensor");
   m.def(
-      "_encode_video_to_file_like(Tensor frames, int frame_rate, str format, int file_like_context, str? codec=None, str? pixel_format=None, float? crf=None) -> ()");
+      "_encode_video_to_file_like(Tensor frames, int frame_rate, str format, int file_like_context, str? codec=None, str? pixel_format=None, float? crf=None, str? preset=None) -> ()");
   m.def(
       "create_from_tensor(Tensor video_tensor, str? seek_mode=None) -> Tensor");
   m.def(
@@ -604,12 +604,14 @@ void encode_video_to_file(
     int64_t frame_rate,
     std::string_view file_name,
     std::optional<std::string> codec = std::nullopt,
-    std::optional<std::string> pixel_format = std::nullopt,
-    std::optional<double> crf = std::nullopt) {
+    std::optional<std::string_view> pixel_format = std::nullopt,
+    std::optional<double> crf = std::nullopt,
+    std::optional<std::string_view> preset = std::nullopt) {
   VideoStreamOptions videoStreamOptions;
   videoStreamOptions.codec = codec;
   videoStreamOptions.pixelFormat = pixel_format;
   videoStreamOptions.crf = crf;
+  videoStreamOptions.preset = preset;
   VideoEncoder(
       frames,
       validateInt64ToInt(frame_rate, "frame_rate"),
@@ -623,13 +625,15 @@ at::Tensor encode_video_to_tensor(
     int64_t frame_rate,
     std::string_view format,
     std::optional<std::string> codec = std::nullopt,
-    std::optional<std::string> pixel_format = std::nullopt,
-    std::optional<double> crf = std::nullopt) {
+    std::optional<std::string_view> pixel_format = std::nullopt,
+    std::optional<double> crf = std::nullopt,
+    std::optional<std::string_view> preset = std::nullopt) {
   auto avioContextHolder = std::make_unique<AVIOToTensorContext>();
   VideoStreamOptions videoStreamOptions;
   videoStreamOptions.codec = codec;
   videoStreamOptions.pixelFormat = pixel_format;
   videoStreamOptions.crf = crf;
+  videoStreamOptions.preset = preset;
   return VideoEncoder(
              frames,
              validateInt64ToInt(frame_rate, "frame_rate"),
@@ -645,8 +649,9 @@ void _encode_video_to_file_like(
     std::string_view format,
     int64_t file_like_context,
     std::optional<std::string> codec = std::nullopt,
-    std::optional<std::string> pixel_format = std::nullopt,
-    std::optional<double> crf = std::nullopt) {
+    std::optional<std::string_view> pixel_format = std::nullopt,
+    std::optional<double> crf = std::nullopt,
+    std::optional<std::string_view> preset = std::nullopt) {
   auto fileLikeContext =
       reinterpret_cast<AVIOFileLikeContext*>(file_like_context);
   TORCH_CHECK(
@@ -657,6 +662,7 @@ void _encode_video_to_file_like(
   videoStreamOptions.codec = codec;
   videoStreamOptions.pixelFormat = pixel_format;
   videoStreamOptions.crf = crf;
+  videoStreamOptions.preset = preset;
 
   VideoEncoder encoder(
       frames,
diff --git a/src/torchcodec/_core/ops.py b/src/torchcodec/_core/ops.py
@@ -216,6 +216,7 @@ def encode_video_to_file_like(
     codec: Optional[str] = None,
     pixel_format: Optional[str] = None,
     crf: Optional[Union[int, float]] = None,
+    preset: Optional[str] = None,
 ) -> None:
     """Encode video frames to a file-like object.
 
@@ -227,6 +228,7 @@ def encode_video_to_file_like(
         codec: Optional codec name (e.g., "libx264", "h264")
         pixel_format: Optional pixel format (e.g., "yuv420p", "yuv444p")
         crf: Optional constant rate factor for encoding quality
+        preset: Optional encoder preset as string (e.g., "ultrafast", "medium")
     """
     assert _pybind_ops is not None
 
@@ -238,6 +240,7 @@ def encode_video_to_file_like(
         codec,
         pixel_format,
         crf,
+        preset,
     )
 
 
@@ -328,6 +331,7 @@ def encode_video_to_file_abstract(
     codec: Optional[str],
     pixel_format: Optional[str] = None,
     crf: Optional[Union[int, float]] = None,
+    preset: Optional[str] = None,
 ) -> None:
     return
 
@@ -340,6 +344,7 @@ def encode_video_to_tensor_abstract(
     codec: Optional[str],
     pixel_format: Optional[str] = None,
     crf: Optional[Union[int, float]] = None,
+    preset: Optional[str] = None,
 ) -> torch.Tensor:
     return torch.empty([], dtype=torch.long)
 
@@ -350,9 +355,10 @@ def _encode_video_to_file_like_abstract(
     frame_rate: int,
     format: str,
     file_like_context: int,
-    codec: Optional[str] = None,
+    codec: Optional[str],
     pixel_format: Optional[str] = None,
     crf: Optional[Union[int, float]] = None,
+    preset: Optional[str] = None,
 ) -> None:
     return
 
diff --git a/src/torchcodec/decoders/_video_decoder.py b/src/torchcodec/decoders/_video_decoder.py
@@ -8,17 +8,18 @@
 import json
 import numbers
 from pathlib import Path
-from typing import Literal, Optional, Tuple, Union
+from typing import List, Literal, Optional, Sequence, Tuple, Union
 
 import torch
-from torch import device as torch_device, Tensor
+from torch import device as torch_device, nn, Tensor
 
 from torchcodec import _core as core, Frame, FrameBatch
 from torchcodec.decoders._decoder_utils import (
     _get_cuda_backend,
     create_decoder,
     ERROR_REPORTING_INSTRUCTIONS,
 )
+from torchcodec.transforms import DecoderTransform, Resize
 
 
 class VideoDecoder:
@@ -67,6 +68,11 @@ class VideoDecoder:
             probably is. Default: "exact".
             Read more about this parameter in:
             :ref:`sphx_glr_generated_examples_decoding_approximate_mode.py`
+        transforms (sequence of transform objects, optional): Sequence of transforms to be
+            applied to the decoded frames by the decoder itself, in order. Accepts both
+            :class:`~torchcodec.transforms.DecoderTransform` and
+            :class:`~torchvision.transforms.v2.Transform`
+            objects. Read more about this parameter in: TODO_DECODER_TRANSFORMS_TUTORIAL.
         custom_frame_mappings (str, bytes, or file-like object, optional):
             Mapping of frames to their metadata, typically generated via ffprobe.
             This enables accurate frame seeking without requiring a full video scan.
@@ -105,6 +111,7 @@ def __init__(
         num_ffmpeg_threads: int = 1,
         device: Optional[Union[str, torch_device]] = None,
         seek_mode: Literal["exact", "approximate"] = "exact",
+        transforms: Optional[Sequence[Union[DecoderTransform, nn.Module]]] = None,
         custom_frame_mappings: Optional[
             Union[str, bytes, io.RawIOBase, io.BufferedReader]
         ] = None,
@@ -151,13 +158,16 @@ def __init__(
 
         device_variant = _get_cuda_backend()
 
+        transform_specs = _make_transform_specs(transforms)
+
         core.add_video_stream(
             self._decoder,
             stream_index=stream_index,
             dimension_order=dimension_order,
             num_threads=num_ffmpeg_threads,
             device=device,
             device_variant=device_variant,
+            transform_specs=transform_specs,
             custom_frame_mappings=custom_frame_mappings_data,
         )
 
@@ -435,6 +445,78 @@ def _get_and_validate_stream_metadata(
     )
 
 
+def _convert_to_decoder_transforms(
+    transforms: Sequence[Union[DecoderTransform, nn.Module]],
+) -> List[DecoderTransform]:
+    """Convert a sequence of transforms that may contain TorchVision transform
+    objects into a list of only TorchCodec transform objects.
+
+    Args:
+        transforms: Squence of transform objects. The objects can be one of two
+        types:
+                1. torchcodec.transforms.DecoderTransform
+                2. torchvision.transforms.v2.Transform, but our type annotation
+                   only mentions its base, nn.Module. We don't want to take a
+                   hard dependency on TorchVision.
+
+    Returns:
+        List of DecoderTransform objects.
+    """
+    try:
+        from torchvision.transforms import v2
+
+        tv_available = True
+    except ImportError:
+        tv_available = False
+
+    converted_transforms: list[DecoderTransform] = []
+    for transform in transforms:
+        if not isinstance(transform, DecoderTransform):
+            if not tv_available:
+                raise ValueError(
+                    f"The supplied transform, {transform}, is not a TorchCodec "
+                    " DecoderTransform. TorchCodec also accept TorchVision "
+                    "v2 transforms, but TorchVision is not installed."
+                )
+            elif isinstance(transform, v2.Resize):
+                converted_transforms.append(Resize._from_torchvision(transform))
+            else:
+                raise ValueError(
+                    f"Unsupported transform: {transform}. Transforms must be "
+                    "either a TorchCodec DecoderTransform or a TorchVision "
+                    "v2 transform."
+                )
+        else:
+            converted_transforms.append(transform)
+
+    return converted_transforms
+
+
+def _make_transform_specs(
+    transforms: Optional[Sequence[Union[DecoderTransform, nn.Module]]],
+) -> str:
+    """Given a sequence of transforms, turn those into the specification string
+       the core API expects.
+
+    Args:
+        transforms: Optional sequence of transform objects. The objects can be
+            one of two types:
+                1. torchcodec.transforms.DecoderTransform
+                2. torchvision.transforms.v2.Transform, but our type annotation
+                   only mentions its base, nn.Module. We don't want to take a
+                   hard dependency on TorchVision.
+
+    Returns:
+        String of transforms in the format the core API expects: transform
+        specifications separate by semicolons.
+    """
+    if transforms is None:
+        return ""
+
+    transforms = _convert_to_decoder_transforms(transforms)
+    return ";".join([t._make_transform_spec() for t in transforms])
+
+
 def _read_custom_frame_mappings(
     custom_frame_mappings: Union[str, bytes, io.RawIOBase, io.BufferedReader]
 ) -> tuple[Tensor, Tensor, Tensor]:
diff --git a/src/torchcodec/encoders/_video_encoder.py b/src/torchcodec/encoders/_video_encoder.py
diff --git a/src/torchcodec/transforms/__init__.py b/src/torchcodec/transforms/__init__.py
diff --git a/src/torchcodec/transforms/_decoder_transforms.py b/src/torchcodec/transforms/_decoder_transforms.py
diff --git a/test/test_encoders.py b/test/test_encoders.py
diff --git a/test/test_transform_ops.py b/test/test_transform_ops.py
