zea 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

zea/io_lib.py

@@ -19,7 +19,7 @@ from PIL import Image, ImageSequence
 from zea import log
 from zea.data.file import File
 
-_SUPPORTED_VID_TYPES = [".avi", ".mp4", ".gif"]
+_SUPPORTED_VID_TYPES = [".mp4", ".gif"]
 _SUPPORTED_IMG_TYPES = [".jpg", ".png", ".JPEG", ".PNG", ".jpeg"]
 _SUPPORTED_ZEA_TYPES = [".hdf5", ".h5"]
 
@@ -27,7 +27,7 @@ _SUPPORTED_ZEA_TYPES = [".hdf5", ".h5"]
 def load_video(filename, mode="L"):
     """Load a video file and return a numpy array of frames.
 
-    Supported file types: avi, mp4, gif.
+    Supported file types: mp4, gif.
 
     Args:
         filename (str): The path to the video file.
@@ -57,12 +57,24 @@ def load_video(filename, mode="L"):
         with Image.open(filename) as im:
             for frame in ImageSequence.Iterator(im):
                 frames.append(_convert_image_mode(frame, mode=mode))
-    else:  # .mp4, .avi
-        reader = imageio.get_reader(filename)
-        for frame in reader:
-            img = Image.fromarray(frame)
-            frames.append(_convert_image_mode(img, mode=mode))
-        reader.close()
+    elif ext == ".mp4":
+        # Use imageio with FFMPEG format for MP4 files
+        try:
+            reader = imageio.get_reader(filename, format="FFMPEG")
+        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:
+                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)
 
@@ -98,17 +110,151 @@ def load_image(filename, mode="L"):
         return _convert_image_mode(img, mode=mode)
 
 
-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
+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 files, this includes `shared_color_palette` (bool).
+
+    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)
+    elif ext == ".gif":
+        return save_to_gif(images, filename, fps=fps, **kwargs)
+    else:
+        raise ValueError(f"Unsupported file extension: {ext}")
+
+
+def save_to_gif(images, filename, fps=20, shared_color_palette=False):
+    """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 False, which is default behavior
+            of PIL.Image.save. Note: True can cause slow saving for longer
+            sequences, and also lead to larger file sizes in some cases.
+
+    """
+    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]
+
+    if shared_color_palette:
+        # Apply the same palette to all frames without dithering for consistent color mapping
+        # Convert all images to RGB and combine their colors for palette generation
+        all_colors = np.vstack([np.array(img.convert("RGB")).reshape(-1, 3) for img in pillow_imgs])
+        combined_image = Image.fromarray(all_colors.reshape(-1, 1, 3))
+
+        # Generate palette from all frames
+        global_palette = combined_image.quantize(
+            colors=256,
+            method=Image.MEDIANCUT,
+            kmeans=1,
+        )
+
+        # Apply the same palette to all frames without dithering
+        pillow_imgs = [
+            img.convert("RGB").quantize(
+                palette=global_palette,
+                dither=Image.NONE,
+            )
+            for img in pillow_imgs
+        ]
+
+    pillow_img, *pillow_imgs = pillow_imgs
+
+    pillow_img.save(
+        fp=filename,
+        format="GIF",
+        append_images=pillow_imgs,
+        save_all=True,
+        loop=0,
+        duration=duration,
+        interlace=False,
+        optimize=False,
+    )
+    log.success(f"Successfully saved GIF to -> {log.yellow(filename)}")
+
+
+def save_to_mp4(images, filename, fps=20):
+    """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.
+
+    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", 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:
+        for image in images:
+            writer.append_data(image)
+    finally:
+        writer.close()
+
+    return log.success(f"Successfully saved MP4 to -> {filename}")
 
 
 def search_file_tree(
@@ -341,3 +487,85 @@ def retry_on_io_error(max_retries=3, initial_delay=0.5, retry_action=None):
         return wrapper
 
     return decorator
+
+
+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
+
+
+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]}"
+        )
+
+
+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)
+    _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

zea/keras_ops.py

@@ -3,14 +3,14 @@ and :mod:`keras.ops.image` functions.
 
 They can be used in zea pipelines like any other :class:`zea.Operation`, for example:
 
-.. code-block:: python
+.. doctest::
 
-    from zea.keras_ops import Squeeze
+    >>> from zea.keras_ops import Squeeze
 
-    op = Squeeze(axis=1)
+    >>> op = Squeeze(axis=1)
 
 This file is generated automatically. Do not edit manually.
