diff --git a/torchvision/csrc/io/image/cpu/readjpeg_cpu.cpp b/torchvision/csrc/io/image/cpu/readjpeg_cpu.cpp index 7f1a411bc3d..94121b4e78f 100644 --- a/torchvision/csrc/io/image/cpu/readjpeg_cpu.cpp +++ b/torchvision/csrc/io/image/cpu/readjpeg_cpu.cpp @@ -117,7 +117,7 @@ torch::Tensor decodeJPEG(const torch::Tensor& data, ImageReadMode mode) { */ default: jpeg_destroy_decompress(&cinfo); - TORCH_CHECK(false, "Provided mode not supported"); + TORCH_CHECK(false, "The provided mode is not supported for JPEG files"); } jpeg_calc_output_dimensions(&cinfo); diff --git a/torchvision/csrc/io/image/cpu/readpng_cpu.cpp b/torchvision/csrc/io/image/cpu/readpng_cpu.cpp index adbbe52e097..bb2d70b47cd 100644 --- a/torchvision/csrc/io/image/cpu/readpng_cpu.cpp +++ b/torchvision/csrc/io/image/cpu/readpng_cpu.cpp @@ -143,7 +143,7 @@ torch::Tensor decodePNG(const torch::Tensor& data, ImageReadMode mode) { break; default: png_destroy_read_struct(&png_ptr, &info_ptr, nullptr); - TORCH_CHECK(false, "Provided mode not supported"); + TORCH_CHECK(false, "The provided mode is not supported for PNG files"); } png_read_update_info(png_ptr, info_ptr); diff --git a/torchvision/csrc/io/image/image_read_mode.h b/torchvision/csrc/io/image/image_read_mode.h index 00ff4f6b581..5df82d373d9 100644 --- a/torchvision/csrc/io/image/image_read_mode.h +++ b/torchvision/csrc/io/image/image_read_mode.h @@ -2,8 +2,8 @@ /* Should be kept in-sync with Python ImageReadMode enum */ using ImageReadMode = int64_t; -#define IMAGE_READ_MODE_UNCHANGED 0 -#define IMAGE_READ_MODE_GRAY 1 -#define IMAGE_READ_MODE_GRAY_ALPHA 2 -#define IMAGE_READ_MODE_RGB 3 -#define IMAGE_READ_MODE_RGB_ALPHA 4 \ No newline at end of file +const ImageReadMode IMAGE_READ_MODE_UNCHANGED = 0; +const ImageReadMode IMAGE_READ_MODE_GRAY = 1; +const ImageReadMode IMAGE_READ_MODE_GRAY_ALPHA = 2; +const ImageReadMode IMAGE_READ_MODE_RGB = 3; +const ImageReadMode IMAGE_READ_MODE_RGB_ALPHA = 4; diff --git a/torchvision/io/image.py b/torchvision/io/image.py index f11a7ccb634..6ea4f0c3157 100644 --- a/torchvision/io/image.py +++ b/torchvision/io/image.py @@ -50,6 +50,15 @@ class ImageReadMode(Enum): + """ + Support for various modes while reading images. + + Use `ImageReadMode.UNCHANGED` for loading the image as-is, + `ImageReadMode.GRAY` for converting to grayscale, + `ImageReadMode.GRAY_ALPHA` for grayscale with transparency, + `ImageReadMode.RGB` for RGB and `ImageReadMode.RGB_ALPHA` for + RGB with transparency. + """ UNCHANGED = 0 GRAY = 1 GRAY_ALPHA = 2 @@ -94,11 +103,9 @@ def decode_png(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGE input (Tensor[1]): a one dimensional uint8 tensor containing the raw bytes of the PNG image. mode (ImageReadMode): the read mode used for optionally - converting the image. Use `ImageReadMode.UNCHANGED` for loading - the image as-is, `ImageReadMode.GRAY` for converting to grayscale, - `ImageReadMode.GRAY_ALPHA` for grayscale with transparency, - `ImageReadMode.RGB` for RGB and `ImageReadMode.RGB_ALPHA` for - RGB with transparency. Default: `ImageReadMode.UNCHANGED` + converting the image. Default: `ImageReadMode.UNCHANGED`. + See `ImageReadMode` class for more information on various + available modes. Returns: output (Tensor[image_channels, image_height, image_width]) @@ -159,9 +166,9 @@ def decode_jpeg(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANG input (Tensor[1]): a one dimensional uint8 tensor containing the raw bytes of the JPEG image. mode (ImageReadMode): the read mode used for optionally - converting the image. Use `ImageReadMode.UNCHANGED` for loading - the image as-is, `ImageReadMode.GRAY` for converting to grayscale - and `ImageReadMode.RGB` for RGB. Default: `ImageReadMode.UNCHANGED` + converting the image. Default: `ImageReadMode.UNCHANGED`. + See `ImageReadMode` class for more information on various + available modes. Returns: output (Tensor[image_channels, image_height, image_width]) @@ -229,11 +236,10 @@ def decode_image(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHAN a one dimensional uint8 tensor containing the raw bytes of the PNG or JPEG image. mode: ImageReadMode - the read mode used for optionally converting the image. JPEG - and PNG images have different permitted values. The default - value is `ImageReadMode.UNCHANGED` and it keeps the image as-is. - See `decode_jpeg()` and `decode_png()` for more information. - Default: `ImageReadMode.UNCHANGED` + the read mode used for optionally converting the image. + Default: `ImageReadMode.UNCHANGED`. + See `ImageReadMode` class for more information on various + available modes. Returns ------- @@ -254,11 +260,10 @@ def read_image(path: str, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torc path: str path of the JPEG or PNG image. mode: ImageReadMode - the read mode used for optionally converting the image. JPEG - and PNG images have different permitted values. The default - value is `ImageReadMode.UNCHANGED` and it keeps the image as-is. - See `decode_jpeg()` and `decode_png()` for more information. - Default: `ImageReadMode.UNCHANGED` + the read mode used for optionally converting the image. + Default: `ImageReadMode.UNCHANGED`. + See `ImageReadMode` class for more information on various + available modes. Returns -------