"""Input / output functions for reading and writing files.
Use to quickly read and write files or interact with file system.
"""
import os
from io import BytesIO
from pathlib import Path
from typing import Generator
import imageio
import numpy as np
from PIL import Image, ImageSequence
from zea import log
_SUPPORTED_VID_TYPES = [".mp4", ".gif"]
_SUPPORTED_IMG_TYPES = [".jpg", ".png", ".JPEG", ".PNG", ".jpeg"]
_SUPPORTED_ZEA_TYPES = [".hdf5", ".h5"]
[docs]
def load_video(filename, mode="L"):
"""Load a video file and return a numpy array of frames.
Supported file types: mp4, gif.
Args:
filename (str): The path to the video file.
mode (str, optional): Color mode: "L" (grayscale) or "RGB".
Defaults to "L".
Returns:
numpy.ndarray: Array of frames (num_frames, H, W) or (num_frames, H, W, C)
Raises:
ValueError: If the file extension or mode is not supported.
"""
filename = Path(filename)
if not filename.exists():
raise FileNotFoundError(f"File {filename} does not exist")
ext = filename.suffix.lower()
if ext not in _SUPPORTED_VID_TYPES:
raise ValueError(f"Unsupported file extension: {ext}")
if mode not in {"L", "RGB"}:
raise ValueError(f"Unsupported mode: {mode}")
frames = []
if ext == ".gif":
with Image.open(filename) as im:
for frame in ImageSequence.Iterator(im):
frames.append(_convert_image_mode(frame, mode=mode))
elif ext == ".mp4":
# Use imageio with FFMPEG format for MP4 files
try:
reader = imageio.get_reader(
filename,
format="FFMPEG", # ty: ignore[invalid-argument-type]
)
except (ImportError, ValueError) as exc:
raise ImportError(
"FFMPEG plugin is required to load MP4 files. "
"Please install it with 'pip install imageio-ffmpeg'."
) from exc
try:
for frame in reader: # ty: ignore[not-iterable]
img = Image.fromarray(frame)
frames.append(_convert_image_mode(img, mode=mode))
finally:
reader.close()
else:
raise ValueError(f"Unsupported file extension: {ext}")
return np.stack(frames, axis=0)
[docs]
def load_image(filename, mode="L"):
"""Load an image file and return a numpy array.
Supported file types: jpg, png.
Args:
filename (str): The path to the image file.
mode (str, optional): Color mode: "L" (grayscale) or "RGB".
Defaults to "L".
Returns:
numpy.ndarray: A numpy array of the image.
Raises:
ValueError: If the file extension or mode is not supported.
"""
filename = Path(filename)
if not filename.exists():
raise FileNotFoundError(f"File {filename} does not exist")
extension = filename.suffix.lower()
allowed_exts = {ext.lower() for ext in _SUPPORTED_IMG_TYPES}
if extension not in allowed_exts:
raise ValueError(f"File extension {extension} not supported")
if mode not in {"L", "RGB"}:
raise ValueError(f"Unsupported mode: {mode}")
with Image.open(filename) as img:
return _convert_image_mode(img, mode=mode)
[docs]
def save_image(image, filename):
"""Saves a single image to an image file.
Supported file types: jpg, png.
Args:
image (ndarray): A single image of shape (height, width),
(height, width, channels) with 1 (grayscale), 3 (RGB) or 4 (RGBA)
channels. Image should be uint8. JPEG has no alpha channel, so an
RGBA image is written as RGB; PNG keeps the alpha channel.
filename (str or Path): Filename to which the image should be written.
Returns:
Path: The path the image was written to.
Raises:
ValueError: If the file extension is not supported.
"""
filename = Path(filename)
ext = filename.suffix.lower()
if ext not in {suffix.lower() for suffix in _SUPPORTED_IMG_TYPES}:
raise ValueError(f"Unsupported file extension: {ext}")
image = np.asarray(image)
# ``_assert_uint8_images`` validates a sequence, so check the frame as a sequence of one.
_assert_uint8_images(image[np.newaxis])
# Drop the channel axis of a grayscale image so PIL writes a single-channel image
# rather than refusing the (height, width, 1) shape.
if image.ndim == 3 and image.shape[-1] == 1:
image = np.squeeze(image, axis=-1)
# JPEG has no alpha channel, and PIL raises rather than dropping it for us.
# ``preprocess_for_saving`` drops it the same way for videos.
elif ext in (".jpg", ".jpeg") and image.ndim == 3 and image.shape[-1] == 4:
image = image[..., :3]
Image.fromarray(image).save(filename)
return filename
[docs]
def save_video(images, filename, fps=20, **kwargs):
"""Saves a sequence of images to a video file.
Supported file types: mp4, gif.
Args:
images (list or np.ndarray): List or array of images. Must have shape
(n_frames, height, width, channels) or (n_frames, height, width).
If channel axis is not present, or is 1, grayscale image is assumed,
which is then converted to RGB. Images should be uint8.
filename (str or Path): Filename to which data should be written.
fps (int): Frames per second of rendered format.
**kwargs: Additional keyword arguments passed to the specific save function.
For GIF and mp4 files, this includes `shared_color_palette` (bool).
For GIF files, `transparent=True` keeps the alpha channel of RGBA images.
Raises:
ValueError: If the file extension is not supported.
"""
filename = Path(filename)
ext = filename.suffix.lower()
if ext == ".mp4":
return save_to_mp4(images, filename, fps=fps, **kwargs)
elif ext == ".gif":
return save_to_gif(images, filename, fps=fps, **kwargs)
else:
raise ValueError(f"Unsupported file extension: {ext}")
[docs]
def save_to_gif(images, filename, fps=20, shared_color_palette=True, transparent=False, loop=0):
"""Saves a sequence of images to a GIF file.
.. note::
It's recommended to use :func:`save_video` for a more general interface
that supports multiple formats.
Args:
images (list or np.ndarray): List or array of images. Must have shape
(n_frames, height, width, channels) or (n_frames, height, width).
If channel axis is not present, or is 1, grayscale image is assumed,
which is then converted to RGB. Images should be uint8.
filename (str or Path): Filename to which data should be written.
fps (int): Frames per second of rendered format.
shared_color_palette (bool, optional): If True, creates a global
color palette across all frames, ensuring consistent colors
throughout the GIF. Defaults to True, which is default behavior
of PIL.Image.save. Note: True increases speed and shrinks file
size for longer sequences.
loop (int, optional): Number of times to repeat after the first play.
Defaults to 0, which loops forever. Pass None to write no loop
extension at all, so the GIF plays once and stops on the last frame.
transparent (bool, optional): If True, keeps the alpha channel of RGBA
images as GIF transparency. GIF only supports binary transparency,
so pixels with alpha < 128 become fully transparent and all others
fully opaque. One palette entry is reserved for the transparent
color. Requires RGBA input. Defaults to False.
"""
opaque_masks = None
if transparent:
images = np.asarray(images)
if images.ndim != 4 or images.shape[-1] != 4:
raise ValueError(
"transparent=True requires RGBA images of shape "
f"(n_frames, height, width, 4), got {images.shape}."
)
opaque_masks = images[..., 3] >= 128
images = preprocess_for_saving(images)
if fps > 50:
log.warning(f"Cannot set fps ({fps}) > 50. Setting it automatically to 50.")
fps = 50
duration = int(round(1000 / fps)) # milliseconds per frame
pillow_imgs = [Image.fromarray(img) for img in images]
# Reserve the last palette index for transparent pixels
palette_size = 255 if transparent else 256
if shared_color_palette:
# Apply the same palette to all frames without dithering for consistent color mapping
global_palette = compute_global_palette_by_histogram(
pillow_imgs, bits_per_channel=5, palette_size=palette_size
)
# Apply the same palette to all frames without dithering
pillow_imgs = [
img.convert("RGB").quantize(
palette=global_palette,
dither=Image.NONE, # ty: ignore[unresolved-attribute]
)
for img in pillow_imgs
]
transparency_kwargs = {}
if opaque_masks is not None:
transparent_index = palette_size
pillow_imgs = [
_apply_gif_transparency(img, mask, palette_size, transparent_index)
for img, mask in zip(pillow_imgs, opaque_masks)
]
# disposal=2 clears each frame to the background before drawing the next,
# otherwise transparent pixels would show the previous frame through.
transparency_kwargs = {"transparency": transparent_index, "disposal": 2}
pillow_img, *pillow_imgs = pillow_imgs
# Pillow only omits the Netscape loop extension when `loop` is absent entirely
loop_kwarg = {} if loop is None else {"loop": loop}
pillow_img.save(
fp=filename,
format="GIF",
append_images=pillow_imgs,
save_all=True,
duration=duration,
interlace=False,
optimize=False,
**loop_kwarg,
**transparency_kwargs,
)
log.success(f"Successfully saved GIF to -> {log.yellow(filename)}")
def _apply_gif_transparency(img, opaque_mask, palette_size, transparent_index):
"""Quantizes ``img`` to at most ``palette_size`` colors (if not already
palettized) and sets pixels outside ``opaque_mask`` to ``transparent_index``."""
if img.mode != "P":
img = img.convert("RGB").quantize(
colors=palette_size,
dither=Image.NONE, # ty: ignore[unresolved-attribute]
)
indices = np.array(img)
full_palette = np.zeros(768, dtype=np.uint8)
current = np.array(img.getpalette(), dtype=np.uint8)[:768]
full_palette[: current.size] = current
full_palette = full_palette.reshape(256, 3).astype(np.int32)
# Quantizing against a 256-entry palette can still pick the reserved index for
# opaque pixels; remap those to the nearest color among the real entries.
clash = (indices == transparent_index) & opaque_mask
if clash.any():
distances = np.sum((full_palette[:palette_size] - full_palette[transparent_index]) ** 2, -1)
indices[clash] = np.argmin(distances)
indices[~opaque_mask] = transparent_index
palette = np.zeros(768, dtype=np.uint8)
palette[: palette_size * 3] = full_palette[:palette_size].ravel()
out = Image.fromarray(indices, mode="P")
out.putpalette(palette.tolist())
return out
[docs]
def save_to_mp4(images, filename, fps=20, shared_color_palette=False):
"""Saves a sequence of images to an MP4 file.
.. note::
It's recommended to use :func:`save_video` for a more general interface
that supports multiple formats.
Args:
images (list or np.ndarray): List or array of images. Must have shape
(n_frames, height, width, channels) or (n_frames, height, width).
If channel axis is not present, or is 1, grayscale image is assumed,
which is then converted to RGB. Images should be uint8.
filename (str or Path): Filename to which data should be written.
fps (int): Frames per second of rendered format.
shared_color_palette (bool, optional): If True, creates a global
color palette across all frames, ensuring consistent colors
throughout the MP4. Note: True can cause slow saving for longer
sequences.
Raises:
ImportError: If imageio-ffmpeg is not installed.
Returns:
str: Success message.
"""
images = preprocess_for_saving(images)
filename = str(filename)
parent_dir = Path(filename).parent
parent_dir.mkdir(parents=True, exist_ok=True)
# Use imageio with FFMPEG format for MP4 files
try:
writer = imageio.get_writer(
filename,
fps=fps,
format="FFMPEG", # ty: ignore[invalid-argument-type]
codec="libx264",
pixelformat="yuv420p",
)
except (ImportError, ValueError) as exc:
raise ImportError(
"FFMPEG plugin is required to save MP4 files. "
"Please install it with 'pip install imageio-ffmpeg'."
) from exc
try:
if shared_color_palette:
pillow_imgs = [Image.fromarray(img) for img in images]
global_palette = compute_global_palette_by_histogram(
pillow_imgs, bits_per_channel=5, palette_size=256
)
for img in pillow_imgs:
paletted_img = img.convert("RGB").quantize(
palette=global_palette,
dither=Image.NONE, # ty: ignore[unresolved-attribute]
)
writer.append_data(np.array(paletted_img.convert("RGB")))
else:
# Write from numpy arrays directly
for image in images:
writer.append_data(image)
finally:
writer.close()
return log.success(f"Successfully saved MP4 to -> {log.yellow(filename)}")
[docs]
def search_file_tree(
directory, filetypes: list[str] | None = None, verbose=True, relative=False
) -> Generator:
"""Traverse a directory tree and yield file paths matching specified file types.
Args:
directory (str or Path): The root directory to start the search.
filetypes (list of str, optional): List of file extensions to match.
If None, all files are matched. Defaults to None.
verbose (bool, optional): If True, logs the search process. Defaults to True.
relative (bool, optional): If True, yields file paths relative to the
root directory. Defaults to False.
Yields:
Path: Paths of files matching the specified file types.
"""
# Traverse file tree to index all files from filetypes
if verbose:
log.info(f"Searching {log.yellow(directory)} for {filetypes} files...")
for dirpath, _, filenames in os.walk(directory):
for file in filenames:
# Append to file_paths if it is a filetype file
if filetypes is None or Path(file).suffix in filetypes:
file_path = Path(dirpath) / file
if relative:
file_path = file_path.relative_to(directory)
yield file_path
[docs]
def compute_global_palette_by_histogram(pillow_imgs, bits_per_channel=5, palette_size=256):
"""Computes a global color palette for a sequence of images using histogram analysis.
Args:
pillow_imgs (list): List of pillow images. All images should be in RGB mode.
bits_per_channel (int, optional): Number of bits to use per color channel for histogram
binning. Can take values between 1 and 7. Defaults to 5.
palette_size (int, optional): Number of colors in the resulting palette. Defaults to 256.
Returns:
PIL.Image: A PIL 'P' mode image containing the computed color palette.
Raises:
ValueError: If bits_per_channel or palette_size is outside of range.
"""
if not 1 <= bits_per_channel <= 7:
raise ValueError(f"bits_per_channel must be between 1 and 7, got {bits_per_channel}")
if not 1 <= palette_size <= 256:
raise ValueError(f"palette_size must be between 1 and 256, got {palette_size}")
# compute number of bins per channel by bitshift
bins_per = 1 << bits_per_channel
# compute total number of histogram bins for RGB
total_bins = bins_per**3
# counts per bin in the final histogram
counts = np.zeros(total_bins, dtype=np.int64)
# sum of the colors per bin
sums = np.zeros((total_bins, 3), dtype=np.int64)
shift = 8 - bits_per_channel
# Iterate images, accumulate bin counts
for img in pillow_imgs:
arr = np.array(img.convert("RGB"), dtype=np.uint8).reshape(-1, 3)
# reduce bits, compute bin index
r = (arr[:, 0] >> shift).astype(np.int32)
g = (arr[:, 1] >> shift).astype(np.int32)
b = (arr[:, 2] >> shift).astype(np.int32)
idx = (r * bins_per + g) * bins_per + b
# accumulate counts
bincount = np.bincount(idx, minlength=total_bins)
counts += bincount
for channel in range(3):
sums[:, channel] += np.bincount(
idx, weights=arr[:, channel], minlength=total_bins
).astype(np.int64)
# pick top bins
top_idx = np.argpartition(-counts, palette_size - 1)[:palette_size]
# sort top bins by frequency
top_idx = top_idx[np.argsort(-counts[top_idx])]
# represent each bin by the mean color in it, so exact colors are preserved
occupied = counts[top_idx] > 0
palette_colors = np.zeros((len(top_idx), 3), dtype=np.uint8) # shape (k,3)
palette_colors[occupied] = np.round(
sums[top_idx][occupied] / counts[top_idx][occupied, None]
).astype(np.uint8)
# build a PIL 'P' palette image from these colors
pal = np.zeros(768, dtype=np.uint8) # 256*3 entries
pal[: palette_colors.size] = palette_colors.flatten()
palette_img = Image.new("P", (1, 1))
palette_img.putpalette(pal.tolist())
return palette_img
def _convert_image_mode(img, mode="L"):
"""Convert a PIL Image to the specified mode and return as numpy array."""
if mode not in {"L", "RGB"}:
raise ValueError(f"Unsupported mode: {mode}, must be one of: L, RGB")
if mode == "L":
img = img.convert("L")
arr = np.array(img)
elif mode == "RGB":
img = img.convert("RGB")
arr = np.array(img)
return arr
[docs]
def grayscale_to_rgb(image):
"""Converts a grayscale image to an RGB image.
Args:
image (ndarray): Grayscale image. Must have shape (height, width).
Returns:
ndarray: RGB image.
"""
assert image.ndim == 2, "Input image must be grayscale."
# Stack the grayscale image into 3 channels (RGB)
return np.stack([image] * 3, axis=-1)
def _assert_uint8_images(images: np.ndarray):
"""
Asserts that the input images have the correct properties.
Args:
images (np.ndarray): The input images.
Raises:
AssertionError: If the dtype of images is not uint8.
AssertionError: If the shape of images is not (n_frames, height, width, channels)
or (n_frames, height, width) for grayscale images.
AssertionError: If images have anything other than 1 (grayscale),
3 (rgb) or 4 (rgba) channels.
"""
assert images.dtype == np.uint8, f"dtype of images should be uint8, got {images.dtype}"
assert images.ndim in (3, 4), (
"images must have shape (n_frames, height, width, channels),"
f" or (n_frames, height, width) for grayscale images. Got {images.shape}"
)
if images.ndim == 4:
assert images.shape[-1] in (1, 3, 4), (
"Grayscale images must have 1 channel, "
"RGB images must have 3 channels, and RGBA images must have 4 channels. "
f"Got shape: {images.shape}, channels: {images.shape[-1]}"
)
[docs]
def preprocess_for_saving(images):
"""Preprocesses images for saving to GIF or MP4.
Args:
images (ndarray, list[ndarray]): Images. Must have shape (n_frames, height, width, channels)
or (n_frames, height, width).
"""
images = np.array(images)
if len(images) == 0:
raise ValueError("Received empty list of images. Cannot save an empty video.")
_assert_uint8_images(images)
# Remove channel axis if it is 1 (grayscale image)
if images.ndim == 4 and images.shape[-1] == 1:
images = np.squeeze(images, axis=-1)
# convert grayscale images to RGB
if images.ndim == 3:
images = [grayscale_to_rgb(image) for image in images]
images = np.array(images)
# drop alpha channel if present (RGBA -> RGB)
if images.ndim == 4 and images.shape[-1] == 4:
images = images[..., :3]
return images