Add docstrings for Image encoders

Author: NicolasHug

docs/source/api_ref_encoders.rst

@@ -32,3 +32,15 @@ Single-stream encoders
 
     VideoEncoder
     AudioEncoder
+
+
+Image encoders
+--------------
+
+.. autosummary::
+    :toctree: generated/
+    :nosignatures:
+    :template: class.rst
+
+    JpegEncoder
+    PngEncoder

src/torchcodec/encoders/_image_encoders.py

@@ -32,22 +32,69 @@ def _encode_to_tensor_through_bytesio(img, param, to_file_like) -> Tensor:
 
 
 class JpegEncoder:
+    """Encoder for JPEG images.
+
+    Example:
+
+        .. code-block:: python
+
+            from torchcodec.encoders import JpegEncoder
+
+            JpegEncoder(img).to_file("image.jpg")
+            # or encode to a file-like object or to a tensor, see methods below.
+
+    Args:
+        img (``torch.Tensor``): The image to encode, a 3-dimensional uint8 tensor
+            in CHW layout with 1 (grayscale) or 3 (RGB) channels. If on a CUDA
+            device, encoding is performed on the GPU with nvJPEG, and only
+            3-channel RGB is supported.
+    """
+
     def __init__(self, img: Tensor) -> None:
         self._img = img
 
     def to_file(self, dest: str | Path, *, quality: int = 75) -> None:
+        """Encode the image into a JPEG file.
+
+        Args:
+            dest (str or ``pathlib.Path``): The path to the output file, e.g.
+                ``image.jpg``.
+            quality (int, optional): Quality of the output, between 1 and 100.
+                Higher means better quality and larger file size. Default: 75.
+        """
         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:
+        """Encode the image into a file-like object.
+
+        Args:
+            dest: A writable file-like object supporting ``write`` and ``seek``,
+                such as ``io.BytesIO()`` or an open file in binary write mode.
+            quality (int, optional): Quality of the output, between 1 and 100.
+                Higher means better quality and larger file size. Default: 75.
+        """
         self._validate_quality(quality)
         _encode_jpeg_to_file_like(
             self._img, create_file_like_context(dest, True), quality
         )
 
     def to_tensor(self, *, quality: int = 75) -> Tensor:
+        """Encode the image into raw bytes, as a 1D uint8 tensor.
+
+        The returned tensor is on the same device as the input image (a CUDA
+        input yields a CUDA tensor).
+
+        Args:
+            quality (int, optional): Quality of the output, between 1 and 100.
+                Higher means better quality and larger file size. Default: 75.
+
+        Returns:
+            torch.Tensor: The encoded bytes, a 1D uint8 tensor on the same
+            device as the input image.
+        """
         self._validate_quality(quality)
         if self._img.is_cuda:
             return _encode_jpeg_to_tensor_cuda(self._img, quality)
@@ -65,20 +112,64 @@ def _validate_quality(quality: int) -> None:
 
 
 class PngEncoder:
+    """Encoder for PNG images.
+
+    Example:
+
+        .. code-block:: python
+
+            from torchcodec.encoders import PngEncoder
+
+            PngEncoder(img).to_file("image.png")
+            # or encode to a file-like object or to a tensor, see methods below.
+
+    Args:
+        img (``torch.Tensor``): The image to encode, a 3-dimensional uint8 tensor
+            in CHW layout with 1 (grayscale) or 3 (RGB) channels.
+    """
+
     def __init__(self, img: Tensor) -> None:
         self._img = img
 
     def to_file(self, dest: str | Path, *, compression_level: int = 6) -> None:
+        """Encode the image into a PNG file.
+
+        Args:
+            dest (str or ``pathlib.Path``): The path to the output file, e.g.
+                ``image.png``.
+            compression_level (int, optional): zlib compression level between 0
+                (no compression, fastest) and 9 (max compression, slowest).
+                Default: 6.
+        """
         _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 the image into a file-like object.
+
+        Args:
+            dest: A writable file-like object supporting ``write`` and ``seek``,
+                such as ``io.BytesIO()`` or an open file in binary write mode.
+            compression_level (int, optional): zlib compression level between 0
+                (no compression, fastest) and 9 (max compression, slowest).
+                Default: 6.
+        """
         _encode_png_to_file_like(
             self._img, create_file_like_context(dest, True), compression_level
         )
 
     def to_tensor(self, *, compression_level: int = 6) -> Tensor:
+        """Encode the image into raw bytes, as a 1D uint8 tensor.
+
+        Args:
+            compression_level (int, optional): zlib compression level between 0
+                (no compression, fastest) and 9 (max compression, slowest).
+                Default: 6.
+
+        Returns:
+            torch.Tensor: The encoded bytes, a 1D uint8 tensor.
+        """
         return _encode_to_tensor_through_bytesio(
             self._img, compression_level, _encode_png_to_file_like
         )