From 4d6f0c06b45fb2d3183094db1b0c1151058c1103 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 28 Aug 2024 11:56:40 +0100 Subject: [PATCH] Minor cleanup of webp decoding mode logic (#8612) --- torchvision/csrc/io/image/cpu/decode_webp.cpp | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/torchvision/csrc/io/image/cpu/decode_webp.cpp b/torchvision/csrc/io/image/cpu/decode_webp.cpp index c5fcc51cec2..bf115c23c41 100644 --- a/torchvision/csrc/io/image/cpu/decode_webp.cpp +++ b/torchvision/csrc/io/image/cpu/decode_webp.cpp @@ -40,20 +40,22 @@ torch::Tensor decode_webp( TORCH_CHECK( !features.has_animation, "Animated webp files are not supported."); - auto decoding_func = WebPDecodeRGB; - int num_channels = 0; - if (mode == IMAGE_READ_MODE_RGB) { - decoding_func = WebPDecodeRGB; - num_channels = 3; - } else if (mode == IMAGE_READ_MODE_RGB_ALPHA) { - decoding_func = WebPDecodeRGBA; - num_channels = 4; - } else { - // Assume mode is "unchanged" - decoding_func = features.has_alpha ? WebPDecodeRGBA : WebPDecodeRGB; - num_channels = features.has_alpha ? 4 : 3; + if (mode != IMAGE_READ_MODE_UNCHANGED && mode != IMAGE_READ_MODE_RGB && + mode != IMAGE_READ_MODE_RGB_ALPHA) { + // Other modes aren't supported, but we don't error or even warn because we + // have generic entry points like decode_image which may support all modes, + // it just depends on the underlying decoder. + mode = IMAGE_READ_MODE_UNCHANGED; } + // If return_rgb is false it means we return rgba - nothing else. + auto return_rgb = + (mode == IMAGE_READ_MODE_RGB || + (mode == IMAGE_READ_MODE_UNCHANGED && !features.has_alpha)); + + auto decoding_func = return_rgb ? WebPDecodeRGB : WebPDecodeRGBA; + auto num_channels = return_rgb ? 3 : 4; + int width = 0; int height = 0;