Add util for Image.fromarray to deal with mode deprecation

Author: scotts

test/common_utils.py

@@ -12,16 +12,15 @@
 from subprocess import CalledProcessError, check_output, STDOUT
 
 import numpy as np
-import PIL.Image
 import pytest
 import torch
 import torch.testing
-from PIL import Image
 
 from torch.testing._comparison import BooleanPair, NonePair, not_close_error_metas, NumberPair, TensorLikePair
 from torchvision import io, tv_tensors
 from torchvision.transforms._functional_tensor import _max_value as get_max_value
 from torchvision.transforms.v2.functional import to_image, to_pil_image
+from torchvision.utils import _Image_fromarray
 
 
 IN_OSS_CI = any(os.getenv(var) == "true" for var in ["CIRCLECI", "GITHUB_ACTIONS"])
@@ -147,7 +146,7 @@ def _create_data(height=3, width=3, channels=3, device="cpu"):
     if channels == 1:
         mode = "L"
         data = data[..., 0]
-    pil_img = Image.fromarray(data, mode=mode)
+    pil_img = _Image_fromarray(data, mode=mode)
     return tensor, pil_img
 
 

test/test_transforms.py

@@ -13,6 +13,7 @@
 import torchvision.transforms.functional as F
 from PIL import Image
 from torch._utils_internal import get_file_path_2
+from torchvision.utils import _Image_fromarray
 
 try:
     import accimage
@@ -654,7 +655,7 @@ def test_1_channel_float_tensor_to_pil_image(self):
         img_F_mode = transforms.ToPILImage(mode="F")(img_data)
         assert img_F_mode.mode == "F"
         torch.testing.assert_close(
-            np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode="F")), np.array(img_F_mode)
+            np.array(_Image_fromarray(img_data.squeeze(0).numpy(), mode="F")), np.array(img_F_mode)
         )
 
     @pytest.mark.parametrize("with_mode", [False, True])
@@ -895,7 +896,7 @@ def test_adjust_brightness():
     x_shape = [2, 2, 3]
     x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
     x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
-    x_pil = Image.fromarray(x_np, mode="RGB")
+    x_pil = _Image_fromarray(x_np, mode="RGB")
 
     # test 0
     y_pil = F.adjust_brightness(x_pil, 1)
@@ -921,7 +922,7 @@ def test_adjust_contrast():
     x_shape = [2, 2, 3]
     x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
     x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
-    x_pil = Image.fromarray(x_np, mode="RGB")
+    x_pil = _Image_fromarray(x_np, mode="RGB")
 
     # test 0
     y_pil = F.adjust_contrast(x_pil, 1)
@@ -947,7 +948,7 @@ def test_adjust_hue():
     x_shape = [2, 2, 3]
     x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
     x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
-    x_pil = Image.fromarray(x_np, mode="RGB")
+    x_pil = _Image_fromarray(x_np, mode="RGB")
 
     with pytest.raises(ValueError):
         F.adjust_hue(x_pil, -0.7)
@@ -1029,7 +1030,7 @@ def test_adjust_sharpness():
         117,
     ]
     x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
-    x_pil = Image.fromarray(x_np, mode="RGB")
+    x_pil = _Image_fromarray(x_np, mode="RGB")
 
     # test 0
     y_pil = F.adjust_sharpness(x_pil, 1)
@@ -1152,7 +1153,7 @@ def test_adjust_sharpness():
     x_shape = [2, 2, 3]
     x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
     x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
-    x_pil = Image.fromarray(x_np, mode="RGB")
+    x_pil = _Image_fromarray(x_np, mode="RGB")
     x_th = torch.tensor(x_np.transpose(2, 0, 1))
     y_pil = F.adjust_sharpness(x_pil, 2)
     y_np = np.array(y_pil).transpose(2, 0, 1)
@@ -1164,7 +1165,7 @@ def test_adjust_gamma():
     x_shape = [2, 2, 3]
     x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
     x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
-    x_pil = Image.fromarray(x_np, mode="RGB")
+    x_pil = _Image_fromarray(x_np, mode="RGB")
 
     # test 0
     y_pil = F.adjust_gamma(x_pil, 1)
@@ -1190,7 +1191,7 @@ def test_adjusts_L_mode():
     x_shape = [2, 2, 3]
     x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
     x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
-    x_rgb = Image.fromarray(x_np, mode="RGB")
+    x_rgb = _Image_fromarray(x_np, mode="RGB")
 
     x_l = x_rgb.convert("L")
     assert F.adjust_brightness(x_l, 2).mode == "L"
@@ -1320,7 +1321,7 @@ def test_to_grayscale():
     x_shape = [2, 2, 3]
     x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
     x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
