Repository: pytorch/vision

Issue:
Support specifying output channels (RGB vs grayscale) for io.image.read_image
## 🚀 Feature
Similar to `channels` argument for `tf.io.decode_image` or `PIL.Image.open(..).convert('format')` allow specifying the output channel 'format' and have a sensible conversion performed for you when decoding an image if set.

## Motivation

Prevent the need to manually check and do color conversion from RGB -> grayscale or grayscale -> RGB when datasets often have mixes of RGB and grayscale images. The output color format of image loading in data pipeline needs to be consistent for many typical use cases.

The lower level libraries (libjpeg, libpng) etc have facilities for specifying the output color space and will do the conversion for you if you set it up (ie `out_color_space` for libjpeg). This is what tensorflow decode_image does and is the most efficient.

## Pitch

Add a `channels=` arg to read_image that matches Tensorflow semantics.

* channels=0 - leave as original, a grayscale image decodes to 1 channel out, rgb to 3
* channels=1 - grayscale out
* channels=3 - RGB out 
* channels=4 - RGBA out (PNG or formats that support alpha, not valid for jpeg)

## Alternatives

## Additional context




Discussion:
@rwightman Thanks for reporting.

This is also related to #2823.
Thanks for the feature request @rwightman !

I think this is a nice addition to have, and we will be looking into adding support for this feature in torchvision.
We are already planning a few improvements to PNG support as @datumbox mentioned, so we could go one step further and also allow for a finer-grained control on the output format for `read_image`.

Resolving patch:
diff --git a/torchvision/io/image.py b/torchvision/io/image.py
--- a/torchvision/io/image.py
+++ b/torchvision/io/image.py
@@ -74,19 +74,24 @@ def write_file(filename: str, data: torch.Tensor) -> None:
     torch.ops.image.write_file(filename, data)
 
 
-def decode_png(input: torch.Tensor) -> torch.Tensor:
+def decode_png(input: torch.Tensor, channels: int = 0) -> torch.Tensor:
     """
     Decodes a PNG image into a 3 dimensional RGB Tensor.
+    Optionally converts the image to the desired number of color channels.
     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
 
     Returns:
-        output (Tensor[3, image_height, image_width])
+        output (Tensor[image_channels, image_height, image_width])
     """
-    output = torch.ops.image.decode_png(input)
+    output = torch.ops.image.decode_png(input, channels)
     return output
 
 
@@ -132,17 +137,23 @@ def write_png(input: torch.Tensor, filename: str, compression_level: int = 6):
     write_file(filename, output)
 
 
-def decode_jpeg(input: torch.Tensor) -> torch.Tensor:
+def decode_jpeg(input: torch.Tensor, channels: int = 0) -> torch.Tensor:
     """
     Decodes a JPEG image into a 3 dimensional RGB Tensor.
+    Optionally converts the image to the desired number of color channels.
     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
+
     Returns:
-        output (Tensor[3, image_height, image_width])
+        output (Tensor[image_channels, image_height, image_width])
     """
-    output = torch.ops.image.decode_jpeg(input)
+    output = torch.ops.image.decode_jpeg(input, channels)
     return output
 
 
@@ -191,11 +202,12 @@ def write_jpeg(input: torch.Tensor, filename: str, quality: int = 75):
     write_file(filename, output)
 
 
-def decode_image(input: torch.Tensor) -> torch.Tensor:
+def decode_image(input: torch.Tensor, channels: int = 0) -> 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.
     The values of the output tensor are uint8 between 0 and 255.
 
     Parameters
@@ -203,28 +215,39 @@ def decode_image(input: torch.Tensor) -> 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
 
     Returns
     -------
-    output: Tensor[3, image_height, image_width]
+    output: Tensor[image_channels, image_height, image_width]
     """
-    output = torch.ops.image.decode_image(input)
+    output = torch.ops.image.decode_image(input, channels)
     return output
 
 
-def read_image(path: str) -> torch.Tensor:
+def read_image(path: str, channels: int = 0) -> 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.
     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
 
     Returns
     -------
-    output: Tensor[3, image_height, image_width]
+    output: Tensor[image_channels, image_height, image_width]
     """
     data = read_file(path)
-    return decode_image(data)
+    return decode_image(data, channels)


