TorchCodec 0.16 - Image decoding and encoding

TorchCodec 0.16 is out! It is compatible with torch >= 2.11. The headline feature of this release is image decoding and encoding: TorchCodec now natively decodes and encodes JPEG (CPU and CUDA), PNG, WebP, GIF, AVIF and HEIC. These image decoders and encoders replace their torchvision counterparts, which are now deprecated.

TorchCodec is the recommended way to decode and encode images in the PyTorch ecosystem. If you are coming from torchvision, we wrote a migration guide

Image decoding

TorchCodec exposes one entry-point per format, plus a generic decode_image() that automatically detects the format. The API is largely backward-compatible with TorchVision:

from torchcodec.decoders import decode_image, decode_jpeg, decode_png

img = decode_image("image.jpg")  # CHW uint8 tensor, format auto-detected
img = decode_image("image.avif")
img = decode_image("image.heic")

# Or use the format-specific decoders for format-specific options
img = decode_jpeg("image.jpg", device="cuda")

Sources can be a path (str or pathlib.Path), bytes, or a 1D uint8 tensor of encoded bytes:

img = decode_image(open("image.png", "rb").read())
img = decode_image(torch.frombuffer(encoded_bytes, dtype=torch.uint8))

Animated and multi-image formats (WebP, GIF, AVIF, HEIC) decode into an (N, C, H, W) tensor:

from torchcodec.decoders import decode_gif

frames = decode_gif("animated.gif")  # (N, C, H, W)

JPEG decoding is also supported on CUDA, through nvJPEG. For CUDA, prefer passing a batch of sources: the whole batch is decoded in a single nvJPEG call, which is much faster than decoding images one at a time.

from torchcodec.decoders import decode_jpeg

imgs = decode_jpeg(["a.jpg", "b.jpg", "c.jpg"], device="cuda")  # list of CUDA tensors

Read more in our image decoding tutorial

Image encoding

Image encoders follow the same class-based design as our video and audio encoders: build the encoder from a CHW uint8 tensor, then choose where the encoded bytes go: a file, a file-like object, or a tensor.

from torchcodec.encoders import JpegEncoder, PngEncoder

JpegEncoder(img).to_file("image.jpg", quality=90)
PngEncoder(img).to_file("image.png", compression_level=9)

# ... or to a file-like object
import io
buffer = io.BytesIO()
JpegEncoder(img).to_file_like(buffer)

# ... or to a 1D uint8 tensor of encoded bytes
encoded = PngEncoder(img).to_tensor()

JPEG encoding is supported on CUDA as well: pass a CUDA tensor and the encoding happens on the GPU with nvJPEG, with to_tensor() returning a CUDA tensor (no host round-trip).

encoded = JpegEncoder(img_on_cuda).to_tensor(quality=90)  # CUDA uint8 tensor

Read more in our image encoding tutorial

Improvements over torchvision's decoders / encoders

The image decoders and encoders were migrated from torchvision and torchvision-extra-decoders, with the same performance, and they are significantly more capable:

• All color modes for every codec: UNCHANGED, GRAY, GRAY_ALPHA, RGB, RGB_ALPHA. torchvision only supports GRAY for PNG and JPEG, and rejects or ignores it elsewhere.
• Animation and multi-image support: animated WebP, GIF and AVIF, and multi-image HEIC, all decode to (N, C, H, W). torchvision rejects animated WebP, errors on multi-image AVIF, and only decodes the primary HEIC image.
• EXIF orientation applied by default, for JPEG (CPU and CUDA), PNG, WebP, AVIF and HEIC. In torchvision it is opt-in, PNG/JPEG-only, and ignored on CUDA.
• output_dtype control (torch.uint8, torch.uint16, or "auto") on every decoder. torchvision has no equivalent: the output dtype is dictated by the source.
• scalability of JPEG encoding and decoding: TorchCodec allows multiple NVJPEG decoders and encoders instances per process, allowing to scale decoding and encoding throughput in multi-threaded pipelines.
• decode_image() auto-detects all six formats, including AVIF and HEIC. torchvision only handles four.
• Richer inputs: str/Path/bytes/Tensor everywhere, non-contiguous encoded input accepted, and batched input for JPEG on both CPU and CUDA.
• No extra package needed: AVIF works out of the box (libavif is bundled). HEIC works if libheif is found at runtime. We don't bundle it because it is LGPL, so install it yourself (e.g. conda install -c conda-forge libheif). Torchvision required the separate torchvision-extra-decoders package for both, and its decode_image couldn't dispatch to them.
• file-like support for encoders - not supported by torchvision.

Along the way we fixed a number of correctness bugs inherited from torchvision, among them: PNG palette and tRNS transparency handling, GIF frame disposal (now aligned with Pillow), truncated JPEGs erroring instead of returning garbage, correct CMYK/YCCK handling, real grayscale for WebP, progressive AVIF stills, and full-range >8-bit HEIC output.

If you are coming from torchvision, we wrote a migration guide

FFmpeg is now an optional dependency

import torchcodec no longer fails at import time if FFmpeg cannot be found. FFmpeg is still required for video and audio decoding and encoding, but the image decoders and encoders don't need FFmpeg and work in FFmpeg-free environments.

FFmpeg 9 support

TorchCodec now support the recently released FFmpeg 9!

Bug Fixes

• Audio resampling correctness. Decoding a resampled audio stream in chunks now returns exactly the same samples as decoding it in one go. (#1604, #1614, #1615, #1616).
• MPEG-PS seeking. Fixed AudioDecoder seeks on MPEG-PS files (#1619).
• Encoders: to_tensor() no longer emits a spurious warning (#1510).