Fix png palette decoding when mode is UNCHANGED.

src/torchcodec/_core/DecodeAvif.cpp

@@ -90,14 +90,19 @@ ExifOrientation avif_exif_orientation(const avifImage* image) {
 torch::stable::Tensor decode_avif(
     const torch::stable::Tensor& input,
     int64_t mode,
-    int64_t output_dtype) {
+    int64_t output_dtype,
+    int64_t num_threads) {
   // Based on
   // https://github.com/AOMediaCodec/libavif/blob/main/examples/avif_example_decode_memory.c
   validate_encoded_data(input);
+  STD_TORCH_CHECK(
+      num_threads >= 1, "num_threads must be >= 1, got ", num_threads);
 
   DecoderPtr decoder(avifDecoderCreate());
   STD_TORCH_CHECK(decoder != nullptr, "Failed to create avif decoder.");
 
+  decoder->maxThreads = static_cast<int>(num_threads);
+
   auto result = avifDecoderSetIOMemory(
       decoder.get(), input.const_data_ptr<uint8_t>(), input.numel());
   STD_TORCH_CHECK(

src/torchcodec/_core/DecodeAvif.h

@@ -13,6 +13,7 @@ namespace facebook::torchcodec {
 FORCE_PUBLIC_VISIBILITY torch::stable::Tensor decode_avif(
     const torch::stable::Tensor& input,
     int64_t mode,
-    int64_t output_dtype);
+    int64_t output_dtype,
+    int64_t num_threads);
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/DecodePng.cpp

@@ -175,14 +175,30 @@ PngHeader read_header_and_configure(
     num_passes = 1;
   }
 
-  if (read_mode != ImageReadMode::UNCHANGED) {
-    bool is_palette = (color_type & PNG_COLOR_MASK_PALETTE) != 0;
+  bool is_palette = (color_type & PNG_COLOR_MASK_PALETTE) != 0;
+  // A tRNS chunk encodes transparency without a dedicated alpha channel.
+  // png_set_tRNS_to_alpha() expands it into a real alpha channel (it must be
+  // called after png_set_palette_to_rgb()).
+  bool has_trns = png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS) != 0;
+
+  bool expanded_palette = false;
+  if (read_mode == ImageReadMode::UNCHANGED) {
+    if (is_palette) {
+      // A PNG palette (PLTE chunk) is always RGB triplets
+      // Note that this path is buggy in torchvision, as TV returns the raw
+      // palette indices instead of the expanded RGB values. What we do here is
+      // correct, and matches the behavior of PIL.
+      png_set_palette_to_rgb(png_ptr);
+      num_output_channels = 3;
+      if (has_trns) {
+        png_set_tRNS_to_alpha(png_ptr);
+        num_output_channels = 4;
+      }
+      expanded_palette = true;
+    }
+  } else {
     bool has_color = (color_type & PNG_COLOR_MASK_COLOR) != 0;
     bool has_alpha = (color_type & PNG_COLOR_MASK_ALPHA) != 0;
-    // A tRNS chunk encodes transparency without a dedicated alpha channel.
-    // png_set_tRNS_to_alpha() expands it into a real alpha channel (it must be
-    // called after png_set_palette_to_rgb()).
-    bool has_trns = png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS) != 0;
 
     png_uint_32 opaque_alpha = output_16 ? 65535 : 255;
 
@@ -272,7 +288,8 @@ PngHeader read_header_and_configure(
     need_scaling = true;
   }
 
-  if (read_mode != ImageReadMode::UNCHANGED || need_scaling) {
+  if (read_mode != ImageReadMode::UNCHANGED || need_scaling ||
+      expanded_palette) {
     png_read_update_info(png_ptr, info_ptr);
   }
 

src/torchcodec/_core/image_custom_ops.cpp

@@ -24,7 +24,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(torchcodec_ns, m) {
   m.def("decode_png(Tensor input, int mode, int output_dtype=0) -> Tensor");
   m.def("decode_webp(Tensor input, int mode) -> Tensor");
   m.def("decode_gif(Tensor input, int mode) -> Tensor");
-  m.def("decode_avif(Tensor input, int mode, int output_dtype=0) -> Tensor");
+  m.def(
+      "decode_avif(Tensor input, int mode, int output_dtype=0, int num_threads=1) -> Tensor");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) {

src/torchcodec/decoders/_image_decoders.py

@@ -22,10 +22,6 @@
 # TODO_IMAGE: we need to make FFmpeg an optional dependency, and probably want a
 # CI job for that.
 
-# TODO_IMAGE Some codecs expose threading options - we should make sure the
-# default is 1 thread and allow the user to override it. (similar to
-# num_ffmpeg_threads)
-
 # TODO_IMAGE: I think there are some tests for corrupted images in TV? We
 # should port those.
 
@@ -319,6 +315,7 @@ def decode_avif(
         Literal["UNCHANGED", "GRAY", "GRAY_ALPHA", "RGB", "RGB_ALPHA"] | ImageReadMode
     ) = "RGB",
     output_dtype: torch.dtype | Literal["auto"] = torch.uint8,
+    num_threads: int = 1,
 ) -> torch.Tensor:
     """Decode an AVIF into a tensor of shape ``(C, H, W)``.
 
@@ -334,7 +331,10 @@ def decode_avif(
     data = _source_to_tensor(source)
     code = _OUTPUT_DTYPE_TO_CODE[output_dtype]
     return _decode_with_mode(
-        lambda d, m: _decode_avif(d, m, code), data, mode, _AVIF_NATIVE_OUTPUT_MODES
+        lambda d, m: _decode_avif(d, m, code, num_threads),
+        data,
+        mode,
+        _AVIF_NATIVE_OUTPUT_MODES,
     )
 
 
@@ -378,6 +378,14 @@ def _detect_image_format(data: torch.Tensor) -> str:
     )
 
 
+# Design note: the parameters of decode_image must apply to *all* codecs
+# uniformly. That's why all modes are supported by decode_image even though not
+# all codec would natively support all mode - e.g. jpeg has no alpha support, so
+# we prepend an opaque alpha channel as a post-processing step.  As a resut, all
+# codec-specific entry points like decode_jpeg, decode_png etc.  must still
+# expose the same parameters that decode_image exposes.  The codec-specific
+# parameters should live in the codec-specific entry points, e.g. decode_avif
+# has its `num_threads`, decode_jpeg has `device`, etc.
 def decode_image(
     source: str | Path | bytes | torch.Tensor,
     *,

test/test_decoders.py

@@ -3730,10 +3730,6 @@ def test_all_source_to_all_output_modes(
     ):
         # Test that every input color mode is decodable to every output mode.
 
-        if fmt == "PNG" and source_mode == "P" and output_mode == "UNCHANGED":
-            # TODO_IMAGE figure out what to do here.
-            pytest.skip("UNCHANGED on a palette PNG returns raw palette indices")
-
         h, w = 40, 60
         xs = numpy.linspace(0, 255, w)
         ys = numpy.linspace(0, 255, h)
@@ -4192,6 +4188,17 @@ def test_avif_against_pil(self, asset, mode, pil_mode):
         assert decoded.shape == reference.shape
         assert_tensor_close_on_at_least(decoded, reference, percentage=99, atol=2)
 
+    @needs_avif
+    def test_avif_num_threads(self):
+        reference = decode_avif(GRADIENT_AVIF.path)
+        for num_threads in (1, 2, 4):
+            decoded = decode_avif(GRADIENT_AVIF.path, num_threads=num_threads)
+            assert torch.equal(decoded, reference)
+
+        for bad in (0, -1):
+            with pytest.raises(RuntimeError, match="num_threads must be >= 1"):
+                decode_avif(GRADIENT_AVIF.path, num_threads=bad)
+
     @needs_avif
     @pytest.mark.parametrize(
         "mode, pil_mode",