.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "generated_examples/encoding/image_encoding.py" .. LINE NUMBERS ARE GIVEN BELOW. .. rst-class:: sphx-glr-example-title .. _sphx_glr_generated_examples_encoding_image_encoding.py: =============== Encoding images =============== In this example, we'll learn how to encode an image tensor to JPEG or PNG using the :class:`~torchcodec.encoders.JpegEncoder` and :class:`~torchcodec.encoders.PngEncoder` classes. .. note:: These encoders supersede the ones from ``torchvision.io``: they are more robust and support more features. See :ref:`sphx_glr_generated_examples_migration_torchvision_migration.py` for a migration guide. .. GENERATED FROM PYTHON SOURCE LINES 25-28 First, a bit of boilerplate: we'll download an image from the web and define a plotting utility. You can ignore that part and jump right below to :ref:`encoding_image`. .. GENERATED FROM PYTHON SOURCE LINES 28-59 .. code-block:: Python import requests import torch from torchcodec.decoders import decode_image url = "https://raw.githubusercontent.com/meta-pytorch/torchcodec/refs/heads/main/docs/source/_static/thumbnails/pigeon_encoding.jpeg" response = requests.get(url, headers={"User-Agent": ""}) if response.status_code != 200: raise RuntimeError(f"Failed to download image. {response.status_code = }.") # The image to encode, a CHW uint8 tensor. It could come from anywhere (e.g. a # model output); here we just decode one. image = decode_image(response.content) def plot(image: torch.Tensor): try: import matplotlib.pyplot as plt from torchvision.transforms.v2.functional import to_pil_image except ImportError: print("Cannot plot, please run `pip install torchvision matplotlib`") return pil_image = to_pil_image(image) fig = plt.figure(figsize=(pil_image.width / 100, pil_image.height / 100)) ax = fig.add_axes([0, 0, 1, 1]) ax.imshow(pil_image) ax.axis("off") .. GENERATED FROM PYTHON SOURCE LINES 60-67 .. _encoding_image: Encoding an image ----------------- Encoders expect a 3D uint8 tensor in CHW layout (1 or 3 channels), which is exactly what our image is: .. GENERATED FROM PYTHON SOURCE LINES 67-70 .. code-block:: Python print(f"{image.shape = }, {image.dtype = }") plot(image) .. image-sg:: /generated_examples/encoding/images/sphx_glr_image_encoding_001.png :alt: image encoding :srcset: /generated_examples/encoding/images/sphx_glr_image_encoding_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none image.shape = torch.Size([3, 288, 300]), image.dtype = torch.uint8 .. GENERATED FROM PYTHON SOURCE LINES 71-76 We instantiate a :class:`~torchcodec.encoders.JpegEncoder` with the image, and encode it. Three destinations are supported: a file with :meth:`~torchcodec.encoders.JpegEncoder.to_file`, a file-like object with :meth:`~torchcodec.encoders.JpegEncoder.to_file_like`, or a 1D uint8 tensor of raw bytes with :meth:`~torchcodec.encoders.JpegEncoder.to_tensor`. .. GENERATED FROM PYTHON SOURCE LINES 76-88 .. code-block:: Python import io from torchcodec.encoders import JpegEncoder encoder = JpegEncoder(image) encoder.to_file("image.jpg") # to a file encoder.to_file_like(io.BytesIO()) # to a file-like object encoded = encoder.to_tensor() # to a tensor print(f"{encoded.shape = }, {encoded.dtype = }") .. rst-class:: sphx-glr-script-out .. code-block:: none encoded.shape = torch.Size([12146]), encoded.dtype = torch.uint8 .. GENERATED FROM PYTHON SOURCE LINES 89-90 That's it! We can decode the encoded bytes back to make sure everything worked: .. GENERATED FROM PYTHON SOURCE LINES 90-96 .. code-block:: Python from torchcodec.decoders import decode_jpeg decoded = decode_jpeg(encoded) print(f"{decoded.shape = }") plot(decoded) .. image-sg:: /generated_examples/encoding/images/sphx_glr_image_encoding_002.png :alt: image encoding :srcset: /generated_examples/encoding/images/sphx_glr_image_encoding_002.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none decoded.shape = torch.Size([3, 288, 300]) .. GENERATED FROM PYTHON SOURCE LINES 97-99 :class:`~torchcodec.encoders.PngEncoder` works exactly the same way, and PNG is lossless (unlike JPEG): .. GENERATED FROM PYTHON SOURCE LINES 99-104 .. code-block:: Python from torchcodec.encoders import PngEncoder encoded = PngEncoder(image).to_tensor() print(f"{encoded.shape = }") .. rst-class:: sphx-glr-script-out .. code-block:: none encoded.shape = torch.Size([104350]) .. GENERATED FROM PYTHON SOURCE LINES 105-108 Both encoders support encoding options: ``JpegEncoder`` takes a ``quality`` (1-100), and ``PngEncoder`` takes a ``compression_level`` (0-9). For example, a lower JPEG quality yields a smaller output: .. GENERATED FROM PYTHON SOURCE LINES 108-112 .. code-block:: Python small = JpegEncoder(image).to_tensor(quality=10) large = JpegEncoder(image).to_tensor(quality=95) print(f"{small.numel() = }, {large.numel() = }") .. rst-class:: sphx-glr-script-out .. code-block:: none small.numel() = 3824, large.numel() = 28251 .. GENERATED FROM PYTHON SOURCE LINES 113-131 Encoding JPEGs on GPU --------------------- ``JpegEncoder`` can encode directly on a CUDA device with nvJPEG: just pass it an image that already lives on the GPU, and the encoding happens there. Only 3-channel RGB images are supported on CUDA. With :meth:`to_tensor `, the encoded bytes stay on the GPU (call ``.cpu()`` to bring them back to the host). .. code-block:: python from torchcodec.encoders import JpegEncoder encoded = JpegEncoder(image.cuda()).to_tensor() # encoded bytes on the GPU # you can still use to_file and to_file_like, but the encoded bytes will # be copied back to the CPU first. PNG encoding is CPU-only. .. GENERATED FROM PYTHON SOURCE LINES 133-135 Check the docstrings of the encoding methods to learn about the different encoding options. .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.204 seconds) .. _sphx_glr_download_generated_examples_encoding_image_encoding.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: image_encoding.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: image_encoding.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: image_encoding.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_