Return rgb when rgba is loaded

chainercv/utils/image/read_image.py

@@ -1,3 +1,5 @@
+from __future__ import division
+
 import chainer
 import numpy as np
 from PIL import Image
@@ -10,9 +12,35 @@
     _cv2_available = False
 
 
-def _read_image_cv2(path, dtype, color):
+def _handle_four_channel_image(img, alpha):
+    if alpha is None:
+        raise ValueError(
+            'An RGBA image is read by chainercv.utils.read_image, '
+            'but the `alpha` option is not set. Please set the option so that '
+            'the function knows how to handle RGBA images.'
+        )
+    elif alpha == 'ignore':
+        img = img[:, :, :3]
+    elif alpha == 'blend_with_white':
+        color_channel = img[:, :, :3]
+        alpha_channel = img[:, :, 3:] / 255
+        img = (color_channel * alpha_channel +
+               255 * np.ones_like(color_channel) * (1 - alpha_channel))
+    elif alpha == 'blend_with_black':
+        color_channel = img[:, :, :3]
+        alpha_channel = img[:, :, 3:] / 255
+        img = color_channel * alpha_channel
+    return img
+
+
+def _read_image_cv2(path, dtype, color, alpha):
     if color:
-        color_option = cv2.IMREAD_COLOR
+        if alpha is None:
+            color_option = cv2.IMREAD_COLOR
+        else:
+            # Images with alpha channel are read as (H, W, 4) by cv2.imread.
+            # Images without alpha channel are read as (H, W, 3).
+            color_option = cv2.IMREAD_UNCHANGED
     else:
         color_option = cv2.IMREAD_GRAYSCALE
 
@@ -22,19 +50,28 @@ def _read_image_cv2(path, dtype, color):
         # reshape (H, W) -> (1, H, W)
         return img[np.newaxis].astype(dtype)
     else:
+        # alpha channel is included
+        if img.shape[-1] == 4:
+            img = _handle_four_channel_image(img, alpha)
         img = img[:, :, ::-1]  # BGR -> RGB
         img = img.transpose((2, 0, 1))  # HWC -> CHW
     return img.astype(dtype)
 
 
-def _read_image_pil(path, dtype, color):
+def _read_image_pil(path, dtype, color, alpha):
     f = Image.open(path)
     try:
         if color:
-            img = f.convert('RGB')
+            if f.mode == 'RGBA':
+                img = f.convert('RGBA')
+            else:
+                img = f.convert('RGB')
         else:
             img = f.convert('L')
         img = np.asarray(img, dtype=dtype)
+        if img.shape[-1] == 4:
+            img = _handle_four_channel_image(
+                img, alpha).astype(dtype, copy=False)
         img.flags.writeable = True
     finally:
         if hasattr(f, 'close'):
@@ -48,7 +85,7 @@ def _read_image_pil(path, dtype, color):
         return img.transpose((2, 0, 1))
 
 
-def read_image(path, dtype=np.float32, color=True):
+def read_image(path, dtype=np.float32, color=True, alpha=None):
     """Read an image from a file.
 
     This function reads an image from given file. The image is CHW format and
@@ -66,23 +103,32 @@ def read_image(path, dtype=np.float32, color=True):
             If :obj:`True`, the number of channels is three. In this case,
             the order of the channels is RGB. This is the default behaviour.
             If :obj:`False`, this function returns a grayscale image.
+        alpha (None or {'ignore', 'blend_with_white', 'blend_with_black'}): \
+            Choose how RGBA images are handled. By default, an error is raised.
+            Here are the other possible behaviors:
+
+            * `'ignore'`: Ignore alpha channel.
+            * `'blend_with_white'`: Blend RGB image multiplied by alpha on \
+                 a white image.
+            * `'blend_with_black'`: Blend RGB image multiplied by alpha on \
+                 a black image.
 
     Returns:
         ~numpy.ndarray: An image.
     """
     if chainer.config.cv_read_image_backend == 'cv2':
         if _cv2_available:
-            return _read_image_cv2(path, dtype, color)
+            return _read_image_cv2(path, dtype, color, alpha)
         else:
             warnings.warn(
                 'Although `chainer.config.cv_read_image_backend == "cv2"`, '
                 'cv2 is not found. As a fallback option, read_image uses '
                 'PIL. Either install cv2 or set '
                 '`chainer.global_config.cv_read_image_backend = "PIL"` '
                 'to suppress this warning.')
-            return _read_image_pil(path, dtype, color)
+            return _read_image_pil(path, dtype, color, alpha)
     elif chainer.config.cv_read_image_backend == 'PIL':
-        return _read_image_pil(path, dtype, color)
+        return _read_image_pil(path, dtype, color, alpha)
     else:
         raise ValueError('chainer.config.cv_read_image_backend should be '
                          'either "cv2" or "PIL".')

tests/utils_tests/image_tests/test_read_image.py

@@ -2,6 +2,8 @@
 import tempfile
 import unittest
 
+from PIL import Image
+
 import chainer
 from chainer import testing
 
@@ -15,13 +17,40 @@
     _cv2_available = False
 
 
