[io] Place transparent images on a white background in read_img_as_numpy

doctr/io/image/base.py

@@ -3,16 +3,56 @@
 # This program is licensed under the Apache License 2.0.
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
+from io import BytesIO
 from pathlib import Path
 
 import cv2
 import numpy as np
+from PIL import Image, ImageOps
 
 from doctr.utils.common_types import AbstractFile
 
 __all__ = ["read_img_as_numpy"]
 
 
+def _pil_to_rgb_array(pil_img: Image.Image) -> np.ndarray:
+    """Convert a PIL image to an RGB array, with its EXIF orientation applied (like `cv2.imdecode` does)
+    and transparent pixels placed on a white background
+
+    Args:
+        pil_img: a PIL image
+
+    Returns:
+        the image as a writable numpy array of shape H x W x 3 in np.uint8
+    """
+    has_transparency = pil_img.has_transparency_data
+    ImageOps.exif_transpose(pil_img, in_place=True)
+    if not has_transparency:
+        return np.array(pil_img.convert("RGB"))
+    rgba = np.asarray(pil_img if pil_img.mode == "RGBA" else pil_img.convert("RGBA"))
+    rgb = cv2.cvtColor(rgba, cv2.COLOR_RGBA2RGB)
+    alpha = rgba[..., 3]
+    # Fully opaque images (e.g. screenshots) need no compositing
+    if alpha.min() == 255:
+        return rgb
+    # Dropping the alpha channel would reveal the color stored under the transparent pixels (often black),
+    # so blend with white in place: 255 - (255 - rgb) * alpha / 255
+    cv2.bitwise_not(rgb, dst=rgb)
+    cv2.multiply(rgb, cv2.cvtColor(np.ascontiguousarray(alpha), cv2.COLOR_GRAY2RGB), dst=rgb, scale=1 / 255)
+    return cv2.bitwise_not(rgb, dst=rgb)
+
+
+def _decode_transparent_img(content: bytes) -> np.ndarray | None:
+    """Decode an image with transparency over a white background, in RGB order (None if it has no transparency)"""
+    try:
+        with Image.open(BytesIO(content)) as pil_img:
+            if not pil_img.has_transparency_data:
+                return None
+            return _pil_to_rgb_array(pil_img)
+    except (OSError, ValueError, Image.DecompressionBombError):  # left to OpenCV
+        return None
+
+
 def read_img_as_numpy(
     file: AbstractFile,
     output_size: tuple[int, int] | None = None,
@@ -29,26 +69,29 @@
         rgb_output: whether the output ndarray channel order should be RGB instead of BGR.
 
     Returns:
-        the page decoded as numpy ndarray of shape H x W x 3
+        the page decoded as numpy ndarray of shape H x W x 3, transparent images are placed on a white background
     """
     if isinstance(file, (str, Path)):
         if not Path(file).is_file():
             raise FileNotFoundError(f"unable to access {file}")
-        _file: np.ndarray = np.frombuffer(Path(file).read_bytes(), np.uint8)
-        img = cv2.imdecode(_file, cv2.IMREAD_COLOR)
+        content = Path(file).read_bytes()
     elif isinstance(file, bytes):
-        _file = np.frombuffer(file, np.uint8)
-        img = cv2.imdecode(_file, cv2.IMREAD_COLOR)
+        content = file
     else:
         raise TypeError("unsupported object type for argument 'file'")
 
+    img = _decode_transparent_img(content)
+    is_rgb = img is not None
+    if img is None:
+        img = cv2.imdecode(np.frombuffer(content, np.uint8), cv2.IMREAD_COLOR)
+
     # Validity check
     if img is None:
         raise ValueError("unable to read file.")
     # Resizing
     if isinstance(output_size, tuple):
         img = cv2.resize(img, output_size[::-1], interpolation=cv2.INTER_LINEAR)
-    # Switch the channel order
-    if rgb_output:
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    # Switch the channel order (OpenCV decodes as BGR, Pillow as RGB)
+    if rgb_output != is_rgb:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB if rgb_output else cv2.COLOR_RGB2BGR)
     return img

doctr/io/image/pytorch.py

@@ -7,11 +7,13 @@
 
 import numpy as np
 import torch
-from PIL import Image, ImageOps
+from PIL import Image
 from torchvision.transforms.functional import to_tensor
 
 from doctr.utils.common_types import AbstractPath
 
+from .base import _pil_to_rgb_array
+
 __all__ = ["tensor_from_pil", "read_img_as_tensor", "decode_img_as_tensor", "tensor_from_numpy", "get_img_shape"]
 
 
@@ -47,9 +49,8 @@ def read_img_as_tensor(img_path: AbstractPath, dtype: torch.dtype = torch.float3
         raise ValueError("insupported value for dtype")
 
     with Image.open(img_path, mode="r") as pil_img:
-        # Apply the EXIF orientation before decoding, to stay consistent with `read_img_as_numpy`
-        # (OpenCV applies EXIF orientation automatically, PIL does not)
-        return tensor_from_pil(ImageOps.exif_transpose(pil_img).convert("RGB"), dtype)
+        # Apply the EXIF orientation and place transparent pixels on white, like `read_img_as_numpy`
+        return tensor_from_numpy(_pil_to_rgb_array(pil_img), dtype)
 
 
 def decode_img_as_tensor(img_content: bytes, dtype: torch.dtype = torch.float32) -> torch.Tensor:
@@ -66,7 +67,7 @@ def decode_img_as_tensor(img_content: bytes, dtype: torch.dtype = torch.float32)
         raise ValueError("insupported value for dtype")
 
     with Image.open(BytesIO(img_content), mode="r") as pil_img:
-        return tensor_from_pil(ImageOps.exif_transpose(pil_img).convert("RGB"), dtype)
+        return tensor_from_numpy(_pil_to_rgb_array(pil_img), dtype)
 
 
 def tensor_from_numpy(npy_img: np.ndarray, dtype: torch.dtype = torch.float32) -> torch.Tensor:

tests/common/test_io.py

@@ -77,6 +77,30 @@ def test_read_img_as_numpy(tmpdir_factory, mock_pdf):
     resized_page = io.read_img_as_numpy(tmp_path, target_size)
     assert resized_page.shape[:2] == target_size
 
+    # Transparency: placed on a white background, not on the color stored under the transparent pixels
+    rgba = np.zeros((40, 60, 4), dtype=np.uint8)
+    rgba[10:30, 20:40] = (255, 0, 0, 255)  # opaque red square on a transparent black background
+    png_stream = BytesIO()
+    Image.fromarray(rgba, "RGBA").save(png_stream, format="PNG")
+    page = io.read_img_as_numpy(png_stream.getvalue())
+    assert page.shape == (40, 60, 3)
+    assert np.all(page[0, 0] == 255) and np.all(page[20, 30] == (255, 0, 0))
+    assert np.all(io.read_img_as_numpy(png_stream.getvalue(), rgb_output=False)[20, 30] == (0, 0, 255))
+    # Semi-transparent pixels are blended with white
+    rgba[0, 0] = (0, 0, 0, 128)
+    png_stream = BytesIO()
+    Image.fromarray(rgba, "RGBA").save(png_stream, format="PNG")
+    assert np.all(io.read_img_as_numpy(png_stream.getvalue())[0, 0] == 127)
+
+    # Fully opaque RGBA images are kept as is
+    rgba[..., 3] = 255
+    png_stream = BytesIO()
+    Image.fromarray(rgba, "RGBA").save(png_stream, format="PNG")
+    page = io.read_img_as_numpy(png_stream.getvalue())
+    assert page.shape == (40, 60, 3) and page.flags.c_contiguous
+    assert np.all(page == rgba[..., :3])
+    assert np.all(io.read_img_as_numpy(png_stream.getvalue(), rgb_output=False) == rgba[..., 2::-1])
+
 
 def test_read_html():
     url = "https://www.google.com"
@@ -112,3 +136,10 @@ def test_read_img_as_numpy_exif_orientation(tmpdir_factory):
     assert io.read_img_as_numpy(path).shape == (60, 40, 3)
     with open(path, "rb") as f:
         assert io.read_img_as_numpy(f.read()).shape == (60, 40, 3)
+
+    # Same for an image with transparency
+    path = str(folder.join("exif_o6.png"))
+    Image.fromarray(np.zeros((40, 60, 4), dtype=np.uint8), "RGBA").save(path, format="PNG", exif=exif)
+    page = io.read_img_as_numpy(path)
+    assert page.shape == (60, 40, 3)
+    assert np.all(page[0, 0] == 255)

tests/pytorch/test_io_image_pt.py

@@ -68,3 +68,20 @@ def test_read_img_as_tensor_exif_orientation(tmpdir_factory):
     with open(path, "rb") as f:
         img_stream = decode_img_as_tensor(f.read())
     assert img_stream.shape == (3, 60, 40)
+
+
+def test_img_as_tensor_transparency(tmpdir_factory):
+    # Opaque red square on a transparent black background, placed on white like `read_img_as_numpy` does
+    rgba = np.zeros((40, 60, 4), dtype=np.uint8)
+    rgba[10:30, 20:40] = (255, 0, 0, 255)
+    path = str(tmpdir_factory.mktemp("images").join("transparent.png"))
+    Image.fromarray(rgba, "RGBA").save(path, format="PNG")
+    with open(path, "rb") as f:
+        img_stream = f.read()
+
+    for img in (read_img_as_tensor(path, dtype=torch.uint8), decode_img_as_tensor(img_stream, dtype=torch.uint8)):
+        assert img.shape == (3, 40, 60)
+        assert torch.all(img[:, 0, 0] == 255)
+        assert img[:, 20, 30].tolist() == [255, 0, 0]
+    assert torch.all(read_img_as_tensor(path)[:, 0, 0] == 1)
+    assert torch.all(decode_img_as_tensor(img_stream)[:, 0, 0] == 1)