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)
artifact_store_path = artifact_store_path.replace("\\", "/")
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 |
PIL.Image.Image |
The |
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 |
---|---|
PIL.Image.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 |
PIL.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 |
PIL.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)
artifact_store_path = artifact_store_path.replace("\\", "/")
return {artifact_store_path: VisualizationType.IMAGE}