-@testing.parameterize(*testing.product({
-    'size': [(48, 32)],
-    'color': [True, False],
-    'suffix': ['bmp', 'jpg', 'png'],
-    'dtype': [np.float32, np.uint8, bool],
-    'backend': ['cv2', 'PIL'],
-}))
+def _write_rgba_image(rgba, path):
+    rgba = rgba.transpose((1, 2, 0))
+    rgba = Image.fromarray(rgba, 'RGBA')
+    canvas = Image.new('RGBA', rgba.size, (255, 255, 255, 255))
+    # Paste the image onto the canvas, using it's alpha channel as mask
+    canvas.paste(rgba, mask=rgba)
+    canvas.save(path)
+
+
+def _create_parameters():
+    params = testing.product({
+        'size': [(48, 32)],
+        'dtype': [np.float32, np.uint8, bool]})
+    no_color_params = testing.product({
+        'color': [False],
+        'alpha': [None],
+        'suffix': ['bmp', 'jpg', 'png']})
+    no_alpha_params = testing.product({
+        'color': [True],
+        'alpha': [None],
+        'suffix': ['bmp', 'jpg', 'png']})
+    alpha_params = testing.product({
+        'color': [True],
+        'alpha': ['ignore', 'blend_with_white', 'blend_with_black'],
+        'suffix': ['png']  # writing alpha image with jpg encoding didn't work
+    })
+    params = testing.product_dict(
+        params,
+        no_color_params + no_alpha_params + alpha_params)
+    return params
+
+
+@testing.parameterize(*testing.product_dict(
+    _create_parameters(), [{'backend': 'cv2'}, {'backend': 'PIL'}]))
 class TestReadImage(unittest.TestCase):
 
     def setUp(self):
@@ -31,70 +60,83 @@ def setUp(self):
             suffix='.' + self.suffix, delete=False)
         self.path = self.file.name
 
-        if self.color:
-            self.img = np.random.randint(
-                0, 255, size=(3,) + self.size, dtype=np.uint8)
+        if self.alpha is None:
+            if self.color:
+                self.img = np.random.randint(
+                    0, 255, size=(3,) + self.size, dtype=np.uint8)
+            else:
+                self.img = np.random.randint(
+                    0, 255, size=(1,) + self.size, dtype=np.uint8)
+            write_image(self.img, self.path)
         else:
             self.img = np.random.randint(
-                0, 255, size=(1,) + self.size, dtype=np.uint8)
-        write_image(self.img, self.path)
+                0, 255, size=(4,) + self.size, dtype=np.uint8)
+            _write_rgba_image(self.img, self.path)
 
     def test_read_image_as_color(self):
-        img = read_image(self.path, dtype=self.dtype)
+        img = read_image(self.path, dtype=self.dtype, alpha=self.alpha)
 
         self.assertEqual(img.shape, (3,) + self.size)
         self.assertEqual(img.dtype, self.dtype)
 
-        if self.suffix in {'bmp', 'png'}:
+        if self.suffix in {'bmp', 'png'} and self.alpha is None:
             np.testing.assert_equal(
                 img,
                 np.broadcast_to(self.img, (3,) + self.size).astype(self.dtype))
 
     def test_read_image_as_grayscale(self):
-        img = read_image(self.path, dtype=self.dtype, color=False)
+        img = read_image(
+            self.path, dtype=self.dtype, color=False, alpha=self.alpha)
 
         self.assertEqual(img.shape, (1,) + self.size)
         self.assertEqual(img.dtype, self.dtype)
 
-        if self.suffix in {'bmp', 'png'} and not self.color:
+        if (self.suffix in {'bmp', 'png'}
+                and not self.color and self.alpha is None):
             np.testing.assert_equal(img, self.img.astype(self.dtype))
 
     def test_read_image_mutable(self):
-        img = read_image(self.path, dtype=self.dtype, color=self.color)
+        img = read_image(self.path, dtype=self.dtype, alpha=self.alpha)
         img[:] = 0
         np.testing.assert_equal(img, 0)
 
 
-@testing.parameterize(*testing.product({
-    'size': [(48, 32)],
-    'color': [True, False],
-    'suffix': ['bmp', 'jpg', 'png'],
-    'dtype': [np.float32, np.uint8, bool]}))
+@testing.parameterize(*_create_parameters())
 class TestReadImageDifferentBackends(unittest.TestCase):
 
     def setUp(self):
         self.file = tempfile.NamedTemporaryFile(
             suffix='.' + self.suffix, delete=False)
         self.path = self.file.name
 
-        if self.color:
-            self.img = np.random.randint(
-                0, 255, size=(3,) + self.size, dtype=np.uint8)
+        if self.alpha is None:
+            if self.color:
+                self.img = np.random.randint(
+                    0, 255, size=(3,) + self.size, dtype=np.uint8)
+            else:
+                self.img = np.random.randint(
+                    0, 255, size=(1,) + self.size, dtype=np.uint8)
+            write_image(self.img, self.path)
         else:
             self.img = np.random.randint(
-                0, 255, size=(1,) + self.size, dtype=np.uint8)
-        write_image(self.img, self.path)
+                0, 255, size=(4,) + self.size, dtype=np.uint8)
+            _write_rgba_image(self.img, self.path)
 
     @unittest.skipUnless(_cv2_available, 'cv2 is not installed')
     def test_read_image_different_backends_as_color(self):
         chainer.config.cv_read_image_backend = 'cv2'
-        cv2_img = read_image(self.path, dtype=self.dtype, color=True)
+        cv2_img = read_image(
+            self.path, dtype=self.dtype, color=self.color, alpha=self.alpha)
 
         chainer.config.cv_read_image_backend = 'PIL'
-        pil_img = read_image(self.path, dtype=self.dtype, color=True)
+        pil_img = read_image(
+            self.path, dtype=self.dtype, color=self.color, alpha=self.alpha)
 
         if self.suffix != 'jpg':
-            np.testing.assert_equal(cv2_img, pil_img)
+            if self.dtype == np.float32 and self.alpha is not None:
+                np.testing.assert_almost_equal(cv2_img, pil_img, decimal=4)
+            else:
+                np.testing.assert_equal(cv2_img, pil_img)
         else:
             # jpg decoders are differnet, so they produce different results
             assert np.mean(cv2_img == pil_img) > 0.99