diff --git a/test/test_image.py b/test/test_image.py index d7c7bb542bb..b11dd67ca12 100644 --- a/test/test_image.py +++ b/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 diff --git a/torchvision/csrc/io/image/cpu/decode_webp.cpp b/torchvision/csrc/io/image/cpu/decode_webp.cpp index 4c13c5c2b1a..5276bf2d9f2 100644 --- a/torchvision/csrc/io/image/cpu/decode_webp.cpp +++ b/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);