diff --git a/docarray/typing/bytes/image_bytes.py b/docarray/typing/bytes/image_bytes.py index f8aaf3c6caf..e5d1d400d99 100644 --- a/docarray/typing/bytes/image_bytes.py +++ b/docarray/typing/bytes/image_bytes.py @@ -10,6 +10,7 @@ from docarray.utils._internal.misc import import_library if TYPE_CHECKING: + from PIL import Image as PILImage from pydantic.fields import BaseConfig, ModelField from docarray.proto import NodeProto @@ -43,6 +44,38 @@ def _to_node_protobuf(self: T) -> 'NodeProto': return NodeProto(blob=self, type=self._proto_type_name) + def load_pil( + self, + ) -> 'PILImage.Image': + """ + Load the image from the bytes into a `PIL.Image.Image` instance + + --- + + ```python + from pydantic import parse_obj_as + + from docarray import BaseDoc + from docarray.typing import ImageUrl + + img_url = "https://upload.wikimedia.org/wikipedia/commons/8/80/Dag_Sebastian_Ahlander_at_G%C3%B6teborg_Book_Fair_2012b.jpg" + + img_url = parse_obj_as(ImageUrl, img_url) + img = img_url.load_pil() + + from PIL.Image import Image + + assert isinstance(img, Image) + ``` + + --- + :return: a Pillow image + """ + PIL = import_library('PIL', raise_error=True) # noqa: F841 + from PIL import Image as PILImage + + return PILImage.open(BytesIO(self)) + def load( self, width: Optional[int] = None, @@ -89,13 +122,7 @@ class MyDoc(BaseDoc): :return: np.ndarray representing the image as RGB values """ - if TYPE_CHECKING: - from PIL import Image as PILImage - else: - PIL = import_library('PIL', raise_error=True) # noqa: F841 - from PIL import Image as PILImage - - raw_img = PILImage.open(BytesIO(self)) + raw_img = self.load_pil() if width or height: new_width = width or raw_img.width new_height = height or raw_img.height diff --git a/docarray/typing/url/image_url.py b/docarray/typing/url/image_url.py index 4c25e68094e..7fa0f58af80 100644 --- a/docarray/typing/url/image_url.py +++ b/docarray/typing/url/image_url.py @@ -8,9 +8,11 @@ from docarray.utils._internal.misc import is_notebook if TYPE_CHECKING: + from PIL import Image as PILImage from pydantic import BaseConfig from pydantic.fields import ModelField + T = TypeVar('T', bound='ImageUrl') IMAGE_FILE_FORMATS = ('png', 'jpeg', 'jpg') @@ -39,6 +41,35 @@ def validate( ) return cls(str(url), scheme=None) + def load_pil(self, timeout: Optional[float] = None) -> 'PILImage.Image': + """ + Load the image from the bytes into a `PIL.Image.Image` instance + + --- + + ```python + from pydantic import parse_obj_as + + from docarray import BaseDoc + from docarray.typing import ImageUrl + + img_url = "https://upload.wikimedia.org/wikipedia/commons/8/80/Dag_Sebastian_Ahlander_at_G%C3%B6teborg_Book_Fair_2012b.jpg" + + img_url = parse_obj_as(ImageUrl, img_url) + img = img_url.load_pil() + + from PIL.Image import Image + + assert isinstance(img, Image) + ``` + + --- + :return: a Pillow image + """ + from docarray.typing.bytes.image_bytes import ImageBytes + + return ImageBytes(self.load_bytes(timeout=timeout)).load_pil() + def load( self, width: Optional[int] = None, diff --git a/tests/units/typing/url/test_image_url.py b/tests/units/typing/url/test_image_url.py index b425e498ab2..b0a6048e69a 100644 --- a/tests/units/typing/url/test_image_url.py +++ b/tests/units/typing/url/test_image_url.py @@ -2,6 +2,7 @@ import urllib import numpy as np +import PIL import pytest from PIL import Image from pydantic.tools import parse_obj_as, schema_json_of @@ -66,6 +67,23 @@ def test_load(image_format, path_to_img): assert isinstance(tensor, np.ndarray) +@pytest.mark.slow +@pytest.mark.internet +@pytest.mark.parametrize( + 'image_format,path_to_img', + [ + ('png', IMAGE_PATHS['png']), + ('jpg', IMAGE_PATHS['jpg']), + ('jpeg', IMAGE_PATHS['jpeg']), + ('remote-jpg', REMOTE_JPG), + ], +) +def test_load_pil(image_format, path_to_img): + url = parse_obj_as(ImageUrl, path_to_img) + img = url.load_pil() + assert isinstance(img, PIL.Image.Image) + + @pytest.mark.slow @pytest.mark.internet @pytest.mark.parametrize(