Pillow

zenml.integrations.pillow special

Initialization of the Pillow integration.

PillowIntegration (Integration)

Definition of Pillow integration for ZenML.

Source code in zenml/integrations/pillow/__init__.py

class PillowIntegration(Integration):
    """Definition of Pillow integration for ZenML."""

    NAME = PILLOW
    REQUIREMENTS = ["Pillow>=9.2.0"]

    @classmethod
    def activate(cls) -> None:
        """Activates the integration."""
        from zenml.integrations.pillow import materializers  # noqa

activate() classmethod

Activates the integration.

Source code in zenml/integrations/pillow/__init__.py

@classmethod
def activate(cls) -> None:
    """Activates the integration."""
    from zenml.integrations.pillow import materializers  # noqa

materializers special

Initialization of the Pillow materializer.

pillow_image_materializer

Materializer for Pillow Image objects.

PillowImageMaterializer (BaseMaterializer)

Materializer for Image.Image objects.

This materializer takes a PIL image object and returns a PIL image object. It handles all the source image formats supported by PIL as listed here: https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html.

Source code in zenml/integrations/pillow/materializers/pillow_image_materializer.py

class PillowImageMaterializer(BaseMaterializer):
    """Materializer for Image.Image objects.

    This materializer takes a PIL image object and returns a PIL image object.
    It handles all the source image formats supported by PIL as listed here:
    https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html.
    """

    ASSOCIATED_TYPES = (Image.Image,)
    ASSOCIATED_ARTIFACT_TYPES = (DataArtifact,)

    def handle_input(self, data_type: Type[Image.Image]) -> Image.Image:
        """Read from artifact store.

        Args:
            data_type: An Image.Image type.

        Returns:
            An Image.Image object.
        """
        super().handle_input(data_type)
        files = io_utils.find_files(
            self.artifact.uri, f"{DEFAULT_IMAGE_FILENAME}.*"
        )
        filepath = [file for file in files if not fileio.isdir(file)][0]

        # # FAILING OPTION 1: temporary directory
        # # create a temporary folder
        temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
        temp_file = os.path.join(
            temp_dir.name,
            f"{DEFAULT_IMAGE_FILENAME}{os.path.splitext(filepath)[1]}",
        )

        # copy from artifact store to temporary file
        fileio.copy(filepath, temp_file)
        return Image.open(temp_file)

    def handle_return(self, image: Image.Image) -> None:
        """Write to artifact store.

        Args:
            image: An Image.Image object.
        """
        # # FAILING OPTION 1: temporary directory
        super().handle_return(image)
        temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
        file_extension = image.format or DEFAULT_IMAGE_EXTENSION
        full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
        temp_image_path = os.path.join(temp_dir.name, full_filename)

        # save the image in a temporary directory
        image.save(temp_image_path)

        # copy the saved image to the artifact store
        artifact_store_path = os.path.join(self.artifact.uri, full_filename)
        io_utils.copy(temp_image_path, artifact_store_path, overwrite=True)  # type: ignore[attr-defined]
        temp_dir.cleanup()

handle_input(self, data_type)

Read from artifact store.

Parameters:

Name	Type	Description	Default
data_type	Type[PIL.Image.Image]	An Image.Image type.	required

Returns:

Type	Description
Image	An Image.Image object.

Source code in zenml/integrations/pillow/materializers/pillow_image_materializer.py

def handle_input(self, data_type: Type[Image.Image]) -> Image.Image:
    """Read from artifact store.

    Args:
        data_type: An Image.Image type.

    Returns:
        An Image.Image object.
    """
    super().handle_input(data_type)
    files = io_utils.find_files(
        self.artifact.uri, f"{DEFAULT_IMAGE_FILENAME}.*"
    )
    filepath = [file for file in files if not fileio.isdir(file)][0]

    # # FAILING OPTION 1: temporary directory
    # # create a temporary folder
    temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
    temp_file = os.path.join(
        temp_dir.name,
        f"{DEFAULT_IMAGE_FILENAME}{os.path.splitext(filepath)[1]}",
    )

    # copy from artifact store to temporary file
    fileio.copy(filepath, temp_file)
    return Image.open(temp_file)

handle_return(self, image)

Write to artifact store.

Parameters:

Name	Type	Description	Default
image	Image	An Image.Image object.	required

Source code in zenml/integrations/pillow/materializers/pillow_image_materializer.py

def handle_return(self, image: Image.Image) -> None:
    """Write to artifact store.

    Args:
        image: An Image.Image object.
    """
    # # FAILING OPTION 1: temporary directory
    super().handle_return(image)
    temp_dir = tempfile.TemporaryDirectory(prefix="zenml-temp-")
    file_extension = image.format or DEFAULT_IMAGE_EXTENSION
    full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
    temp_image_path = os.path.join(temp_dir.name, full_filename)

    # save the image in a temporary directory
    image.save(temp_image_path)

    # copy the saved image to the artifact store
    artifact_store_path = os.path.join(self.artifact.uri, full_filename)
    io_utils.copy(temp_image_path, artifact_store_path, overwrite=True)  # type: ignore[attr-defined]
    temp_dir.cleanup()