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: ClassVar[Tuple[Type[Any], ...]] = (Image.Image,)
    ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA

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

        Args:
            data_type: An Image.Image type.

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

        # 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 save(self, image: Image.Image) -> None:
        """Write to artifact store.

        Args:
            image: An Image.Image object.
        """
        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.uri, full_filename)
        io_utils.copy(temp_image_path, artifact_store_path, overwrite=True)  # type: ignore[attr-defined]
        temp_dir.cleanup()

    def save_visualizations(
        self, image: Image.Image
    ) -> Dict[str, VisualizationType]:
        """Finds and saves the given image as a visualization.

        Args:
            image: The image to save as a visualization.

        Returns:
            A dictionary of visualization URIs and their types.
        """
        file_extension = image.format or DEFAULT_IMAGE_EXTENSION
        full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
        artifact_store_path = os.path.join(self.uri, full_filename)
        return {artifact_store_path: VisualizationType.IMAGE}

    def extract_metadata(
        self, image: Image.Image
    ) -> Dict[str, "MetadataType"]:
        """Extract metadata from the given `Image` object.

        Args:
            image: The `Image` object to extract metadata from.

        Returns:
            The extracted metadata as a dictionary.
        """
        metadata = {
            "width": image.width,
            "height": image.height,
            "mode": str(image.mode),
        }
        if hasattr(image, "filename"):
            metadata["original_filename"] = str(image.filename)
        return metadata  # type: ignore[return-value]

extract_metadata(self, image)

Extract metadata from the given Image object.

Parameters:

Name	Type	Description	Default
image	Image	The Image object to extract metadata from.	required

Returns:

Type	Description
Dict[str, MetadataType]	The extracted metadata as a dictionary.

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

def extract_metadata(
    self, image: Image.Image
) -> Dict[str, "MetadataType"]:
    """Extract metadata from the given `Image` object.

    Args:
        image: The `Image` object to extract metadata from.

    Returns:
        The extracted metadata as a dictionary.
    """
    metadata = {
        "width": image.width,
        "height": image.height,
        "mode": str(image.mode),
    }
    if hasattr(image, "filename"):
        metadata["original_filename"] = str(image.filename)
    return metadata  # type: ignore[return-value]

load(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 load(self, data_type: Type[Image.Image]) -> Image.Image:
    """Read from artifact store.

    Args:
        data_type: An Image.Image type.

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

    # 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)

save(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 save(self, image: Image.Image) -> None:
    """Write to artifact store.

    Args:
        image: An Image.Image object.
    """
    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.uri, full_filename)
    io_utils.copy(temp_image_path, artifact_store_path, overwrite=True)  # type: ignore[attr-defined]
    temp_dir.cleanup()

save_visualizations(self, image)

Finds and saves the given image as a visualization.

Parameters:

Name	Type	Description	Default
image	Image	The image to save as a visualization.	required

Returns:

Type	Description
Dict[str, zenml.enums.VisualizationType]	A dictionary of visualization URIs and their types.

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

def save_visualizations(
    self, image: Image.Image
) -> Dict[str, VisualizationType]:
    """Finds and saves the given image as a visualization.

    Args:
        image: The image to save as a visualization.

    Returns:
        A dictionary of visualization URIs and their types.
    """
    file_extension = image.format or DEFAULT_IMAGE_EXTENSION
    full_filename = f"{DEFAULT_IMAGE_FILENAME}.{file_extension}"
    artifact_store_path = os.path.join(self.uri, full_filename)
    return {artifact_store_path: VisualizationType.IMAGE}