Support specifying output channels in io.image.read_image (pytorch#2988)

* Adding output channels implementation for pngs.

* Adding tests for png.

* Adding channels in the API and documentation.

* Fixing formatting.

* Refactoring test_image.py to remove huge grace_hopper_517x606.pth file from assets and reduce duplicate code. Moving jpeg assets used by encode and write unit-tests on their separate folders.

* Adding output channels implementation for jpegs. Fix asset locations.

* Add tests for JPEG, adding the channels in the API and documentation and adding checks for inputs.

* Changing folder for unit-test.

* Fixing windows flakiness, removing duplicate test.

* Replacing components to channels.

* Adding reference for supporting CMYK.

* Minor changes: num_components to output_components, adding comments, fixing variable name etc.

* Reverting output_components to num_components.

* Replacing decoding with generic method on tests.

* Palette converted to Gray.

Author: datumbox

test/assets/grace_hopper_517x606.jpg renamed to test/assets/encode_jpeg/grace_hopper_517x606.jpg

@@ -25,15 +25,17 @@ 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)
     return x.view(1, 3, 224, 224)
 
 
 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)

test/assets/jpeg_write/grace_hopper_517x606_pil.jpg renamed to test/assets/encode_jpeg/jpeg_write/grace_hopper_517x606_pil.jpg

@@ -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):

test/assets/fakedata/logos/cmyk_pytorch.jpg

@@ -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'

test/assets/fakedata/logos/gray_pytorch.jpg

@@ -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):