Repository: pytorch/vision

Issue:
Improved format conversion in io.image.read_image
## 🚀 Feature

Replace the `channels` parameter of the `read_image()`, `decode_image()`, `decode_png()` and `decode_jpeg()` methods with a `mode` parameter that gives users better control on image conversions.

## Motivation

Issue #2948 proposed specifying the number of output channels when loading images and PR #2988 introduced the change on the above methods. Though the API provides similar functionality to other libraries, it requires making implicit assumptions in relation to the mapping between # channels and output formats. For example, we assumed that `channel=1` means Grayscale. Nevertheless since both Palette images and Gayscale images use 1 channel, we had to introduce logic for handling such corner-cases. 

A better approach would be to give control to the users to explicitly define what type of conversions they want to make. 

## Pitch

Building on top of @vfdev-5's proposal, we could replace `channels` with a `mode` parameter. For example:
```python
# old:
def read_image(path: str, channels: int = 0) -> torch.Tensor

# new:
def read_image(path: str, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor
```

The `mode` will be an enum with the following values:
```
ImageReadMode.UNCHANGED
ImageReadMode.GRAY
ImageReadMode.GRAY_ALPHA
ImageReadMode.RGB
ImageReadMode.RGB_ALPHA
```

The default value of `mode` will be ImageReadMode.UNCHANGED and it will have similar behaviour as the current `channels=0`. It will load the image without making any modification and to ensure BC it will additionally support Palette, CMYK and other currently supported formats. 

**Note:** The scope of this proposal is to change this experimental API to allow for better image format support on the future. Adding support for converting images to Palette, from/to CMYK, etc is not within the scope of this proposal. Many of such conversions are not supported by LibJPEG and LibPNG and require writing custom conversion code which should be handled in a separate ticket.


Discussion:
The proposal sounds great to me.
One thing we should decide is what to do if the original input format doesn't support one specific configuration. For example, if the user passes `RGB_ALPHA` to a `JPEG` image, should we return a `RGB_ALPHA` with a all-ones (or zeros) alpha channel, or error out?
Thanks for the feedback Francisco. 

I had in mind to raise an error. This is exactly what we currently do if someone passes `channels=4` for a JPEG:
https://github.com/pytorch/vision/blob/8c281757a0daf3a8e92cbb4bded0e5e6b389a375/torchvision/csrc/cpu/image/readjpeg_cpu.cpp#L108-L123

If on the future we decide that RGB_ALPHA is one of the supported conversion formats for jpeg images, we will simply add an entry on the above snippet along with the conversion code. In this case, as you hinted, the conversion code is simply adding a 4th channel with value 255 everywhere.

Resolving patch:
diff --git a/torchvision/io/image.py b/torchvision/io/image.py
--- a/torchvision/io/image.py
+++ b/torchvision/io/image.py
@@ -4,6 +4,8 @@
 import os.path as osp
 import importlib.machinery
 
+from enum import Enum
+
 _HAS_IMAGE_OPT = False
 
 try:
@@ -47,6 +49,14 @@
     pass
 
 
+class ImageReadMode(Enum):
+    UNCHANGED = 0
+    GRAY = 1
+    GRAY_ALPHA = 2
+    RGB = 3
+    RGB_ALPHA = 4
+
+
 def read_file(path: str) -> torch.Tensor:
     """
     Reads and outputs the bytes contents of a file as a uint8 Tensor
@@ -74,24 +84,26 @@ def write_file(filename: str, data: torch.Tensor) -> None:
     torch.ops.image.write_file(filename, data)
 
 
-def decode_png(input: torch.Tensor, channels: int = 0) -> torch.Tensor:
+def decode_png(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
     """
     Decodes a PNG image into a 3 dimensional RGB Tensor.
-    Optionally converts the image to the desired number of color channels.
+    Optionally converts the image to the desired format.
     The values of the output tensor are uint8 between 0 and 255.
 
     Arguments:
         input (Tensor[1]): a one dimensional uint8 tensor containing
     the raw bytes of the PNG image.
-        channels (int): the number of output channels for the decoded
-    image. 0 keeps the original number of channels, 1 converts to Grayscale
-    2 converts to Grayscale with Alpha, 3 converts to RGB and 4 coverts to
-    RGB with Alpha. Default: 0
+        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`
 
     Returns:
         output (Tensor[image_channels, image_height, image_width])
     """
-    output = torch.ops.image.decode_png(input, channels)
+    output = torch.ops.image.decode_png(input, mode.value)
     return output
 
 
@@ -137,23 +149,24 @@ def write_png(input: torch.Tensor, filename: str, compression_level: int = 6):
     write_file(filename, output)
 
 
-def decode_jpeg(input: torch.Tensor, channels: int = 0) -> torch.Tensor:
+def decode_jpeg(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
     """
     Decodes a JPEG image into a 3 dimensional RGB Tensor.
-    Optionally converts the image to the desired number of color channels.
+    Optionally converts the image to the desired format.
     The values of the output tensor are uint8 between 0 and 255.
 
     Arguments:
         input (Tensor[1]): a one dimensional uint8 tensor containing
     the raw bytes of the JPEG image.
-        channels (int): the number of output channels for the decoded
-    image. 0 keeps the original number of channels, 1 converts to Grayscale
-    and 3 converts to RGB. Default: 0
+        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`
 
     Returns:
         output (Tensor[image_channels, image_height, image_width])
     """
-    output = torch.ops.image.decode_jpeg(input, channels)
+    output = torch.ops.image.decode_jpeg(input, mode.value)
     return output
 
 
@@ -202,12 +215,12 @@ def write_jpeg(input: torch.Tensor, filename: str, quality: int = 75):
     write_file(filename, output)
 
 
-def decode_image(input: torch.Tensor, channels: int = 0) -> torch.Tensor:
+def decode_image(input: torch.Tensor, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
     """
     Detects whether an image is a JPEG or PNG and performs the appropriate
     operation to decode the image into a 3 dimensional RGB Tensor.
 
