Rework public API of image encoders (#1586)

Author: NicolasHug

benchmarks/encoders/benchmark_image_encoders.py

@@ -42,9 +42,9 @@
 
 torch.set_num_threads(1)
 
-from torchcodec.encoders._image_encoders import (  # noqa: E402
-    encode_jpeg as tc_encode_jpeg,
-    encode_png as tc_encode_png,
+from torchcodec.encoders import (  # noqa: E402
+    JpegEncoder as TCJpegEncoder,
+    PngEncoder as TCPngEncoder,
 )
 from torchvision.io import (  # noqa: E402
     encode_jpeg as tv_encode_jpeg,
@@ -111,18 +111,18 @@ def load_source_image(override: Path | None) -> Tensor:
 def make_encode_fn(backend, fmt, dest, img, path, param):
     """Return a zero-arg callable that encodes `img` (in `fmt`) to `dest`, or
     None if the backend doesn't support that destination. `param` is the quality
-    (jpeg) or compression_level (png). torchcodec/torchvision take the encode
-    param positionally so this stays format-agnostic."""
+    (jpeg) or compression_level (png)."""
     is_png = fmt == "png"
     if backend == "torchcodec":
-        tc_encode = tc_encode_png if is_png else tc_encode_jpeg
+        Encoder = TCPngEncoder if is_png else TCJpegEncoder
+        kwargs = {"compression_level" if is_png else "quality": param}
         if dest == "file":
-            return lambda: tc_encode(img, path, param)
+            return lambda: Encoder(img).to_file(path, **kwargs)
         if dest == "file_like":
-            return lambda: tc_encode(img, io.BytesIO(), param)
-        # tensor: native dest=None. On CUDA it returns a device tensor
+            return lambda: Encoder(img).to_file_like(io.BytesIO(), **kwargs)
+        # tensor: native to_tensor(). On CUDA it returns a device tensor
         # (zero-copy); .cpu() brings it to host to match torchvision's path.
-        return lambda: tc_encode(img, None, param).cpu()
+        return lambda: Encoder(img).to_tensor(**kwargs).cpu()
 
     if backend == "torchvision":
         if is_png:
@@ -177,7 +177,9 @@ def make_batch_encode_fn(backend, imgs, quality):
     return the encoded bytes on the host (torchvision's live on the input
     device, so we .cpu() them)."""
     if backend == "torchcodec":
-        return lambda: [tc_encode_jpeg(img, quality=quality).cpu() for img in imgs]
+        return lambda: [
+            TCJpegEncoder(img).to_tensor(quality=quality).cpu() for img in imgs
+        ]
 
     if backend == "torchvision":
         return lambda: [t.cpu() for t in tv_encode_jpeg(list(imgs), quality=quality)]

src/torchcodec/encoders/__init__.py

@@ -1,3 +1,4 @@
 from ._audio_encoder import AudioEncoder  # noqa
+from ._image_encoders import JpegEncoder, PngEncoder  # noqa
 from ._multi_stream_encoder import AudioStream, Encoder, VideoStream  # noqa
 from ._video_encoder import VideoEncoder  # noqa

src/torchcodec/encoders/_image_encoders.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 
 import torch
+from torch import Tensor
 
 from torchcodec._core.ops import (
     create_file_like_context,
@@ -19,112 +20,65 @@
 )
 
 
-def _encode_to_dest(input, dest, param, *, to_file, to_file_like) -> None:
-    if isinstance(dest, (str, Path)):
-        to_file(input, str(dest), param)
-    else:
-        # Assume file-like, it gets validated in C++ (it's tested).
-        to_file_like(input, create_file_like_context(dest, True), param)
-
-
-def _encode_to_tensor_through_bytesio(input, param, to_file_like) -> torch.Tensor:
+def _encode_to_tensor_through_bytesio(img, param, to_file_like) -> Tensor:
     # Encode into an in-memory BytesIO and wrap its buffer as a 1-D uint8 tensor.
     # getbuffer() (unlike getvalue()) exposes the buffer without copying. We
     # could have native C++ implementation for that in each encoder, but it's
     # not always worth it (based on benchmarks). Currently, the only encoder
     # that really needs a dedicated C++ path is JPEG on CUDA.
     buf = io.BytesIO()
-    to_file_like(input, create_file_like_context(buf, True), param)
+    to_file_like(img, create_file_like_context(buf, True), param)
     return torch.frombuffer(buf.getbuffer(), dtype=torch.uint8)
 
 
-def encode_png(
-    input: torch.Tensor,
-    dest: str | Path | None = None,
-    compression_level: int = 6,
-) -> torch.Tensor | None:
-    """Encode a CHW uint8 image tensor into a PNG.
-
-    Args:
-        input (``torch.Tensor``): The image to encode, a 3-dimensional uint8
-            tensor in CHW layout with 1 (grayscale) or 3 (RGB) channels.
-        dest (str, ``pathlib.Path``, file-like object, or ``None``): The
-            destination to write the encoded PNG to. Either a path to the output
-            file, or a file-like object that supports
-            ``write(data: bytes) -> int`` and
-            ``seek(offset: int, whence: int = 0) -> int``, such as
-            ``io.BytesIO()`` or an open file in binary write mode. If ``None``
-            (the default), the encoded bytes are returned as a 1-D uint8 tensor
-            instead of being written anywhere.
-        compression_level (int): zlib compression level between 0 (no
-            compression, fastest) and 9 (max compression, slowest). Default: 6.
-
-    Returns:
-        ``None`` if ``dest`` is a path or file-like object, otherwise a 1-D uint8
-        tensor of the encoded bytes.
-    """
-    if dest is None:
-        return _encode_to_tensor_through_bytesio(
-            input, compression_level, _encode_png_to_file_like
-        )
-    else:
-        _encode_to_dest(
-            input,
-            dest,
-            compression_level,
-            to_file=_encode_png_to_file,
-            to_file_like=_encode_png_to_file_like,
+class JpegEncoder:
+    def __init__(self, img: Tensor) -> None:
+        self._img = img
+
+    def to_file(self, dest: str | Path, *, quality: int = 75) -> None:
+        self._validate_quality(quality)
+        _encode_jpeg_to_file(self._img, str(dest), quality)
+
+    def to_file_like(
+        self, dest: io.RawIOBase | io.BufferedIOBase, *, quality: int = 75
+    ) -> None:
+        self._validate_quality(quality)
+        _encode_jpeg_to_file_like(
+            self._img, create_file_like_context(dest, True), quality
         )
-        return None
-
-
-def encode_jpeg(
-    input: torch.Tensor,
-    dest: str | Path | None = None,
-    quality: int = 75,
-) -> torch.Tensor | None:
-    """Encode a CHW uint8 image tensor into a JPEG.
-
-    Args:
-        input (``torch.Tensor``): The image to encode, a 3-dimensional uint8
-            tensor in CHW layout with 1 (grayscale) or 3 (RGB) channels.
-        dest (str, ``pathlib.Path``, file-like object, or ``None``): The
-            destination to write the encoded JPEG to. Either a path to the output
-            file, or a file-like object that supports
-            ``write(data: bytes) -> int`` and
-            ``seek(offset: int, whence: int = 0) -> int``, such as
-            ``io.BytesIO()`` or an open file in binary write mode. If ``None``
-            (the default), the encoded bytes are returned as a 1-D uint8 tensor
-            instead of being written anywhere.
-        quality (int): Quality of the resulting JPEG, between 1 and 100. Higher
-            means better quality and larger file size. Default: 75.
-
-    Returns:
-        ``None`` if ``dest`` is a path or file-like object. If ``dest`` is
-        ``None``, a 1-D uint8 tensor of the encoded bytes, on the same device as
-        ``input`` (a CUDA input yields a CUDA tensor; call ``.cpu()`` for host
-        bytes).
-
-    If ``input`` is on a CUDA device, encoding is performed on the GPU with
-    nvJPEG. Only 3-channel RGB tensors are supported on CUDA (grayscale must be
-    encoded on the CPU).
-    """
-    if quality < 1 or quality > 100:
-        raise ValueError("Image quality should be a positive number between 1 and 100")
-
-    if dest is None:
-        if input.is_cuda:
-            return _encode_jpeg_to_tensor_cuda(input, quality)
+
+    def to_tensor(self, *, quality: int = 75) -> Tensor:
+        self._validate_quality(quality)
+        if self._img.is_cuda:
+            return _encode_jpeg_to_tensor_cuda(self._img, quality)
         else:
             return _encode_to_tensor_through_bytesio(
-                input, quality, _encode_jpeg_to_file_like
+                self._img, quality, _encode_jpeg_to_file_like
+            )
+
+    @staticmethod
+    def _validate_quality(quality: int) -> None:
+        if quality < 1 or quality > 100:
+            raise ValueError(
+                "Image quality should be a positive number between 1 and 100"
             )
-    else:
-        _encode_to_dest(
-            input,
-            dest,
-            quality,
-            to_file=_encode_jpeg_to_file,
-            to_file_like=_encode_jpeg_to_file_like,
+
+
+class PngEncoder:
+    def __init__(self, img: Tensor) -> None:
+        self._img = img
+
+    def to_file(self, dest: str | Path, *, compression_level: int = 6) -> None:
+        _encode_png_to_file(self._img, str(dest), compression_level)
+
+    def to_file_like(
+        self, dest: io.RawIOBase | io.BufferedIOBase, *, compression_level: int = 6
+    ) -> None:
+        _encode_png_to_file_like(
+            self._img, create_file_like_context(dest, True), compression_level
+        )
+
+    def to_tensor(self, *, compression_level: int = 6) -> Tensor:
+        return _encode_to_tensor_through_bytesio(
+            self._img, compression_level, _encode_png_to_file_like
         )
-        return None

test/smoke_test.py

@@ -34,8 +34,13 @@
     decode_webp,
     VideoDecoder,
 )
-from torchcodec.encoders import AudioEncoder, Encoder, VideoEncoder
-from torchcodec.encoders._image_encoders import encode_jpeg, encode_png
+from torchcodec.encoders import (
+    AudioEncoder,
+    Encoder,
+    JpegEncoder,
+    PngEncoder,
+    VideoEncoder,
+)
 
 
 @pytest.fixture(autouse=True)
@@ -336,22 +341,30 @@ def test_decode_heic(self):
         assert img.shape == (3, h, w)
 
 
-# CUDA JPEG encoding is triggered by the input tensor's device, so we wrap
-# encode_jpeg to move the input to the GPU.
-def _encode_jpeg_cuda(input, dest):
-    return encode_jpeg(input.cuda(), dest)
-
-
-# Each backend: (encode_fn, decode_fn, suffix, lossless).
+# Each backend: (Encoder, device, decode_fn, suffix, lossless). CUDA JPEG
+# encoding is triggered by the input tensor's device.
 _IMAGE_ENCODERS = (
     pytest.param(
-        encode_png, decode_png, "png", True, marks=pytest.mark.needs_png, id="png"
+        PngEncoder,
+        "cpu",
+        decode_png,
+        "png",
+        True,
+        marks=pytest.mark.needs_png,
+        id="png",
     ),
     pytest.param(
-        encode_jpeg, decode_jpeg, "jpg", False, marks=pytest.mark.needs_jpeg, id="jpeg"
+        JpegEncoder,
+        "cpu",
+        decode_jpeg,
+        "jpg",
+        False,
+        marks=pytest.mark.needs_jpeg,
+        id="jpeg",
     ),
     pytest.param(
-        _encode_jpeg_cuda,
+        JpegEncoder,
+        "cuda",
         decode_jpeg,
         "jpg",
         False,
@@ -362,26 +375,34 @@ def _encode_jpeg_cuda(input, dest):
 
 
 class TestImageEncoder:
-    def _make_image(self):
-        return torch.randint(0, 256, (3, HEIGHT, WIDTH), dtype=torch.uint8)
+    def _make_image(self, device):
+        return torch.randint(
+            0, 256, (3, HEIGHT, WIDTH), dtype=torch.uint8, device=device
+        )
 
-    @pytest.mark.parametrize("encode_fn,decode_fn,suffix,lossless", _IMAGE_ENCODERS)
-    def test_encode_to_file(self, tmp_path, encode_fn, decode_fn, suffix, lossless):
-        img = self._make_image()
+    @pytest.mark.parametrize(
+        "Encoder,device,decode_fn,suffix,lossless", _IMAGE_ENCODERS
+    )
+    def test_encode_to_file(
+        self, tmp_path, Encoder, device, decode_fn, suffix, lossless
+    ):
+        img = self._make_image(device)
         path = tmp_path / f"out.{suffix}"
-        encode_fn(img, path)
+        Encoder(img).to_file(path)
         assert path.stat().st_size > 0
 
         decoded = decode_fn(str(path))
         assert decoded.shape == (3, HEIGHT, WIDTH)
         if lossless:
-            torch.testing.assert_close(decoded, img, atol=0, rtol=0)
+            torch.testing.assert_close(decoded, img.cpu(), atol=0, rtol=0)
 
-    @pytest.mark.parametrize("encode_fn,decode_fn,suffix,lossless", _IMAGE_ENCODERS)
-    def test_encode_to_file_like(self, encode_fn, decode_fn, suffix, lossless):
-        img = self._make_image()
+    @pytest.mark.parametrize(
+        "Encoder,device,decode_fn,suffix,lossless", _IMAGE_ENCODERS
+    )
+    def test_encode_to_file_like(self, Encoder, device, decode_fn, suffix, lossless):
+        img = self._make_image(device)
         buf = io.BytesIO()
-        encode_fn(img, buf)
+        Encoder(img).to_file_like(buf)
         assert buf.getbuffer().nbytes > 0