Add get_image_metadata()

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:
 

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<int64_t> 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<int64_t> 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<int64_t>(decoder->image->depth),
+      num_frames,
+      has_orientation ? static_cast<int64_t>(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<uint8_t>(),
+      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

src/torchcodec/_core/DecodeAvif.h

@@ -6,6 +6,9 @@
 
 #pragma once
 
+#include <string>
+#include <vector>
+
 #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<int64_t> probe_avif(
+    const torch::stable::Tensor& input);
+
+FORCE_PUBLIC_VISIBILITY std::vector<int64_t> probe_avif_from_file(
+    std::string filename);
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/DecodeGif.cpp

@@ -26,16 +26,36 @@ torch::stable::Tensor decode_gif(
       "prebuilt wheel, please report it to the TorchCodec repo.");
 }
 
+std::vector<int64_t> 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<int64_t> 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
 
 #include <torch/headeronly/core/MemoryFormat.h>
 
 #include <algorithm>
+#include <climits>
 #include <cstdint>
 #include <optional>
 
+#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<int>(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<FileIO*>(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<int64_t>(gif_file->SWidth, first_image_width);
+  int64_t height = std::max<int64_t>(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<uint8_t>(),
+      .count = static_cast<size_t>(contig_input.numel())};
+  GifFileType* gif_file =
+      DGifOpen(static_cast<void*>(&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<void*>(&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

src/torchcodec/_core/DecodeGif.h

@@ -6,6 +6,9 @@
 
 #pragma once
 
+#include <string>
+#include <vector>
+
 #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<int64_t> probe_gif(
+    const torch::stable::Tensor& input);
+
+FORCE_PUBLIC_VISIBILITY std::vector<int64_t> probe_gif_from_file(
+    std::string filename);
+
 } // namespace facebook::torchcodec

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) {

src/torchcodec/_core/ImageCommon.h

@@ -9,6 +9,8 @@
 #include <torch/csrc/stable/ops.h>
 #include <torch/headeronly/util/Exception.h>
 
+#include <vector>
+
 #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<int64_t>;
+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(

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