From f39cb3bcaa45bfca55fa21ace78f6b677148e4c2 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Mon, 2 Sep 2024 02:40:53 -0700 Subject: [PATCH] [fbsync] Minor cleanup of webp decoding mode logic (#8612) Reviewed By: ahmadsharif1 Differential Revision: D62032044 fbshipit-source-id: b453a8b82d2d386f5f3e75c632f4605ac63368c2 --- 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;