Test patch:
diff --git a/test/assets/grace_hopper_517x606.jpg b/test/assets/encode_jpeg/grace_hopper_517x606.jpg
similarity index 100%
rename from test/assets/grace_hopper_517x606.jpg
rename to test/assets/encode_jpeg/grace_hopper_517x606.jpg
diff --git a/test/assets/jpeg_write/grace_hopper_517x606_pil.jpg b/test/assets/encode_jpeg/jpeg_write/grace_hopper_517x606_pil.jpg
similarity index 100%
rename from test/assets/jpeg_write/grace_hopper_517x606_pil.jpg
rename to test/assets/encode_jpeg/jpeg_write/grace_hopper_517x606_pil.jpg
diff --git a/test/assets/fakedata/logos/cmyk_pytorch.jpg b/test/assets/fakedata/logos/cmyk_pytorch.jpg
new file mode 100644
Binary files /dev/null and b/test/assets/fakedata/logos/cmyk_pytorch.jpg differ
diff --git a/test/assets/fakedata/logos/gray_pytorch.jpg b/test/assets/fakedata/logos/gray_pytorch.jpg
new file mode 100644
Binary files /dev/null and b/test/assets/fakedata/logos/gray_pytorch.jpg differ
diff --git a/test/assets/fakedata/logos/rgb_pytorch.jpg b/test/assets/fakedata/logos/rgb_pytorch.jpg
new file mode 100644
Binary files /dev/null and b/test/assets/fakedata/logos/rgb_pytorch.jpg differ
diff --git a/test/assets/grace_hopper_517x606.pth b/test/assets/grace_hopper_517x606.pth
deleted file mode 100644
Binary files a/test/assets/grace_hopper_517x606.pth and /dev/null differ
diff --git a/test/test_cpp_models.py b/test/test_cpp_models.py
--- a/test/test_cpp_models.py
+++ b/test/test_cpp_models.py
@@ -25,7 +25,8 @@ def process_model(model, tensor, func, name):
 
 
 def read_image1():
-    image_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'assets', 'grace_hopper_517x606.jpg')
+    image_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'assets', 'encode_jpeg',
+                              'grace_hopper_517x606.jpg')
     image = Image.open(image_path)
     image = image.resize((224, 224))
     x = F.to_tensor(image)
@@ -33,7 +34,8 @@ def read_image1():
 
 
 def read_image2():
-    image_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'assets', 'grace_hopper_517x606.jpg')
+    image_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'assets', 'encode_jpeg',
+                              'grace_hopper_517x606.jpg')
     image = Image.open(image_path)
     image = image.resize((299, 299))
     x = F.to_tensor(image)
diff --git a/test/test_datasets_utils.py b/test/test_datasets_utils.py
--- a/test/test_datasets_utils.py
+++ b/test/test_datasets_utils.py
@@ -14,7 +14,7 @@
 
 
 TEST_FILE = get_file_path_2(
-    os.path.dirname(os.path.abspath(__file__)), 'assets', 'grace_hopper_517x606.jpg')
+    os.path.dirname(os.path.abspath(__file__)), 'assets', 'encode_jpeg', 'grace_hopper_517x606.jpg')
 
 
 class Tester(unittest.TestCase):
diff --git a/test/test_image.py b/test/test_image.py
--- a/test/test_image.py
+++ b/test/test_image.py
@@ -19,6 +19,7 @@
 FAKEDATA_DIR = os.path.join(IMAGE_ROOT, "fakedata")
 IMAGE_DIR = os.path.join(FAKEDATA_DIR, "imagefolder")
 DAMAGED_JPEG = os.path.join(IMAGE_ROOT, 'damaged_jpeg')
+ENCODE_JPEG = os.path.join(IMAGE_ROOT, "encode_jpeg")
 
 
 def get_images(directory, img_ext):
@@ -33,14 +34,44 @@ def get_images(directory, img_ext):
                 yield os.path.join(root, fl)
 
 
+def pil_read_image(img_path):
+    with Image.open(img_path) as img:
+        return torch.from_numpy(np.array(img))
+
+
+def normalize_dimensions(img_pil):
+    if len(img_pil.shape) == 3:
+        img_pil = img_pil.permute(2, 0, 1)
+    else:
+        img_pil = img_pil.unsqueeze(0)
+    return img_pil
+
+
 class ImageTester(unittest.TestCase):
     def test_decode_jpeg(self):
+        conversion = [(None, 0), ("L", 1), ("RGB", 3)]
         for img_path in get_images(IMAGE_ROOT, ".jpg"):