-    x_pil = Image.fromarray(x_np, mode="RGB")
+    x_pil = _Image_fromarray(x_np, mode="RGB")
     x_pil_2 = x_pil.convert("L")
     gray_np = np.array(x_pil_2)
 
@@ -1769,7 +1770,7 @@ def test_color_jitter():
     x_shape = [2, 2, 3]
     x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
     x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
-    x_pil = Image.fromarray(x_np, mode="RGB")
+    x_pil = _Image_fromarray(x_np, mode="RGB")
     x_pil_2 = x_pil.convert("L")
 
     for _ in range(10):

torchvision/datasets/mnist.py

@@ -11,8 +11,8 @@
 
 import numpy as np
 import torch
-from PIL import Image
 
+from ..utils import _Image_fromarray
 from .utils import _flip_byte_order, check_integrity, download_and_extract_archive, extract_archive, verify_str_arg
 from .vision import VisionDataset
 
@@ -140,7 +140,7 @@ def __getitem__(self, index: int) -> tuple[Any, Any]:
 
         # doing this so that it is consistent with all other datasets
         # to return a PIL Image
-        img = Image.fromarray(img.numpy(), mode="L")
+        img = _Image_fromarray(img.numpy(), mode="L")
 
         if self.transform is not None:
             img = self.transform(img)
@@ -478,7 +478,7 @@ def download(self) -> None:
     def __getitem__(self, index: int) -> tuple[Any, Any]:
         # redefined to handle the compat flag
         img, target = self.data[index], self.targets[index]
-        img = Image.fromarray(img.numpy(), mode="L")
+        img = _Image_fromarray(img.numpy(), mode="L")
         if self.transform is not None:
             img = self.transform(img)
         if self.compat:

torchvision/datasets/semeion.py

@@ -3,8 +3,8 @@
 from typing import Any, Callable, Optional, Union
 
 import numpy as np
-from PIL import Image
 
+from ..utils import _Image_fromarray
 from .utils import check_integrity, download_url
 from .vision import VisionDataset
 
@@ -64,7 +64,7 @@ def __getitem__(self, index: int) -> tuple[Any, Any]:
 
         # doing this so that it is consistent with all other datasets
         # to return a PIL Image
-        img = Image.fromarray(img, mode="L")
+        img = _Image_fromarray(img, mode="L")
 
         if self.transform is not None:
             img = self.transform(img)

torchvision/datasets/usps.py

@@ -3,8 +3,8 @@
 from typing import Any, Callable, Optional, Union
 
 import numpy as np
-from PIL import Image
 
+from ..utils import _Image_fromarray
 from .utils import download_url
 from .vision import VisionDataset
 
@@ -82,7 +82,7 @@ def __getitem__(self, index: int) -> tuple[Any, Any]:
 
         # doing this so that it is consistent with all other datasets
         # to return a PIL Image
-        img = Image.fromarray(img, mode="L")
+        img = _Image_fromarray(img, mode="L")
 
         if self.transform is not None:
             img = self.transform(img)

torchvision/transforms/functional.py

@@ -16,7 +16,7 @@
 except ImportError:
     accimage = None
 
-from ..utils import _log_api_usage_once
+from ..utils import _Image_fromarray, _log_api_usage_once
 from . import _functional_pil as F_pil, _functional_tensor as F_t
 
 
@@ -321,7 +321,7 @@ def to_pil_image(pic, mode=None):
     if mode is None:
         raise TypeError(f"Input type {npimg.dtype} is not supported")
 
-    return Image.fromarray(npimg, mode=mode)
+    return _Image_fromarray(npimg, mode=mode)
 
 
 def normalize(tensor: Tensor, mean: list[float], std: list[float], inplace: bool = False) -> Tensor:

torchvision/utils.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 import torch
+from packaging import version
 from PIL import Image, ImageColor, ImageDraw, ImageFont
 
 
@@ -174,6 +175,21 @@ def dashed_line(self, xy, fill=None, width=0, joint=None, dash_length=5, space_l
                 current_dash = not current_dash
 
 
+def _Image_fromarray(
+    obj: Union[torch.Tensor, np.ndarray],
+    mode: Optional[str],
+) -> Image:
+    """
+    A wrapper around PIL.Image.fromarray to mitigate the deprecation of the
+    mode paramter. See:
+      https://pillow.readthedocs.io/en/stable/releasenotes/11.3.0.html#image-fromarray-mode-parameter
+    """
+    if version.parse(Image.__version__) >= version.parse("11.3.0"):
+        return Image.fromarray(obj)
+    else:
+        return Image.fromarray(obj, mode)
+
+
 @torch.no_grad()
 def save_image(
     tensor: Union[torch.Tensor, list[torch.Tensor]],