Warn when webp is asked to decode into grayscale (#9101)

Author: scotts

test/test_image.py

@@ -1,4 +1,5 @@
 import concurrent.futures
+import contextlib
 import glob
 import io
 import os
@@ -934,6 +935,32 @@ def test_decode_webp(decode_fun, scripted):
     img += 123  # make sure image buffer wasn't freed by underlying decoding lib
 
 
+@pytest.mark.parametrize("decode_fun", (decode_webp, decode_image))
+def test_decode_webp_grayscale(decode_fun, capfd):
+    encoded_bytes = read_file(next(get_images(FAKEDATA_DIR, ".webp")))
+
+    # We warn at the C++ layer because for decode_image(), we don't do the image
+    # type dispatch until we get to the C++ version of decode_image(). We could
+    # warn at the Python layer in decode_webp(), but then users would get a
+    # double wanring: one from the Python layer and one from the C++ layer.
+    #
+    # Because we use the TORCH_WARN_ONCE macro, we need to do this dance to
+    # temporarily always warn so we can test.
+    @contextlib.contextmanager
+    def set_always_warn():
+        torch._C._set_warnAlways(True)
+        yield
+        torch._C._set_warnAlways(False)
+
+    with set_always_warn():
+        img = decode_fun(encoded_bytes, mode=ImageReadMode.GRAY)
+        assert "Webp does not support grayscale conversions" in capfd.readouterr().err
+
+        # Note that because we do not support grayscale conversions, we expect
+        # that the number of color channels is still 3.
+        assert img.shape == (3, 100, 100)
+
+
 # This test is skipped by default because it requires webp images that we're not
 # including within the repo. The test images were downloaded manually from the
 # different pages of https://developers.google.com/speed/webp/gallery

torchvision/csrc/io/image/cpu/decode_webp.cpp

@@ -33,6 +33,13 @@ torch::Tensor decode_webp(
   TORCH_CHECK(
       !features.has_animation, "Animated webp files are not supported.");
 
+  if (mode == IMAGE_READ_MODE_GRAY ||
+      mode == IMAGE_READ_MODE_GRAY_ALPHA) {
+    TORCH_WARN_ONCE(
+        "Webp does not support grayscale conversions. "
+        "The returned tensor will be in the colorspace of the original image.");
+  }
+
   auto return_rgb =
       should_this_return_rgb_or_rgba_let_me_know_in_the_comments_down_below_guys_see_you_in_the_next_video(
           mode, features.has_alpha);