diff --git a/docs/source/api_ref_decoders.rst b/docs/source/api_ref_decoders.rst index 1186cd424..684edb3aa 100644 --- a/docs/source/api_ref_decoders.rst +++ b/docs/source/api_ref_decoders.rst @@ -88,6 +88,7 @@ Image Decoding decode_gif decode_avif decode_heic + get_image_metadata .. autosummary:: :toctree: generated/ @@ -96,6 +97,13 @@ Image Decoding ImageReadMode +.. autosummary:: + :toctree: generated/ + :nosignatures: + :template: dataclass.rst + + ImageMetadata + .. _decoders_low_level: diff --git a/src/torchcodec/_core/DecodeAvif.cpp b/src/torchcodec/_core/DecodeAvif.cpp index b40b9b97c..afadca7e8 100644 --- a/src/torchcodec/_core/DecodeAvif.cpp +++ b/src/torchcodec/_core/DecodeAvif.cpp @@ -28,6 +28,24 @@ torch::stable::Tensor decode_avif( "prebuilt wheel, please report it to the TorchCodec repo."); } +std::vector probe_avif( + [[maybe_unused]] const torch::stable::Tensor& input) { + STD_TORCH_CHECK( + false, + "get_image_metadata: torchcodec was not compiled with libavif support. " + "Rebuild torchcodec with TORCHCODEC_BUILD_AVIF=1. If you see this error " + "in a prebuilt wheel, please report it to the TorchCodec repo."); +} + +std::vector probe_avif_from_file( + [[maybe_unused]] std::string filename) { + STD_TORCH_CHECK( + false, + "get_image_metadata: torchcodec was not compiled with libavif support. " + "Rebuild torchcodec with TORCHCODEC_BUILD_AVIF=1. If you see this error " + "in a prebuilt wheel, please report it to the TorchCodec repo."); +} + } // namespace facebook::torchcodec #else @@ -193,6 +211,67 @@ torch::stable::Tensor decode_avif( return output; } +namespace { + +ImageProbeResult parse_and_probe(avifDecoder* decoder) { + auto result = avifDecoderParse(decoder); + STD_TORCH_CHECK( + result == AVIF_RESULT_OK, + "avifDecoderParse failed: ", + avifResultToString(result)); + + int64_t width = decoder->image->width; + int64_t height = decoder->image->height; + ExifOrientation exif_orientation = avif_exif_orientation(decoder->image); + if (swaps_width_and_height(exif_orientation)) { + std::swap(width, height); + } + bool has_orientation = (decoder->image->transformFlags & + (AVIF_TRANSFORM_IROT | AVIF_TRANSFORM_IMIR)) != 0; + // Same frame count logic as decode_avif(). + int64_t num_frames = + decoder->imageSequenceTrackPresent ? decoder->imageCount : 1; + return { + width, + height, + decoder->alphaPresent ? 4 : 3, + static_cast(decoder->image->depth), + num_frames, + has_orientation ? static_cast(exif_orientation) + : NO_ORIENTATION}; +} + +} // namespace + +ImageProbeResult probe_avif(const torch::stable::Tensor& input) { + auto contig_input = validate_encoded_data(input); + + DecoderPtr decoder(avifDecoderCreate()); + STD_TORCH_CHECK(decoder != nullptr, "Failed to create avif decoder."); + + auto result = avifDecoderSetIOMemory( + decoder.get(), + contig_input.const_data_ptr(), + contig_input.numel()); + STD_TORCH_CHECK( + result == AVIF_RESULT_OK, + "avifDecoderSetIOMemory failed: ", + avifResultToString(result)); + return parse_and_probe(decoder.get()); +} + +ImageProbeResult probe_avif_from_file(std::string filename) { + DecoderPtr decoder(avifDecoderCreate()); + STD_TORCH_CHECK(decoder != nullptr, "Failed to create avif decoder."); + + auto result = avifDecoderSetIOFile(decoder.get(), filename.c_str()); + STD_TORCH_CHECK( + result == AVIF_RESULT_OK, + "avifDecoderSetIOFile failed: ", + avifResultToString(result)); + return parse_and_probe(decoder.get()); +} + } // namespace facebook::torchcodec #endif // !TORCHCODEC_ENABLE_AVIF diff --git a/src/torchcodec/_core/DecodeAvif.h b/src/torchcodec/_core/DecodeAvif.h index e3e2a3f59..4e8f22e99 100644 --- a/src/torchcodec/_core/DecodeAvif.h +++ b/src/torchcodec/_core/DecodeAvif.h @@ -6,6 +6,9 @@ #pragma once +#include +#include + #include "StableABICompat.h" namespace facebook::torchcodec { @@ -16,4 +19,10 @@ FORCE_PUBLIC_VISIBILITY torch::stable::Tensor decode_avif( int64_t output_dtype, int64_t num_threads); +FORCE_PUBLIC_VISIBILITY std::vector probe_avif( + const torch::stable::Tensor& input); + +FORCE_PUBLIC_VISIBILITY std::vector probe_avif_from_file( + std::string filename); + } // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/DecodeGif.cpp b/src/torchcodec/_core/DecodeGif.cpp index 44e9fae4c..febc85177 100644 --- a/src/torchcodec/_core/DecodeGif.cpp +++ b/src/torchcodec/_core/DecodeGif.cpp @@ -26,6 +26,24 @@ torch::stable::Tensor decode_gif( "prebuilt wheel, please report it to the TorchCodec repo."); } +std::vector probe_gif( + [[maybe_unused]] const torch::stable::Tensor& input) { + STD_TORCH_CHECK( + false, + "get_image_metadata: torchcodec was not compiled with GIF support. " + "Rebuild torchcodec with TORCHCODEC_BUILD_GIF=1. If you see this error " + "in a prebuilt wheel, please report it to the TorchCodec repo."); +} + +std::vector probe_gif_from_file( + [[maybe_unused]] std::string filename) { + STD_TORCH_CHECK( + false, + "get_image_metadata: torchcodec was not compiled with GIF support. " + "Rebuild torchcodec with TORCHCODEC_BUILD_GIF=1. If you see this error " + "in a prebuilt wheel, please report it to the TorchCodec repo."); +} + } // namespace facebook::torchcodec #else @@ -33,9 +51,11 @@ torch::stable::Tensor decode_gif( #include #include +#include #include #include +#include "FileIO.h" #include "ImageCommon.h" #include "giflib/gif_lib.h" @@ -62,6 +82,12 @@ int read_callback(GifFileType* gif_file, GifByteType* output, int bytes) { return static_cast(num_bytes_to_read); } +// Same as read_callback, reading from a file instead of a tensor. +int file_read_callback(GifFileType* gif_file, GifByteType* output, int bytes) { + auto* file = static_cast(gif_file->UserData); + return std::max(file->read(output, bytes), 0); +} + } // namespace torch::stable::Tensor decode_gif( @@ -332,6 +358,121 @@ torch::stable::Tensor decode_gif( return output; } +namespace { + +// Walks the records like DGifSlurp() does, but skips the image data instead of +// LZW-decoding it. Closes gif_file. +ImageProbeResult probe_and_close(GifFileType* gif_file) { + int error = D_GIF_SUCCEEDED; + + auto fail = [&]() { + auto gif_error = gif_file->Error; + DGifCloseFile(gif_file, &error); + STD_TORCH_CHECK(false, "Failed to parse GIF - ", gif_error); + }; + + int64_t num_images = 0; + int64_t first_image_width = 0; + int64_t first_image_height = 0; + // Like DGifSavedExtensionToGCB(), only the first graphics control extension + // preceding an image applies to it. + bool seen_gcb_for_next_image = false; + bool has_transparency = false; + GifRecordType record_type; + do { + if (DGifGetRecordType(gif_file, &record_type) == GIF_ERROR) { + fail(); + } + if (record_type == IMAGE_DESC_RECORD_TYPE) { + if (DGifGetImageDesc(gif_file) == GIF_ERROR) { + fail(); + } + const GifImageDesc& desc = gif_file->Image; + if (desc.Width <= 0 || desc.Height <= 0 || + desc.Width > (INT_MAX / desc.Height)) { + fail(); + } + if (num_images == 0) { + first_image_width = desc.Width; + first_image_height = desc.Height; + } + ++num_images; + seen_gcb_for_next_image = false; + + int code_size = 0; + GifByteType* code_block = nullptr; + if (DGifGetCode(gif_file, &code_size, &code_block) == GIF_ERROR) { + fail(); + } + while (code_block != nullptr) { + if (DGifGetCodeNext(gif_file, &code_block) == GIF_ERROR) { + fail(); + } + } + } else if (record_type == EXTENSION_RECORD_TYPE) { + int ext_code = 0; + GifByteType* ext_data = nullptr; + if (DGifGetExtension(gif_file, &ext_code, &ext_data) == GIF_ERROR) { + fail(); + } + if (ext_data != nullptr && ext_code == GRAPHICS_EXT_FUNC_CODE && + !seen_gcb_for_next_image) { + seen_gcb_for_next_image = true; + GraphicsControlBlock gcb; + if (DGifExtensionToGCB(ext_data[0], ext_data + 1, &gcb) == GIF_OK && + gcb.TransparentColor != NO_TRANSPARENT_COLOR) { + has_transparency = true; + } + } + while (ext_data != nullptr) { + if (DGifGetExtensionNext(gif_file, &ext_data) == GIF_ERROR) { + fail(); + } + } + } + } while (record_type != TERMINATE_RECORD_TYPE); + + int64_t width = std::max(gif_file->SWidth, first_image_width); + int64_t height = std::max(gif_file->SHeight, first_image_height); + DGifCloseFile(gif_file, &error); + STD_TORCH_CHECK( + num_images > 0, "GIF file should contain at least one image!"); + + return { + width, height, has_transparency ? 4 : 3, 8, num_images, NO_ORIENTATION}; +} + +} // namespace + +ImageProbeResult probe_gif(const torch::stable::Tensor& input) { + auto contig_input = validate_encoded_data(input); + + int error = D_GIF_SUCCEEDED; + SourceCtx source_ctx{ + .ptr = contig_input.const_data_ptr(), + .count = static_cast(contig_input.numel())}; + GifFileType* gif_file = + DGifOpen(static_cast(&source_ctx), read_callback, &error); + STD_TORCH_CHECK( + (gif_file != nullptr) && (error == D_GIF_SUCCEEDED), + "DGifOpen() failed - ", + error); + return probe_and_close(gif_file); +} + +ImageProbeResult probe_gif_from_file(std::string filename) { + // We don't use DGifOpenFileName(): it opens the file in text mode on Windows. + FileIO file(filename, FileIO::Mode::Read); + int error = D_GIF_SUCCEEDED; + GifFileType* gif_file = + DGifOpen(static_cast(&file), file_read_callback, &error); + STD_TORCH_CHECK( + (gif_file != nullptr) && (error == D_GIF_SUCCEEDED), + "DGifOpen() failed - ", + error); + return probe_and_close(gif_file); +} + } // namespace facebook::torchcodec #endif // !TORCHCODEC_ENABLE_GIF diff --git a/src/torchcodec/_core/DecodeGif.h b/src/torchcodec/_core/DecodeGif.h index 7873761ab..623ab14d7 100644 --- a/src/torchcodec/_core/DecodeGif.h +++ b/src/torchcodec/_core/DecodeGif.h @@ -6,6 +6,9 @@ #pragma once +#include +#include + #include "StableABICompat.h" namespace facebook::torchcodec { @@ -20,4 +23,10 @@ FORCE_PUBLIC_VISIBILITY torch::stable::Tensor decode_gif( const torch::stable::Tensor& input, int64_t mode); +FORCE_PUBLIC_VISIBILITY std::vector probe_gif( + const torch::stable::Tensor& input); + +FORCE_PUBLIC_VISIBILITY std::vector probe_gif_from_file( + std::string filename); + } // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/Exif.h b/src/torchcodec/_core/Exif.h index 13c1724ef..843620d57 100644 --- a/src/torchcodec/_core/Exif.h +++ b/src/torchcodec/_core/Exif.h @@ -192,6 +192,15 @@ inline ExifOrientation fetch_exif_orientation( return exif_orientation; } +// Whether applying the orientation involves a 90 or 270 degree rotation, which +// swaps the image's width and height. +inline bool swaps_width_and_height(ExifOrientation orientation) { + return orientation == ExifOrientation::LeftTop || + orientation == ExifOrientation::RightTop || + orientation == ExifOrientation::RightBottom || + orientation == ExifOrientation::LeftBottom; +} + inline torch::stable::Tensor exif_orientation_transform( const torch::stable::Tensor& image, ExifOrientation orientation) { diff --git a/src/torchcodec/_core/ImageCommon.h b/src/torchcodec/_core/ImageCommon.h index b286d18e6..3f6ad7636 100644 --- a/src/torchcodec/_core/ImageCommon.h +++ b/src/torchcodec/_core/ImageCommon.h @@ -9,6 +9,8 @@ #include #include +#include + #include "StableABICompat.h" namespace facebook::torchcodec { @@ -81,6 +83,12 @@ inline bool should_return_rgb(ImageReadMode mode, bool has_alpha) { } } +// The probe_* functions return {width, height, num_channels, bit_depth, +// num_frames, orientation}: see ImageMetadata in decoders/_image_metadata.py. +// orientation is NO_ORIENTATION when the image has no orientation metadata. +using ImageProbeResult = std::vector; +constexpr int64_t NO_ORIENTATION = -1; + // Validates that `input` is a non-empty, 1-dimensional, uint8 CPU tensor of // encoded image bytes, and returns a contiguous version of it. inline torch::stable::Tensor validate_encoded_data( diff --git a/src/torchcodec/_core/image_custom_ops.cpp b/src/torchcodec/_core/image_custom_ops.cpp index b6d1b56ef..8f6e03686 100644 --- a/src/torchcodec/_core/image_custom_ops.cpp +++ b/src/torchcodec/_core/image_custom_ops.cpp @@ -104,6 +104,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(torchcodec_ns, m) { m.def( "encode_jpeg_to_file_like(Tensor img, int file_like_context, int quality) -> ()"); m.def("encode_jpeg_to_tensor_cuda(Tensor img, int quality) -> Tensor"); + m.def("probe_gif(Tensor input) -> int[]"); + m.def("probe_avif(Tensor input) -> int[]"); + m.def("probe_avif_from_file(str filename) -> int[]"); + m.def("probe_gif_from_file(str filename) -> int[]"); } STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) { @@ -114,6 +118,13 @@ STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) { m.impl("decode_avif", TORCH_BOX(&decode_avif)); m.impl("encode_png_to_file", TORCH_BOX(&encode_png_to_file)); m.impl("encode_png_to_file_like", TORCH_BOX(&encode_png_to_file_like)); + m.impl("probe_gif", TORCH_BOX(&probe_gif)); + m.impl("probe_avif", TORCH_BOX(&probe_avif)); +} + +STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, BackendSelect, m) { + m.impl("probe_avif_from_file", TORCH_BOX(&probe_avif_from_file)); + m.impl("probe_gif_from_file", TORCH_BOX(&probe_gif_from_file)); } // The JPEG encoders are registered as CompositeExplicitAutograd (rather than diff --git a/src/torchcodec/_core/ops.py b/src/torchcodec/_core/ops.py index 7d68de275..4212fbd9b 100644 --- a/src/torchcodec/_core/ops.py +++ b/src/torchcodec/_core/ops.py @@ -48,6 +48,10 @@ def expose_ffmpeg_dlls(): # noqa: F811 encode_jpeg_to_file = torch.ops.torchcodec_ns.encode_jpeg_to_file.default encode_jpeg_to_file_like = torch.ops.torchcodec_ns.encode_jpeg_to_file_like.default encode_jpeg_to_tensor_cuda = torch.ops.torchcodec_ns.encode_jpeg_to_tensor_cuda.default +probe_gif = torch.ops.torchcodec_ns.probe_gif.default +probe_gif_from_file = torch.ops.torchcodec_ns.probe_gif_from_file.default +probe_avif = torch.ops.torchcodec_ns.probe_avif.default +probe_avif_from_file = torch.ops.torchcodec_ns.probe_avif_from_file.default def get_decode_heic(): diff --git a/src/torchcodec/decoders/__init__.py b/src/torchcodec/decoders/__init__.py index 160f81a54..d8507d145 100644 --- a/src/torchcodec/decoders/__init__.py +++ b/src/torchcodec/decoders/__init__.py @@ -40,6 +40,7 @@ decode_webp, ImageReadMode, ) +from ._image_metadata import get_image_metadata, ImageMetadata # noqa from ._video_decoder import CpuFallbackStatus, VideoDecoder # noqa from ._wav_decoder import WavDecoder # noqa diff --git a/src/torchcodec/decoders/_image_decoders.py b/src/torchcodec/decoders/_image_decoders.py index e116d9ad5..0a5904475 100644 --- a/src/torchcodec/decoders/_image_decoders.py +++ b/src/torchcodec/decoders/_image_decoders.py @@ -536,7 +536,10 @@ def _detect_image_format(data: torch.Tensor) -> str: # Sniff the codec from the leading "magic" bytes of the encoded data. # This used to be implemented in C++ in torchvision, but benchmarks show # this is negligible in Python - header = bytes(data[:64].tolist()) + return _detect_image_format_from_header(bytes(data[:64].tolist())) + + +def _detect_image_format_from_header(header: bytes) -> str: if header[:3] == b"\xff\xd8\xff": return "jpeg" if header[:8] == b"\x89PNG\r\n\x1a\n": diff --git a/src/torchcodec/decoders/_image_metadata.py b/src/torchcodec/decoders/_image_metadata.py new file mode 100644 index 000000000..0a8c82fd4 --- /dev/null +++ b/src/torchcodec/decoders/_image_metadata.py @@ -0,0 +1,836 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import struct +from dataclasses import dataclass +from pathlib import Path + +import torch +from torchcodec._core import ops +from torchcodec.decoders._image_decoders import ( + _detect_image_format_from_header, + _source_to_tensor, +) + + +@dataclass(frozen=True) +class ImageMetadata: + """Metadata of an encoded image, as returned by :func:`get_image_metadata`. + + Every attribute describes the tensor that :func:`decode_image` (or the + matching format-specific decoder) returns for this image. + """ + + format: str + """The image format: one of ``"jpeg"``, ``"png"``, ``"webp"``, ``"gif"``, + ``"avif"`` or ``"heic"``.""" + width: int + """Width of the decoded image, in pixels. This is the **post-rotation** + width: for 90 or -90 degree :attr:`rotation`, width and height are swapped + compared to the dimensions stored in the file.""" + height: int + """Height of the decoded image, in pixels. Like :attr:`width`, this is the + **post-rotation** height.""" + num_channels: int + """Number of channels of the decoded image when decoding with + ``mode="UNCHANGED"``.""" + bit_depth: int + """Number of bits per channel in the encoded image. When greater than 8, + decoding with ``output_dtype="auto"`` returns a ``torch.uint16`` tensor.""" + num_frames: int + """Number of frames (or images) in the file: 1 for still images, and the + size of the leading ``N`` dimension of the decoded tensor for animated or + multi-image files.""" + rotation: float | None + """Rotation angle in degrees (counter-clockwise: one of 0, 90, -90 or 180) + from the orientation metadata (EXIF orientation, or the AVIF/HEIC rotation + and mirroring properties), or None if the image has no orientation + metadata. The decoders always apply it, so the returned images are in the + correct orientation. This uses the same convention as + :attr:`VideoStreamMetadata.rotation + `. + + The orientation metadata can also specify a mirroring, which the decoders + apply too but isn't reported here, since it doesn't change the image + dimensions. In that case, this is the rotation applied after mirroring the + image left to right.""" + + +class _ByteSource: + # Random access over the encoded bytes. Paths are read lazily, so probing a + # large file usually only reads its first few KB. + + def __init__(self, source): + self._file = None + if isinstance(source, (str, Path)): + self._file = open(source, "rb") + self._buf = None + self.size = Path(source).stat().st_size + elif isinstance(source, (bytes, bytearray)): + self._buf = memoryview(source) + self.size = len(source) + elif isinstance(source, torch.Tensor): + if ( + source.dtype != torch.uint8 + or source.ndim != 1 + or source.device.type != "cpu" + ): + raise ValueError( + "Input tensor must be a 1-dimensional uint8 CPU tensor, got " + f"dtype={source.dtype}, ndim={source.ndim}, device={source.device}." + ) + self._buf = memoryview(source.contiguous().numpy()) + self.size = len(self._buf) + else: + raise TypeError( + f"Unknown source type: {type(source)}. " + "Supported types are str, Path, bytes and torch.Tensor." + ) + + def read(self, offset: int, size: int) -> bytes: + # Returns fewer than `size` bytes when reading past the end. + if self._buf is not None: + return bytes(self._buf[offset : offset + size]) + assert self._file is not None + self._file.seek(offset) + return self._file.read(size) + + def read_exact(self, offset: int, size: int) -> bytes: + data = self.read(offset, size) + if len(data) != size: + raise RuntimeError("Image is truncated or corrupted.") + return data + + def close(self): + if self._file is not None: + self._file.close() + + +def _swaps_width_and_height(orientation: int | None) -> bool: + # EXIF orientations 5 to 8 involve a 90 or 270 degree rotation. + return orientation is not None and 5 <= orientation <= 8 + + +_EXIF_INCORRECT_TAG = 0xFFFF + + +def _exif_orientation(tiff: bytes) -> int | None: + # Port of fetch_exif_orientation() from Exif.h, quirks included, so that we + # report exactly the orientation the decoders apply. `tiff` starts at the + # TIFF header (i.e. after any "Exif\0\0" prefix). Returns the EXIF + # orientation in [1, 8], or None if there's none (or it's invalid, in which + # case the decoders don't apply anything either). + n = len(tiff) + byte_order = "<" if tiff[:2] == b"II" else ">" + + def u16(offset): + if offset + 1 >= n: + return _EXIF_INCORRECT_TAG + return struct.unpack_from(byte_order + "H", tiff, offset)[0] + + if u16(2) != 0x2A: + return None + offset = ( + _EXIF_INCORRECT_TAG + if 7 >= n + else struct.unpack_from(byte_order + "I", tiff, 4)[0] + ) + num_entries = u16(offset) + offset += 2 + for _ in range(num_entries): + tag = u16(offset) + if tag == _EXIF_INCORRECT_TAG: + break + if tag == 0x0112: + orientation = u16(offset + 8) + return orientation if 1 <= orientation <= 8 else None + offset += 12 + return None + + +# ---------------------------------------------------------------------------- +# JPEG +# ---------------------------------------------------------------------------- + +# SOFn markers that libjpeg-turbo can decode. The other SOFn markers +# (hierarchical, and the reserved JPG marker) make jpeg_read_header() fail. +_JPEG_SUPPORTED_SOF = frozenset((0xC0, 0xC1, 0xC2, 0xC3, 0xC9, 0xCA, 0xCB)) +_JPEG_UNSUPPORTED_SOF = frozenset((0xC5, 0xC6, 0xC7, 0xC8, 0xCD, 0xCE, 0xCF)) + + +def _probe_jpeg(src: _ByteSource) -> tuple: + # Walks the markers up to the first SOS, like jpeg_read_header(). The + # decoder takes the EXIF orientation from the first APP1 segment, whatever + # it contains. + pos = 2 # skip SOI, already checked by the format sniffer + sof = None + orientation = None + seen_app1 = False + while True: + chunk = src.read(pos, 4) + if len(chunk) < 2: + raise RuntimeError("JPEG is truncated: no SOS marker found.") + if chunk[0] != 0xFF: + # libjpeg skips garbage bytes between markers. + window = src.read(pos, 4096) + idx = window.find(b"\xff") + pos += idx if idx > 0 else len(window) + continue + marker = chunk[1] + if marker == 0xFF: # fill byte + pos += 1 + continue + if marker == 0xDA: # SOS + break + if marker == 0xD9: # EOI + raise RuntimeError("JPEG has no image data.") + if marker == 0x01 or 0xD0 <= marker <= 0xD7: # standalone markers + pos += 2 + continue + if len(chunk) < 4: + raise RuntimeError("JPEG is truncated.") + segment_length = (chunk[2] << 8) | chunk[3] + if segment_length < 2: + raise RuntimeError("Invalid JPEG marker segment length.") + if marker in _JPEG_SUPPORTED_SOF: + if sof is not None: + raise RuntimeError("JPEG has more than one SOF marker.") + precision, height, width, num_components = struct.unpack( + ">BHHB", src.read_exact(pos + 4, 6) + ) + if segment_length != 8 + 3 * num_components: + raise RuntimeError("Invalid JPEG SOF marker length.") + if height == 0 or width == 0 or num_components == 0: + raise RuntimeError("Empty JPEG image (DNL not supported).") + sof = (precision, height, width, num_components) + elif marker in _JPEG_UNSUPPORTED_SOF: + raise RuntimeError(f"Unsupported JPEG process: SOF type 0x{marker:02x}.") + elif marker == 0xE1 and not seen_app1: + seen_app1 = True + payload = src.read_exact(pos + 4, segment_length - 2) + if len(payload) > 6: + orientation = _exif_orientation(payload[6:]) + pos += 2 + segment_length + + if sof is None: + raise RuntimeError("JPEG has no SOF marker before SOS.") + precision, height, width, num_components = sof + if _swaps_width_and_height(orientation): + width, height = height, width + return width, height, num_components, precision, 1, orientation + + +# ---------------------------------------------------------------------------- +# PNG +# ---------------------------------------------------------------------------- + +_PNG_VALID_BIT_DEPTHS = { + 0: (1, 2, 4, 8, 16), # gray + 2: (8, 16), # RGB + 3: (1, 2, 4, 8), # palette + 4: (8, 16), # gray + alpha + 6: (8, 16), # RGBA +} + + +def _probe_png(src: _ByteSource) -> tuple: + # Reads IHDR and walks the chunks up to the first IDAT, which is what + # png_read_info() sees. In particular, an eXIf chunk after IDAT is ignored + # by the decoder. + ihdr = src.read_exact(8, 8 + 13) + length, chunk_type = struct.unpack_from(">I4s", ihdr) + if chunk_type != b"IHDR" or length != 13: + raise RuntimeError("PNG is missing its IHDR chunk.") + width, height, bit_depth, color_type = struct.unpack_from(">IIBB", ihdr, 8) + if not (0 < width < 2**31 and 0 < height < 2**31): + raise RuntimeError(f"Invalid PNG dimensions: {width}x{height}.") + if bit_depth not in _PNG_VALID_BIT_DEPTHS.get(color_type, ()): + raise RuntimeError( + f"Invalid PNG color type ({color_type}) and bit depth ({bit_depth})." + ) + + pos = 8 + 12 + 13 + num_palette = None + has_trns = False + exif = None + while True: + length, chunk_type = struct.unpack(">I4s", src.read_exact(pos, 8)) + if chunk_type == b"IDAT": + break + if chunk_type == b"IEND": + raise RuntimeError("PNG has no IDAT chunk.") + if chunk_type == b"PLTE": + num_palette = length // 3 + elif chunk_type == b"tRNS" and not has_trns: + # libpng ignores (with a warning) tRNS chunks that are invalid for + # the color type. + if color_type == 3: + has_trns = num_palette is not None and 0 < length <= num_palette + elif color_type == 0: + has_trns = length == 2 + elif color_type == 2: + has_trns = length == 6 + elif chunk_type == b"eXIf" and exif is None and length >= 2: + payload = src.read_exact(pos + 8, length) + if payload[0] == payload[1] and payload[0] in b"MI": + exif = payload + pos += 12 + length + if color_type == 3 and num_palette is None: + raise RuntimeError("PNG palette image has no PLTE chunk before IDAT.") + + orientation = _exif_orientation(exif) if exif is not None else None + if _swaps_width_and_height(orientation): + width, height = height, width + if color_type == 3: + num_channels = 4 if has_trns else 3 + bit_depth = 8 # palette entries are 8-bit RGB(A) + else: + num_channels = {0: 1, 2: 3, 4: 2, 6: 4}[color_type] + return width, height, num_channels, bit_depth, 1, orientation + + +# ---------------------------------------------------------------------------- +# WebP +# ---------------------------------------------------------------------------- + +_WEBP_ANIMATION_FLAG = 0x02 +_WEBP_EXIF_FLAG = 0x08 +_WEBP_ALPHA_FLAG = 0x10 + + +def _vp8_dimensions(src: _ByteSource, payload_pos: int) -> tuple[int, int]: + header = src.read_exact(payload_pos, 10) + if header[3:6] != b"\x9d\x01\x2a": + raise RuntimeError("Invalid VP8 bitstream signature.") + width, height = struct.unpack_from(" tuple[int, int, bool]: + header = src.read_exact(payload_pos, 5) + if header[0] != 0x2F: + raise RuntimeError("Invalid VP8L bitstream signature.") + (bits,) = struct.unpack_from("> 29 != 0: + raise RuntimeError("Invalid VP8L version.") + width = (bits & 0x3FFF) + 1 + height = ((bits >> 14) & 0x3FFF) + 1 + has_alpha = bool((bits >> 28) & 1) + return width, height, has_alpha + + +def _probe_webp(src: _ByteSource) -> tuple: + # Mirrors WebPGetFeatures() (dimensions and alpha), WebPAnimDecoder (frame + # count and canvas size) and WebPDemux (the EXIF chunk, which libwebpdemux + # only exposes when the VP8X EXIF flag is set). + first_type, first_size = struct.unpack("<4sI", src.read_exact(12, 8)) + orientation = None + if first_type == b"VP8 ": + width, height = _vp8_dimensions(src, 20) + has_alpha = False + num_frames = 1 + elif first_type == b"VP8L": + width, height, has_alpha = _vp8l_info(src, 20) + num_frames = 1 + elif first_type == b"VP8X": + vp8x = src.read_exact(20, 10) + flags = vp8x[0] + canvas_width = int.from_bytes(vp8x[4:7], "little") + 1 + canvas_height = int.from_bytes(vp8x[7:10], "little") + 1 + is_animated = bool(flags & _WEBP_ANIMATION_FLAG) + + num_anmf = 0 + seen_alph = False + bitstream = None # (width, height, vp8l_alpha_or_None) of first image + exif = None + pos = 20 + first_size + (first_size & 1) + while True: + header = src.read(pos, 8) + if len(header) < 8: + break + chunk_type, chunk_size = struct.unpack("<4sI", header) + if chunk_type == b"ANMF": + num_anmf += 1 + elif not is_animated and bitstream is None: + if chunk_type == b"ALPH": + seen_alph = True + elif chunk_type == b"VP8 ": + bitstream = (*_vp8_dimensions(src, pos + 8), None) + elif chunk_type == b"VP8L": + bitstream = _vp8l_info(src, pos + 8) + if chunk_type == b"EXIF" and exif is None and flags & _WEBP_EXIF_FLAG: + exif = src.read_exact(pos + 8, chunk_size) + pos += 8 + chunk_size + (chunk_size & 1) + + if exif is not None: + if exif[:6] == b"Exif\x00\x00": + exif = exif[6:] + if len(exif) > 0: + orientation = _exif_orientation(exif) + + if is_animated: + if num_anmf == 0: + raise RuntimeError("Animated WebP has no frames.") + width, height = canvas_width, canvas_height + has_alpha = bool(flags & _WEBP_ALPHA_FLAG) + num_frames = num_anmf + else: + if bitstream is None: + raise RuntimeError("WebP has no VP8 or VP8L bitstream.") + width, height, vp8l_alpha = bitstream + if (width, height) != (canvas_width, canvas_height): + raise RuntimeError("WebP canvas and bitstream dimensions differ.") + if vp8l_alpha is None: + has_alpha = bool(flags & _WEBP_ALPHA_FLAG) or seen_alph + else: + has_alpha = vp8l_alpha or seen_alph + num_frames = 1 + else: + raise RuntimeError(f"Unexpected first WebP chunk: {first_type!r}.") + + if _swaps_width_and_height(orientation): + width, height = height, width + return width, height, 4 if has_alpha else 3, 8, num_frames, orientation + + +# ---------------------------------------------------------------------------- +# HEIC (ISOBMFF / HEIF) +# ---------------------------------------------------------------------------- + + +def _iter_boxes(data: bytes, start: int, end: int): + # Yields (box_type, payload_start, payload_end) for the boxes in + # data[start:end]. + pos = start + while pos + 8 <= end: + size, box_type = struct.unpack_from(">I4s", data, pos) + header_size = 8 + if size == 1: + if pos + 16 > end: + raise RuntimeError("ISOBMFF box is truncated.") + (size,) = struct.unpack_from(">Q", data, pos + 8) + header_size = 16 + elif size == 0: + size = end - pos + if box_type == b"uuid": + header_size += 16 + if size < header_size or pos + size > end: + raise RuntimeError("Invalid or truncated ISOBMFF box.") + yield box_type, pos + header_size, pos + size + pos += size + + +def _iter_top_level_boxes(src: _ByteSource): + # Like _iter_boxes, over the whole file, reading only box headers and the + # payloads we ask for. Yields (box_type, payload_start, payload_end). + pos = 0 + while True: + header = src.read(pos, 16) + if len(header) < 8: + return + size, box_type = struct.unpack_from(">I4s", header) + header_size = 8 + if size == 1: + if len(header) < 16: + raise RuntimeError("ISOBMFF box is truncated.") + (size,) = struct.unpack_from(">Q", header, 8) + header_size = 16 + elif size == 0: # the box extends to the end of the file + size = src.size - pos + if size < header_size: + raise RuntimeError("Invalid ISOBMFF box size.") + yield box_type, pos + header_size, pos + size + pos += size + + +class _HeifMeta: + # The parts of a 'meta' box we need: items, their references, and their + # associated properties (in association order). + + def __init__(self, data: bytes): + self.primary_id = None + self.item_types: dict[int, bytes] = {} + self.hidden: set[int] = set() + # refs[ref_type][from_id] = [to_ids...], in file order + self.refs: dict[bytes, dict[int, list[int]]] = {} + self._refs_to: dict[bytes, dict[int, list[int]]] = {} + self._properties: list[tuple[bytes, bytes]] = [] + self._associations: dict[int, list[int]] = {} + + for box_type, start, end in _iter_boxes(data, 4, len(data)): + if box_type == b"pitm": + version = data[start] + fmt = ">H" if version == 0 else ">I" + (self.primary_id,) = struct.unpack_from(fmt, data, start + 4) + elif box_type == b"iinf": + version = data[start] + entries_start = start + 4 + (2 if version == 0 else 4) + for infe_type, s, e in _iter_boxes(data, entries_start, end): + if infe_type == b"infe": + self._parse_infe(data, s, e) + elif box_type == b"iref": + version = data[start] + id_fmt, id_size = (">H", 2) if version == 0 else (">I", 4) + for ref_type, s, e in _iter_boxes(data, start + 4, end): + (from_id,) = struct.unpack_from(id_fmt, data, s) + (count,) = struct.unpack_from(">H", data, s + id_size) + first = s + id_size + 2 + if first + count * id_size > e: + raise RuntimeError("Invalid 'iref' box.") + to_ids = [ + struct.unpack_from(id_fmt, data, first + i * id_size)[0] + for i in range(count) + ] + self.refs.setdefault(ref_type, {}).setdefault(from_id, []).extend( + to_ids + ) + for to_id in to_ids: + self._refs_to.setdefault(ref_type, {}).setdefault( + to_id, [] + ).append(from_id) + elif box_type == b"iprp": + for iprp_type, s, e in _iter_boxes(data, start, end): + if iprp_type == b"ipco": + self._properties = [ + (t, data[ps:pe]) for t, ps, pe in _iter_boxes(data, s, e) + ] + elif iprp_type == b"ipma": + self._parse_ipma(data, s, e) + + def _parse_infe(self, data, start, end): + version = data[start] + if version < 2: + return + flags = int.from_bytes(data[start + 1 : start + 4], "big") + if version == 2: + (item_id,) = struct.unpack_from(">H", data, start + 4) + item_type = data[start + 8 : start + 12] + else: + (item_id,) = struct.unpack_from(">I", data, start + 4) + item_type = data[start + 10 : start + 14] + self.item_types[item_id] = bytes(item_type) + if flags & 1: + self.hidden.add(item_id) + + def _parse_ipma(self, data, start, end): + version = data[start] + flags = int.from_bytes(data[start + 1 : start + 4], "big") + (entry_count,) = struct.unpack_from(">I", data, start + 4) + pos = start + 8 + for _ in range(entry_count): + if version < 1: + (item_id,) = struct.unpack_from(">H", data, pos) + pos += 2 + else: + (item_id,) = struct.unpack_from(">I", data, pos) + pos += 4 + num_associations = data[pos] + pos += 1 + indices = self._associations.setdefault(item_id, []) + for _ in range(num_associations): + if flags & 1: + (value,) = struct.unpack_from(">H", data, pos) + pos += 2 + indices.append(value & 0x7FFF) + else: + indices.append(data[pos] & 0x7F) + pos += 1 + + def properties(self, item_id: int) -> list[tuple[bytes, bytes]]: + # (property_type, property_payload) pairs, in association order. + # Property indices are 1-based; 0 means "no property". + return [ + self._properties[i - 1] + for i in self._associations.get(item_id, ()) + if 0 < i <= len(self._properties) + ] + + def find_property(self, item_id: int, property_type: bytes) -> bytes | None: + for t, payload in self.properties(item_id): + if t == property_type: + return payload + return None + + def refs_from(self, ref_type: bytes, from_id: int) -> list[int]: + return self.refs.get(ref_type, {}).get(from_id, []) + + def refs_to(self, ref_type: bytes, to_id: int) -> list[int]: + return self._refs_to.get(ref_type, {}).get(to_id, []) + + +_HEIC_ALPHA_URNS = ( + b"urn:mpeg:avc:2015:auxid:1", + b"urn:mpeg:hevc:2015:auxid:1", + b"urn:mpeg:mpegB:cicp:systems:auxiliary:alpha", +) + + +def _aux_type(auxc_payload: bytes | None) -> bytes | None: + # auxC is a FullBox holding a null-terminated URN. + if auxc_payload is None: + return None + return auxc_payload[4:].split(b"\x00", 1)[0] + + +def _ispe_size(meta: _HeifMeta, item_id: int) -> tuple[int, int]: + ispe = meta.find_property(item_id, b"ispe") + if ispe is None or len(ispe) < 12: + raise RuntimeError(f"Image item {item_id} has no valid 'ispe' property.") + return struct.unpack_from(">II", ispe, 4) + + +def _read_meta_box(src: _ByteSource) -> bytes: + # Returns the payload of the top-level 'meta' box, skipping everything else + # (e.g. 'mdat') without reading it. + for box_type, start, end in _iter_top_level_boxes(src): + if box_type == b"meta": + return src.read_exact(start, end - start) + raise RuntimeError("HEIC has no 'meta' box.") + + +def _av1c_depth(av1c: bytes) -> int: + flags = av1c[2] + if flags & 0x20: # twelve_bit + return 12 + if flags & 0x40: # high_bitdepth + return 10 + return 8 + + +def _flip_h(grid): + return tuple(row[::-1] for row in grid) + + +def _flip_v(grid): + return grid[::-1] + + +def _transpose(grid): + return tuple(zip(*grid)) + + +def _rotate_ccw(grid): + return _flip_v(_transpose(grid)) + + +# libheif applies the irot and imir properties in association order, so we +# compose them by applying them to a small grid of distinct values, and look up +# which EXIF orientation produces the same result. The transforms below are the +# ones of exif_orientation_transform() in Exif.h. +_ORIENTATION_REFERENCE_GRID = ((0, 1, 2), (3, 4, 5)) +_ORIENTATION_FROM_GRID = { + transform(_ORIENTATION_REFERENCE_GRID): orientation + for orientation, transform in { + 1: lambda g: g, + 2: _flip_h, + 3: lambda g: _flip_h(_flip_v(g)), + 4: _flip_v, + 5: _transpose, + 6: lambda g: _flip_h(_transpose(g)), + 7: lambda g: _flip_h(_flip_v(_transpose(g))), + 8: lambda g: _flip_v(_transpose(g)), + }.items() +} + + +_HEIC_IMAGE_ITEM_TYPES = frozenset( + ( + b"jpeg", + b"hvc1", + b"av01", + b"vvc1", + b"avc1", + b"unci", + b"j2k1", + b"mski", + b"grid", + b"iovl", + b"iden", + b"tili", + ) +) +_HEIC_DERIVED_ITEM_TYPES = frozenset((b"grid", b"iovl", b"iden")) + + +def _probe_heic(src: _ByteSource) -> tuple: + # Mirrors libheif: the top-level images are the non-hidden image items that + # aren't thumbnails or auxiliary images, in item-ID order. The decoder uses + # the first one's geometry (and requires every other one to match). + meta = _HeifMeta(_read_meta_box(src)) + + image_ids = sorted( + i for i, t in meta.item_types.items() if t in _HEIC_IMAGE_ITEM_TYPES + ) + if meta.primary_id not in image_ids or meta.primary_id in meta.hidden: + raise RuntimeError( + "'pitm' box references an unsupported or non-existing image." + ) + top_level = [ + i + for i in image_ids + if i not in meta.hidden + and not meta.refs_from(b"thmb", i) + and not meta.refs_from(b"auxl", i) + ] + if not top_level: + raise RuntimeError("HEIC file should contain at least one top-level image.") + item_id = top_level[0] + + width, height = _ispe_size(meta, item_id) + grid = _ORIENTATION_REFERENCE_GRID + has_orientation = False + for prop_type, payload in meta.properties(item_id): + if prop_type == b"clap": + w_num, w_den, h_num, h_den = struct.unpack_from(">iIiI", payload) + if w_den == 0 or h_den == 0: + raise RuntimeError("Invalid 'clap' property.") + width = (w_num + w_den // 2) // w_den + height = (h_num + h_den // 2) // h_den + elif prop_type == b"irot": + has_orientation = True + angle = payload[0] & 0x03 # counter-clockwise quarter turns + for _ in range(angle): + grid = _rotate_ccw(grid) + if angle in (1, 3): + width, height = height, width + elif prop_type == b"imir": + has_orientation = True + # Axis 0 exchanges the top and bottom, axis 1 the left and right. + grid = _flip_v(grid) if payload[0] & 0x01 == 0 else _flip_h(grid) + orientation = _ORIENTATION_FROM_GRID[grid] if has_orientation else None + + # Luma bit depth, from the first non-derived image down the 'dimg' chain. + coded_id = item_id + for _ in range(len(meta.item_types) + 1): + if meta.item_types.get(coded_id) not in _HEIC_DERIVED_ITEM_TYPES: + break + children = meta.refs_from(b"dimg", coded_id) + if not children: + raise RuntimeError("Derived HEIC image has no input image.") + coded_id = children[0] + coded_type = meta.item_types.get(coded_id) + if coded_type == b"hvc1": + hvcc = meta.find_property(coded_id, b"hvcC") + if hvcc is None or len(hvcc) < 18: + raise RuntimeError("No hvcC property in hvc1 type image.") + bit_depth = (hvcc[17] & 0x07) + 8 + elif coded_type == b"av01": + av1c = meta.find_property(coded_id, b"av1C") + if av1c is None: + raise RuntimeError("No av1C property in av01 type image.") + bit_depth = _av1c_depth(av1c) + elif coded_type == b"jpeg": + bit_depth = 8 + else: + raise RuntimeError(f"Unsupported HEIC image item type: {coded_type!r}.") + + def has_alpha_aux(image_id): + return any( + _aux_type(meta.find_property(i, b"auxC")) in _HEIC_ALPHA_URNS + for i in meta.refs_to(b"auxl", image_id) + ) + + has_alpha = has_alpha_aux(item_id) or ( + meta.item_types[item_id] == b"grid" + and any(has_alpha_aux(t) for t in meta.refs_from(b"dimg", item_id)) + ) + return width, height, 4 if has_alpha else 3, bit_depth, len(top_level), orientation + + +_PY_PROBES = { + "jpeg": _probe_jpeg, + "png": _probe_png, + "webp": _probe_webp, + "heic": _probe_heic, +} + +# GIF and AVIF are probed in C++, with giflib and libavif. Both take a path +# directly, so that they can read the file lazily. +_CPP_PROBES = { + "gif": (ops.probe_gif, ops.probe_gif_from_file), + "avif": (ops.probe_avif, ops.probe_avif_from_file), +} +_CPP_NO_ORIENTATION = -1 + +# Rotation of each EXIF orientation, when decomposed into an optional +# left-to-right mirroring followed by a rotation. +_ROTATION_FROM_EXIF_ORIENTATION = { + 1: 0.0, + 2: 0.0, + 3: 180.0, + 4: 180.0, + 5: 90.0, + 6: -90.0, + 7: -90.0, + 8: 90.0, +} + + +def get_image_metadata(source: str | Path | bytes | torch.Tensor) -> ImageMetadata: + """Get the metadata of an encoded image, without decoding it. + + The format is detected from the encoded data, like in + :func:`~torchcodec.decoders.decode_image`. Supported formats are JPEG, PNG, + WebP, GIF, AVIF and HEIC. Only the image headers are parsed, so this is + much faster than decoding. When ``source`` is a path, only the parts of the + file that hold the headers are read (except for GIF, which has no frame + count in its header). + + See :class:`ImageMetadata` for the returned attributes, which all describe + the output of the decoders. + + Example: + + .. code-block:: python + + from torchcodec.decoders import get_image_metadata + + metadata = get_image_metadata("image.jpg") + print(metadata.width, metadata.height) + + Args: + source (str, ``pathlib.Path``, bytes, or ``torch.Tensor``): + The encoded image data: a path (``str`` or ``pathlib.Path``), a + ``bytes`` object, or a 1-D uint8 ``torch.Tensor`` of the raw encoded + bytes. + + Returns: + ImageMetadata: The image metadata. + """ + src = _ByteSource(source) + try: + fmt = _detect_image_format_from_header(src.read(0, 64)) + if fmt in _PY_PROBES: + probe_result = _PY_PROBES[fmt](src) + else: + probe, probe_from_file = _CPP_PROBES[fmt] + if isinstance(source, (str, Path)): + probe_result = probe_from_file(str(source)) + else: + probe_result = probe(_source_to_tensor(source)) + if probe_result[5] == _CPP_NO_ORIENTATION: + probe_result = (*probe_result[:5], None) + except (struct.error, IndexError): + raise RuntimeError("Image is truncated or corrupted.") from None + finally: + src.close() + width, height, num_channels, bit_depth, num_frames, orientation = probe_result + return ImageMetadata( + format=fmt, + width=width, + height=height, + num_channels=num_channels, + bit_depth=bit_depth, + num_frames=num_frames, + rotation=( + None + if orientation is None + else _ROTATION_FROM_EXIF_ORIENTATION[orientation] + ), + ) diff --git a/test/test_decoders.py b/test/test_decoders.py index 3af00fdf4..b5f8fe3ff 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -45,7 +45,9 @@ Demuxer, FrameIndex, get_container_metadata, + get_image_metadata, get_nvdec_cache_capacity, + ImageMetadata, ImageReadMode, Packet, RawAudioSamples, @@ -7963,3 +7965,287 @@ def boom(path): finally: # Don't leak the (failure-poisoned) cache to other tests. iru.load_heic_library.cache_clear() + + +class TestGetImageMetadata: + @staticmethod + def _assert_matches_decode(metadata, source): + # Every metadata field is defined by what the decoders output. + decoded = decode_image(source, mode="UNCHANGED", output_dtype="auto") + num_frames = decoded.shape[0] if decoded.ndim == 4 else 1 + assert metadata.num_frames == num_frames + assert metadata.num_channels == decoded.shape[-3] + assert metadata.height == decoded.shape[-2] + assert metadata.width == decoded.shape[-1] + if metadata.format in ("png", "avif", "heic"): + assert (metadata.bit_depth > 8) == (decoded.dtype == torch.uint16) + + @pytest.mark.parametrize( + "make_source", + ( + pytest.param(lambda a: a.path, id="path"), + pytest.param(lambda a: str(a.path), id="str"), + pytest.param(lambda a: a.path.read_bytes(), id="bytes"), + pytest.param( + lambda a: torch.frombuffer( + bytearray(a.path.read_bytes()), dtype=torch.uint8 + ), + id="tensor", + ), + ), + ) + @pytest.mark.parametrize( + "asset, fmt, num_channels, bit_depth, num_frames", + ( + pytest.param(GRADIENT_JPEG, "jpeg", 3, 8, 1, marks=pytest.mark.needs_jpeg), + pytest.param(GRAYSCALE_JPEG, "jpeg", 1, 8, 1, marks=pytest.mark.needs_jpeg), + pytest.param(CMYK_JPEG, "jpeg", 4, 8, 1, marks=pytest.mark.needs_jpeg), + pytest.param(GRADIENT_PNG, "png", 3, 8, 1, marks=pytest.mark.needs_png), + pytest.param(GRAYSCALE_PNG, "png", 1, 8, 1, marks=pytest.mark.needs_png), + pytest.param(RGBA_PNG, "png", 4, 8, 1, marks=pytest.mark.needs_png), + pytest.param( + GRAYSCALE_ALPHA_PNG, "png", 2, 8, 1, marks=pytest.mark.needs_png + ), + pytest.param( + GRADIENT_16BIT_PNG, "png", 3, 16, 1, marks=pytest.mark.needs_png + ), + pytest.param( + GRAYSCALE_16BIT_PNG, "png", 1, 16, 1, marks=pytest.mark.needs_png + ), + pytest.param( + GRADIENT_INTERLACED_PNG, "png", 3, 8, 1, marks=pytest.mark.needs_png + ), + pytest.param(GRADIENT_WEBP, "webp", 3, 8, 1, marks=pytest.mark.needs_webp), + pytest.param(RGBA_WEBP, "webp", 4, 8, 1, marks=pytest.mark.needs_webp), + pytest.param(GRADIENT_GIF, "gif", 3, 8, 1), + pytest.param(ANIMATED_GIF, "gif", 4, 8, 4), + pytest.param(TRANSPARENT_GIF, "gif", 4, 8, 1), + pytest.param(FRAME_EXCEEDS_SCREEN_GIF, "gif", 4, 8, 1), + pytest.param(GRADIENT_AVIF, "avif", 3, 8, 1, marks=pytest.mark.needs_avif), + pytest.param(RGBA_AVIF, "avif", 4, 8, 1, marks=pytest.mark.needs_avif), + pytest.param( + GRADIENT_10BIT_AVIF, "avif", 3, 10, 1, marks=pytest.mark.needs_avif + ), + pytest.param( + GRADIENT_12BIT_AVIF, "avif", 3, 12, 1, marks=pytest.mark.needs_avif + ), + pytest.param(GRADIENT_HEIC, "heic", 3, 8, 1, marks=pytest.mark.needs_heic), + pytest.param(RGBA_HEIC, "heic", 4, 8, 1, marks=pytest.mark.needs_heic), + pytest.param( + GRADIENT_10BIT_HEIC, "heic", 3, 10, 1, marks=pytest.mark.needs_heic + ), + pytest.param( + GRADIENT_ROTATED_HEIC, "heic", 3, 8, 1, marks=pytest.mark.needs_heic + ), + pytest.param( + GRADIENT_MIRRORED_HEIC, "heic", 3, 8, 1, marks=pytest.mark.needs_heic + ), + pytest.param(ANIMATED_HEIC, "heic", 3, 8, 3, marks=pytest.mark.needs_heic), + ), + ) + def test_matches_decode( + self, asset, fmt, num_channels, bit_depth, num_frames, make_source + ): + metadata = get_image_metadata(make_source(asset)) + rotation = { + GRADIENT_ROTATED_HEIC.filename: -90.0, + GRADIENT_MIRRORED_HEIC.filename: 0.0, + }.get(asset.filename) + assert metadata == ImageMetadata( + format=fmt, + width=asset.width, + height=asset.height, + num_channels=num_channels, + bit_depth=bit_depth, + num_frames=num_frames, + rotation=rotation, + ) + self._assert_matches_decode(metadata, make_source(asset)) + + @pytest.mark.parametrize( + "fmt, ext, save_kwargs", + ( + pytest.param("JPEG", "jpg", {}, marks=pytest.mark.needs_jpeg, id="jpeg"), + pytest.param("PNG", "png", {}, marks=pytest.mark.needs_png, id="png"), + pytest.param("WEBP", "webp", {}, marks=pytest.mark.needs_webp, id="webp"), + pytest.param("AVIF", "avif", {}, marks=pytest.mark.needs_avif, id="avif"), + ), + ) + @pytest.mark.parametrize("orientation", (0, 1, 2, 3, 4, 5, 6, 7, 8)) + def test_exif_orientation(self, tmp_path, orientation, fmt, ext, save_kwargs): + # Width and height are reported after orientation is applied. + img = Image.fromarray(numpy.zeros((40, 60, 3), dtype=numpy.uint8)) + exif = img.getexif() + exif[0x0112] = orientation # 0x0112 is the EXIF orientation tag + path = tmp_path / f"exif_{orientation}.{ext}" + img.save(path, fmt, exif=exif.tobytes(), **save_kwargs) + + metadata = get_image_metadata(path) + # 0 isn't a valid EXIF orientation, so the decoders ignore it. AVIF + # stores orientation as irot/imir properties, which aren't written for + # the identity orientation 1. + if orientation == 0 or (fmt == "AVIF" and orientation == 1): + assert metadata.rotation is None + else: + expected_rotation = {1: 0, 2: 0, 3: 180, 4: 180, 5: 90, 6: -90} + expected_rotation |= {7: -90, 8: 90} + assert metadata.rotation == expected_rotation[orientation] + if orientation in (5, 6, 7, 8): + assert (metadata.width, metadata.height) == (40, 60) + else: + assert (metadata.width, metadata.height) == (60, 40) + self._assert_matches_decode(metadata, path) + + @needs_png + @pytest.mark.parametrize("orientation", (1, 3, 6, 8)) + def test_rotation_convention(self, tmp_path, orientation): + # Same convention as VideoStreamMetadata.rotation: the decoded image is + # the stored one, rotated counter-clockwise by `rotation` degrees. + arr = torch.randint(0, 256, (40, 60, 3), dtype=torch.uint8).numpy() + img = Image.fromarray(arr) + exif = img.getexif() + exif[0x0112] = orientation # 0x0112 is the EXIF orientation tag + path = tmp_path / "img.png" + img.save(path, exif=exif.tobytes()) + + rotation = get_image_metadata(path).rotation + stored = torch.from_numpy(arr).permute(2, 0, 1) + expected = torch.rot90(stored, k=int(rotation // 90), dims=(-2, -1)) + torch.testing.assert_close(decode_image(path), expected, rtol=0, atol=0) + + @needs_jpeg + def test_jpeg_orientation_comes_from_first_app1(self): + # The decoder only reads the orientation from the first APP1 segment, + # so an XMP segment before the EXIF one hides the orientation. + buf = io.BytesIO() + img = Image.fromarray(numpy.zeros((40, 60, 3), dtype=numpy.uint8)) + exif = img.getexif() + exif[0x0112] = 6 + img.save(buf, "JPEG", exif=exif.tobytes()) + data = buf.getvalue() + xmp = b"http://ns.adobe.com/xap/1.0/\x00" + xmp_segment = b"\xff\xe1" + (len(xmp) + 2).to_bytes(2, "big") + xmp + data = data[:2] + xmp_segment + data[2:] + + metadata = get_image_metadata(data) + assert metadata.rotation is None + assert (metadata.width, metadata.height) == (60, 40) + self._assert_matches_decode(metadata, data) + + @needs_png + @pytest.mark.parametrize( + "pil_mode, save_kwargs, num_channels, bit_depth", + ( + ("P", {}, 3, 8), + ("P", {"transparency": 0}, 4, 8), + ("P", {"bits": 4}, 3, 8), + ("L", {"transparency": 0}, 1, 8), + ("RGB", {"transparency": (0, 0, 0)}, 3, 8), + ("1", {}, 1, 1), + ), + ) + def test_png_color_types( + self, tmp_path, pil_mode, save_kwargs, num_channels, bit_depth + ): + # Palettes are expanded to RGB(A) by the decoder, and a tRNS chunk only + # adds an alpha channel to palette images. + img = Image.fromarray(numpy.zeros((40, 60, 3), dtype=numpy.uint8)) + path = tmp_path / "img.png" + img.convert(pil_mode).save(path, **save_kwargs) + + metadata = get_image_metadata(path) + assert metadata.num_channels == num_channels + assert metadata.bit_depth == bit_depth + self._assert_matches_decode(metadata, path) + + @pytest.mark.parametrize( + "fmt, ext, save_kwargs", + ( + pytest.param("WEBP", "webp", {}, marks=pytest.mark.needs_webp, id="webp"), + pytest.param( + "WEBP", + "webp", + {"lossless": True}, + marks=pytest.mark.needs_webp, + id="webp-lossless", + ), + pytest.param("GIF", "gif", {}, id="gif"), + pytest.param("AVIF", "avif", {}, marks=pytest.mark.needs_avif, id="avif"), + ), + ) + @pytest.mark.parametrize("pil_mode", ("RGB", "RGBA")) + def test_animated(self, tmp_path, fmt, ext, save_kwargs, pil_mode): + frames = [ + Image.fromarray(numpy.full((40, 60, 4), value, dtype=numpy.uint8)).convert( + pil_mode + ) + for value in (0, 100, 200) + ] + path = tmp_path / f"animated.{ext}" + frames[0].save( + path, fmt, save_all=True, append_images=frames[1:], **save_kwargs + ) + + metadata = get_image_metadata(path) + assert metadata.num_frames == 3 + assert (metadata.width, metadata.height) == (60, 40) + self._assert_matches_decode(metadata, path) + + def test_heic_does_not_need_libheif(self, monkeypatch): + # Unlike decode_heic, getting the metadata of a HEIC file works without + # libheif. + from torchcodec import _internally_replaced_utils as iru + + iru.load_heic_library.cache_clear() + + def boom(path): + raise OSError("libheif.so.1: cannot open shared object file") + + monkeypatch.setattr(torch.ops, "load_library", boom) + try: + metadata = get_image_metadata(GRADIENT_HEIC.path) + finally: + iru.load_heic_library.cache_clear() + assert (metadata.width, metadata.height) == (1280, 720) + + def test_unrecognized_format_raises(self): + with pytest.raises(ValueError, match="Unsupported or unrecognized image"): + get_image_metadata(b"not an image at all") + + def test_bad_source_type_raises(self): + with pytest.raises(TypeError, match="Unknown source type"): + get_image_metadata(123) + + @pytest.mark.parametrize( + "asset", + ( + GRADIENT_JPEG, + CMYK_JPEG, + GRADIENT_PNG, + GRADIENT_16BIT_PNG, + GRADIENT_WEBP, + RGBA_WEBP, + GRADIENT_GIF, + ANIMATED_GIF, + GRADIENT_AVIF, + GRADIENT_HEIC, + ANIMATED_HEIC, + ), + ) + def test_truncated(self, asset): + # Truncating the data anywhere must either still work (if the headers + # are intact) or raise a RuntimeError. + data = asset.path.read_bytes() + cut_points = sorted(set(range(1, min(len(data), 300))) | {len(data) // 2}) + for cut in cut_points: + try: + get_image_metadata(data[:cut]) + except RuntimeError: + pass + except ValueError as e: + # Too short for the format to be detected. + assert "Unsupported or unrecognized image" in str(e) + + def test_corrupt_png_raises(self): + with pytest.raises(RuntimeError, match="truncated or corrupted"): + get_image_metadata(SIGSEGV_PNG.path)