-Generated with Keras 3.11.3
+Generated with Keras 3.12.0
 """
 
 import keras
@@ -388,6 +388,16 @@ class Cholesky(Lambda):
         except AttributeError as e:
             raise MissingKerasOps("Cholesky", "keras.ops.cholesky") from e
 
+@ops_registry("keras.ops.cholesky_inverse")
+class CholeskyInverse(Lambda):
+    """Operation wrapping keras.ops.cholesky_inverse."""
+
+    def __init__(self, **kwargs):
+        try:
+            super().__init__(func=keras.ops.cholesky_inverse, **kwargs)
+        except AttributeError as e:
+            raise MissingKerasOps("CholeskyInverse", "keras.ops.cholesky_inverse") from e
+
 @ops_registry("keras.ops.clip")
 class Clip(Lambda):
     """Operation wrapping keras.ops.clip."""
@@ -918,6 +928,36 @@ class Isnan(Lambda):
         except AttributeError as e:
             raise MissingKerasOps("Isnan", "keras.ops.isnan") from e
 
+@ops_registry("keras.ops.isneginf")
+class Isneginf(Lambda):
+    """Operation wrapping keras.ops.isneginf."""
+
+    def __init__(self, **kwargs):
+        try:
+            super().__init__(func=keras.ops.isneginf, **kwargs)
+        except AttributeError as e:
+            raise MissingKerasOps("Isneginf", "keras.ops.isneginf") from e
+
+@ops_registry("keras.ops.isposinf")
+class Isposinf(Lambda):
+    """Operation wrapping keras.ops.isposinf."""
+
+    def __init__(self, **kwargs):
+        try:
+            super().__init__(func=keras.ops.isposinf, **kwargs)
+        except AttributeError as e:
+            raise MissingKerasOps("Isposinf", "keras.ops.isposinf") from e
+
+@ops_registry("keras.ops.isreal")
+class Isreal(Lambda):
+    """Operation wrapping keras.ops.isreal."""
+
+    def __init__(self, **kwargs):
+        try:
+            super().__init__(func=keras.ops.isreal, **kwargs)
+        except AttributeError as e:
+            raise MissingKerasOps("Isreal", "keras.ops.isreal") from e
+
 @ops_registry("keras.ops.istft")
 class Istft(Lambda):
     """Operation wrapping keras.ops.istft."""
@@ -1828,6 +1868,16 @@ class Trunc(Lambda):
         except AttributeError as e:
             raise MissingKerasOps("Trunc", "keras.ops.trunc") from e
 
+@ops_registry("keras.ops.unfold")
+class Unfold(Lambda):
+    """Operation wrapping keras.ops.unfold."""
+
+    def __init__(self, **kwargs):
+        try:
+            super().__init__(func=keras.ops.unfold, **kwargs)
+        except AttributeError as e:
+            raise MissingKerasOps("Unfold", "keras.ops.unfold") from e
+
 @ops_registry("keras.ops.unstack")
 class Unstack(Lambda):
     """Operation wrapping keras.ops.unstack."""
@@ -1848,6 +1898,16 @@ class Var(Lambda):
         except AttributeError as e:
             raise MissingKerasOps("Var", "keras.ops.var") from e
 
+@ops_registry("keras.ops.view")
+class View(Lambda):
+    """Operation wrapping keras.ops.view."""
+
+    def __init__(self, **kwargs):
+        try:
+            super().__init__(func=keras.ops.view, **kwargs)
+        except AttributeError as e:
+            raise MissingKerasOps("View", "keras.ops.view") from e
+
 @ops_registry("keras.ops.view_as_complex")
 class ViewAsComplex(Lambda):
     """Operation wrapping keras.ops.view_as_complex."""
@@ -1987,3 +2047,13 @@ class RgbToHsv(Lambda):
             super().__init__(func=keras.ops.image.rgb_to_hsv, **kwargs)
         except AttributeError as e:
             raise MissingKerasOps("RgbToHsv", "keras.ops.image.rgb_to_hsv") from e
+
+@ops_registry("keras.ops.image.scale_and_translate")
+class ScaleAndTranslate(Lambda):
+    """Operation wrapping keras.ops.image.scale_and_translate."""
+
+    def __init__(self, **kwargs):
+        try:
+            super().__init__(func=keras.ops.image.scale_and_translate, **kwargs)
+        except AttributeError as e:
+            raise MissingKerasOps("ScaleAndTranslate", "keras.ops.image.scale_and_translate") from e

zea/log.py

@@ -6,7 +6,7 @@ to the console and to a file with color support.
 Example usage
 ^^^^^^^^^^^^^^
 
-.. code-block:: python
+.. testsetup::
 
     from zea import log
 
@@ -323,18 +323,20 @@ def set_level(level):
 
     Also sets the log level for the file logger if it exists.
 
-    Example:
-        >>> from zea import log
-        >>> with log.set_level("WARNING"):
-        ...     log.info("This will not be shown")
-        ...     log.warning("This will be shown")
-
     Args:
         level (str or int): The log level to set temporarily
             (e.g., "DEBUG", "INFO", logging.WARNING).
 
     Yields:
         None
+
+    Example:
+        .. doctest::
+
+            >>> from zea import log
+            >>> with log.set_level("ERROR"):
+            ...     _ = log.info("Info messages will not be shown")
+            ...     _ = log.error("Error messages will be shown")
     """
     prev_level = logger.level
     prev_file_level = file_logger.level if file_logger else None