-            img_pil = torch.load(img_path.replace('jpg', 'pth'))
-            img_pil = img_pil.permute(2, 0, 1)
-            data = read_file(img_path)
-            img_ljpeg = decode_jpeg(data)
-            self.assertTrue(img_ljpeg.equal(img_pil))
+            for pil_mode, channels in conversion:
+                with Image.open(img_path) as img:
+                    is_cmyk = img.mode == "CMYK"
+                    if pil_mode is not None:
+                        if is_cmyk:
+                            # libjpeg does not support the conversion
+                            continue
+                        img = img.convert(pil_mode)
+                    img_pil = torch.from_numpy(np.array(img))
+                    if is_cmyk:
+                        # flip the colors to match libjpeg
+                        img_pil = 255 - img_pil
+
+                img_pil = normalize_dimensions(img_pil)
+                data = read_file(img_path)
+                img_ljpeg = decode_image(data, channels=channels)
+
+                # Permit a small variation on pixel values to account for implementation
+                # differences between Pillow and LibJPEG.
+                abs_mean_diff = (img_ljpeg.type(torch.float32) - img_pil).abs().mean().item()
+                self.assertTrue(abs_mean_diff < 2)
 
         with self.assertRaisesRegex(RuntimeError, "Expected a non empty 1-dimensional tensor"):
             decode_jpeg(torch.empty((100, 1), dtype=torch.uint8))
@@ -68,7 +99,7 @@ def test_damaged_images(self):
                 decode_jpeg(data)
 
     def test_encode_jpeg(self):
-        for img_path in get_images(IMAGE_ROOT, ".jpg"):
+        for img_path in get_images(ENCODE_JPEG, ".jpg"):
             dirname = os.path.dirname(img_path)
             filename, _ = os.path.splitext(os.path.basename(img_path))
             write_folder = os.path.join(dirname, 'jpeg_write')
@@ -111,7 +142,7 @@ def test_encode_jpeg(self):
             encode_jpeg(torch.empty((100, 100), dtype=torch.uint8))
 
     def test_write_jpeg(self):
-        for img_path in get_images(IMAGE_ROOT, ".jpg"):
+        for img_path in get_images(ENCODE_JPEG, ".jpg"):
             data = read_file(img_path)
             img = decode_jpeg(data)
 
@@ -134,20 +165,25 @@ 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)]
         for img_path in get_images(FAKEDATA_DIR, ".png"):
-            img_pil = torch.from_numpy(np.array(Image.open(img_path)))
-            if len(img_pil.shape) == 3:
-                img_pil = img_pil.permute(2, 0, 1)
-            else:
-                img_pil = img_pil.unsqueeze(0)
-            data = read_file(img_path)
-            img_lpng = decode_png(data)
-            self.assertTrue(img_lpng.equal(img_pil))
+            for pil_mode, channels in conversion:
+                with Image.open(img_path) as img:
+                    if pil_mode is not None:
+                        img = img.convert(pil_mode)
+                    img_pil = torch.from_numpy(np.array(img))
 
-            with self.assertRaises(RuntimeError):
-                decode_png(torch.empty((), dtype=torch.uint8))
-            with self.assertRaises(RuntimeError):
-                decode_png(torch.randint(3, 5, (300,), dtype=torch.uint8))
+                img_pil = normalize_dimensions(img_pil)
+                data = read_file(img_path)
+                img_lpng = decode_image(data, channels=channels)
+
+                tol = 0 if conversion is None else 1
+                self.assertTrue(img_lpng.allclose(img_pil, atol=tol))
+
+        with self.assertRaises(RuntimeError):
+            decode_png(torch.empty((), dtype=torch.uint8))
+        with self.assertRaises(RuntimeError):
+            decode_png(torch.randint(3, 5, (300,), dtype=torch.uint8))
 
     def test_encode_png(self):
         for img_path in get_images(IMAGE_DIR, '.png'):
@@ -196,19 +232,6 @@ def test_write_png(self):
 
             self.assertTrue(img_pil.equal(saved_image))
 
-    def test_decode_image(self):
-        for img_path in get_images(IMAGE_ROOT, ".jpg"):
-            img_pil = torch.load(img_path.replace('jpg', 'pth'))
-            img_pil = img_pil.permute(2, 0, 1)
-            img_ljpeg = decode_image(read_file(img_path))
-            self.assertTrue(img_ljpeg.equal(img_pil))
-
-        for img_path in get_images(IMAGE_DIR, ".png"):
-            img_pil = torch.from_numpy(np.array(Image.open(img_path)))
-            img_pil = img_pil.permute(2, 0, 1)
-            img_lpng = decode_image(read_file(img_path))
-            self.assertTrue(img_lpng.equal(img_pil))
-
     def test_read_file(self):
         with get_tmp_dir() as d:
             fname, content = 'test1.bin', b'TorchVision\211\n'
diff --git a/test/test_transforms.py b/test/test_transforms.py
--- a/test/test_transforms.py
+++ b/test/test_transforms.py
@@ -24,7 +24,7 @@
 
 
 GRACE_HOPPER = get_file_path_2(
-    os.path.dirname(os.path.abspath(__file__)), 'assets', 'grace_hopper_517x606.jpg')
+    os.path.dirname(os.path.abspath(__file__)), 'assets', 'encode_jpeg', 'grace_hopper_517x606.jpg')
 
 
 class Tester(unittest.TestCase):