-    Optionally converts the image to the desired number of color channels.
+    Optionally converts the image to the desired format.
     The values of the output tensor are uint8 between 0 and 255.
 
     Parameters
@@ -215,39 +228,41 @@ def decode_image(input: torch.Tensor, channels: int = 0) -> torch.Tensor:
     input: Tensor
         a one dimensional uint8 tensor containing the raw bytes of the
         PNG or JPEG image.
-    channels: int
-        the number of output channels of the decoded image. JPEG and PNG images
-        have different permitted values. The default value is 0 and it keeps
-        the original number of channels. See `decode_jpeg()` and `decode_png()`
-        for more information. Default: 0
+    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`
 
     Returns
     -------
     output: Tensor[image_channels, image_height, image_width]
     """
-    output = torch.ops.image.decode_image(input, channels)
+    output = torch.ops.image.decode_image(input, mode.value)
     return output
 
 
-def read_image(path: str, channels: int = 0) -> torch.Tensor:
+def read_image(path: str, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
     """
     Reads a JPEG or PNG image into a 3 dimensional RGB Tensor.
-    Optionally converts the image to the desired number of color channels.
+    Optionally converts the image to the desired format.
     The values of the output tensor are uint8 between 0 and 255.
 
     Parameters
     ----------
     path: str
         path of the JPEG or PNG image.
-    channels: int
-        the number of output channels of the decoded image. JPEG and PNG images
-        have different permitted values. The default value is 0 and it keeps
-        the original number of channels. See `decode_jpeg()` and `decode_png()`
-        for more information. Default: 0
+    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`
 
     Returns
     -------
     output: Tensor[image_channels, image_height, image_width]
     """
     data = read_file(path)
-    return decode_image(data, channels)
+    return decode_image(data, mode)


Test patch:
diff --git a/test/test_image.py b/test/test_image.py
--- a/test/test_image.py
+++ b/test/test_image.py
@@ -2,14 +2,12 @@
 import io
 import glob
 import unittest
-import sys
 
 import torch
-import torchvision
 from PIL import Image
 from torchvision.io.image import (
     decode_png, decode_jpeg, encode_jpeg, write_jpeg, decode_image, read_file,
-    encode_png, write_png, write_file)
+    encode_png, write_png, write_file, ImageReadMode)
 import numpy as np
 
 from common_utils import get_tmp_dir
@@ -49,9 +47,9 @@ def normalize_dimensions(img_pil):
 
 class ImageTester(unittest.TestCase):
     def test_decode_jpeg(self):
-        conversion = [(None, 0), ("L", 1), ("RGB", 3)]
+        conversion = [(None, ImageReadMode.UNCHANGED), ("L", ImageReadMode.GRAY), ("RGB", ImageReadMode.RGB)]
         for img_path in get_images(IMAGE_ROOT, ".jpg"):
-            for pil_mode, channels in conversion:
+            for pil_mode, mode in conversion:
                 with Image.open(img_path) as img:
                     is_cmyk = img.mode == "CMYK"
                     if pil_mode is not None:
@@ -66,7 +64,7 @@ def test_decode_jpeg(self):
 
                 img_pil = normalize_dimensions(img_pil)
                 data = read_file(img_path)
-                img_ljpeg = decode_image(data, channels=channels)
+                img_ljpeg = decode_image(data, mode=mode)
 
                 # Permit a small variation on pixel values to account for implementation
                 # differences between Pillow and LibJPEG.
@@ -165,9 +163,10 @@ def test_write_jpeg(self):
             self.assertEqual(torch_bytes, pil_bytes)
 
     def test_decode_png(self):
-        conversion = [(None, 0), ("L", 1), ("LA", 2), ("RGB", 3), ("RGBA", 4)]
+        conversion = [(None, ImageReadMode.UNCHANGED), ("L", ImageReadMode.GRAY), ("LA", ImageReadMode.GRAY_ALPHA),
+                      ("RGB", ImageReadMode.RGB), ("RGBA", ImageReadMode.RGB_ALPHA)]
         for img_path in get_images(FAKEDATA_DIR, ".png"):
-            for pil_mode, channels in conversion:
+            for pil_mode, mode in conversion:
                 with Image.open(img_path) as img:
                     if pil_mode is not None:
                         img = img.convert(pil_mode)
@@ -175,7 +174,7 @@ def test_decode_png(self):
 
                 img_pil = normalize_dimensions(img_pil)
                 data = read_file(img_path)
-                img_lpng = decode_image(data, channels=channels)
+                img_lpng = decode_image(data, mode=mode)
 
                 tol = 0 if conversion is None else 1
                 self.assertTrue(img_lpng.allclose(img_pil, atol=tol))