Repository: OpenMined/PySyft

Issue:
Add PIL.Image.Image for torchvision
## Description
Add `PIL.Image.Image` return type support for `torchvision`.

## Definition of Done
`torchvision/allowlist.py` can enable `to_pil_image` and `to_grayscale`.


Discussion:

Resolving patch:
diff --git a/src/syft/lib/PIL/__init__.py b/src/syft/lib/PIL/__init__.py
new file mode 100644
--- /dev/null
+++ b/src/syft/lib/PIL/__init__.py
@@ -0,0 +1,48 @@
+# stdlib
+import functools
+from typing import Any as TypeAny
+from typing import List as TypeList
+from typing import Tuple as TypeTuple
+
+# third party
+import PIL
+
+# syft relative
+from . import image  # noqa: 401
+from ...ast import add_classes
+from ...ast import add_methods
+from ...ast import add_modules
+from ...ast.globals import Globals
+from ..util import generic_update_ast
+
+LIB_NAME = "PIL"
+PACKAGE_SUPPORT = {"lib": LIB_NAME}
+
+
+def create_ast(client: TypeAny = None) -> Globals:
+    ast = Globals(client)
+
+    modules: TypeList[TypeTuple[str, TypeAny]] = [
+        ("PIL", PIL),
+        ("PIL.Image", PIL.Image),
+    ]
+
+    classes: TypeList[TypeTuple[str, str, TypeAny]] = [
+        ("PIL.Image.Image", "PIL.Image.Image", PIL.Image.Image)
+    ]
+
+    methods: TypeList[TypeTuple[str, str]] = []
+
+    add_modules(ast, modules)
+    add_classes(ast, classes)
+    add_methods(ast, methods)
+
+    for klass in ast.classes:
+        klass.create_pointer_class()
+        klass.create_send_method()
+        klass.create_storable_object_attr_convenience_methods()
+
+    return ast
+
+
+update_ast = functools.partial(generic_update_ast, LIB_NAME, create_ast)
diff --git a/src/syft/lib/PIL/image.py b/src/syft/lib/PIL/image.py
new file mode 100644
--- /dev/null
+++ b/src/syft/lib/PIL/image.py
@@ -0,0 +1,35 @@
+# third party
+import PIL
+import numpy as np
+import torch
+import torchvision
+
+# syft relative
+from ...generate_wrapper import GenerateWrapper
+from ...lib.torch.tensor_util import protobuf_tensor_deserializer
+from ...lib.torch.tensor_util import protobuf_tensor_serializer
+from ...proto.lib.torch.tensor_pb2 import TensorData
+
+
+def object2proto(obj: PIL.Image.Image) -> TensorData:
+    image_tensor = torch.Tensor(np.array(obj))
+    tensor_proto = protobuf_tensor_serializer(image_tensor)
+
+    return tensor_proto
+
+
+def proto2object(proto: TensorData) -> PIL.Image.Image:
+    image_tensor = protobuf_tensor_deserializer(proto)
+    if image_tensor.dim() == 3:
+        image_tensor = image_tensor.permute(2, 0, 1)
+    image_obj = torchvision.transforms.functional.to_pil_image(image_tensor)
+    return image_obj
+
+
+GenerateWrapper(
+    wrapped_type=PIL.Image.Image,
+    import_path="PIL.Image.Image",
+    protobuf_scheme=TensorData,
+    type_object2proto=object2proto,
+    type_proto2object=proto2object,
+)
diff --git a/src/syft/lib/torchvision/allowlist.py b/src/syft/lib/torchvision/allowlist.py
--- a/src/syft/lib/torchvision/allowlist.py
+++ b/src/syft/lib/torchvision/allowlist.py
@@ -466,18 +466,12 @@
     "min_version": "0.8.0"
     # Torch 1.6 expects input to be PIL image, so minimum version as 0.7 (Torch 1.7.0)
 }
-
-# TODO: Fix when we have PIL support
-# Issue: https://github.com/OpenMined/PySyft/issues/5324
-# Following takes  PIL image as input, currently not supported
-# allowlist["torchvision.transforms.functional.to_grayscale"] = {
-#    "return_type" : "torch.Tensor",
-#    "test_parameters" : "(npy_array)"
-# }
-
-# Following converts image to PIL image, currently not supported
-# allowlist["torchvision.transforms.functional.to_pil_image"] = "PIL.Image.Image"
-
+allowlist["torchvision.transforms.functional.to_grayscale"] = {
+    "return_type": "PIL.Image.Image"
+}
+allowlist["torchvision.transforms.functional.to_pil_image"] = {
+    "return_type": "PIL.Image.Image",
+}
 allowlist["torchvision.transforms.functional.to_tensor"] = {
     "return_type": "torch.Tensor",
 }


Test patch:
diff --git a/tests/syft/lib/PIL/__init__.py b/tests/syft/lib/PIL/__init__.py
new file mode 100644
diff --git a/tests/syft/lib/PIL/pil_test.py b/tests/syft/lib/PIL/pil_test.py
new file mode 100644
--- /dev/null
+++ b/tests/syft/lib/PIL/pil_test.py
@@ -0,0 +1,44 @@
+# third party
+import pytest
+
+# syft absolute
+import syft as sy
+from syft.grid.duet.ui import LOGO_URL
+
+
+@pytest.mark.vendor(lib="PIL")
+def test_send_and_get() -> None:
+    # third party
+    import PIL
+
+    sy.load("PIL")
+
+    data_owner = sy.VirtualMachine().get_root_client()
+
+    im = PIL.Image.open(LOGO_URL)
+    remote_im = im.send(data_owner)
+    received_im = remote_im.get()
+
+    assert PIL.ImageChops.difference(im, received_im).getbbox() is None
+
+
+@pytest.mark.vendor(lib="PIL")
+def test_remote_create() -> None:
+    # third party
+    import PIL
+    import numpy as np
+    import torch
+
+    sy.load("PIL")
+
+    data_owner = sy.VirtualMachine().get_root_client()
+    remote_torchvision = data_owner.torchvision
+
+    im = PIL.Image.open(LOGO_URL)
+    im_array = np.array(im)
+    im_tensor = torch.Tensor(im_array).permute(2, 0, 1)
+    remote_tensor = im_tensor.send(data_owner)
+    remote_im = remote_torchvision.transforms.functional.to_pil_image(remote_tensor)
+    received_im = remote_im.get()
+
+    assert PIL.ImageChops.difference(im, received_im).getbbox() is None