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

zea/io_lib.py

@@ -4,20 +4,17 @@ Use to quickly read and write files or interact with file system.
 """
 
 import functools
-import multiprocessing
 import os
 import time
 from io import BytesIO
 from pathlib import Path
+from typing import Generator
 
 import imageio
 import numpy as np
-import tqdm
-import yaml
 from PIL import Image, ImageSequence
 
 from zea import log
-from zea.data.file import File
 
 _SUPPORTED_VID_TYPES = [".mp4", ".gif"]
 _SUPPORTED_IMG_TYPES = [".jpg", ".png", ".JPEG", ".PNG", ".jpeg"]
@@ -123,7 +120,7 @@ def save_video(images, filename, fps=20, **kwargs):
         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).
+            For GIF and mp4 files, this includes `shared_color_palette` (bool).
 
     Raises:
         ValueError: If the file extension is not supported.
@@ -133,14 +130,14 @@ def save_video(images, filename, fps=20, **kwargs):
     ext = filename.suffix.lower()
 
     if ext == ".mp4":
-        return save_to_mp4(images, filename, fps=fps)
+        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}")
 
 
-def save_to_gif(images, filename, fps=20, shared_color_palette=False):
+def save_to_gif(images, filename, fps=20, shared_color_palette=True):
     """Saves a sequence of images to a GIF file.
 
     .. note::
@@ -156,9 +153,9 @@ def save_to_gif(images, filename, fps=20, shared_color_palette=False):
         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.
+            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.
 
     """
     images = preprocess_for_saving(images)
@@ -173,15 +170,8 @@ def save_to_gif(images, filename, fps=20, shared_color_palette=False):
 
     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,
+        global_palette = compute_global_palette_by_histogram(
+            pillow_imgs, bits_per_channel=5, palette_size=256
         )
 
         # Apply the same palette to all frames without dithering
@@ -208,7 +198,7 @@ def save_to_gif(images, filename, fps=20, shared_color_palette=False):
     log.success(f"Successfully saved GIF to -> {log.yellow(filename)}")
 
 
-def save_to_mp4(images, filename, fps=20):
+def save_to_mp4(images, filename, fps=20, shared_color_palette=False):
     """Saves a sequence of images to an MP4 file.
 
     .. note::
@@ -222,6 +212,10 @@ def save_to_mp4(images, filename, fps=20):
             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.
@@ -249,166 +243,120 @@ def save_to_mp4(images, filename, fps=20):
         ) from exc
 
     try:
-        for image in images:
-            writer.append_data(image)
+        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,
+                )
+                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 -> {filename}")
 
 
-def search_file_tree(
-    directory,
-    filetypes=None,
-    write=True,
-    dataset_info_filename="dataset_info.yaml",
-    hdf5_key_for_length=None,
-    redo=False,
-    parallel=False,
-    verbose=True,
-):
-    """Lists all files in directory and sub-directories.
-
-    If dataset_info.yaml is detected in the directory, that file is read and used
-    to deduce the file paths. If not, the file paths are searched for in the
-    directory and written to a dataset_info.yaml file.
+def search_file_tree(directory, filetypes=None, verbose=True, relative=False) -> Generator:
+    """Traverse a directory tree and yield file paths matching specified file types.
 
     Args:
-        directory (str): Path to base directory to start file search.
-        filetypes (str or list, optional): Filetypes to look for in directory.
-            Defaults to image types (.png etc.). Make sure to include the dot.
-        write (bool, optional): Whether to write to dataset_info.yaml file.
-            Defaults to True. If False, the file paths are not written to file
-            and simply returned.
-        dataset_info_filename (str, optional): Name of dataset info file.
-            Defaults to "dataset_info.yaml", but can be changed to any name.
-        hdf5_key_for_length (str, optional): Key to use for getting length of hdf5 files.
-            Defaults to None. If set, the number of frames in each hdf5 file is
-            calculated and stored in the dataset_info.yaml file. This is extra
-            functionality of ``search_file_tree`` and only works with hdf5 files.
-        redo (bool, optional): Whether to redo the search and overwrite the dataset_info.yaml file.
-        parallel (bool, optional): Whether to use multiprocessing for hdf5 shape reading.
-        verbose (bool, optional): Whether to print progress and info.
-
-    Returns:
-        dict: Dictionary containing file paths and total number of files.
-            Has the following structure:
-
-            .. code-block:: python
-
-                {
-                    "file_paths": list of file paths,
-                    "total_num_files": total number of files,
-                    "file_lengths": list of number of frames in each hdf5 file,
-                    "file_shapes": list of shapes of each image file,
-                    "total_num_frames": total number of frames in all hdf5 files
-                }
-
+        directory (str or Path): The root directory to start the search.
+        filetypes (list of str, optional): List of file extensions to match.
+            If None, file types supported by `zea` 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.
     """
-    directory = Path(directory)
-    if not directory.is_dir():
-        raise ValueError(
-            log.error(f"Directory {directory} does not exist. Please provide a valid directory.")
-        )
-    assert Path(dataset_info_filename).suffix == ".yaml", (
-        "Currently only YAML files are supported for dataset info file when "
-        f"using `search_file_tree`, got {dataset_info_filename}"
-    )
-
-    if (directory / dataset_info_filename).is_file() and not redo:
-        with open(directory / dataset_info_filename, "r", encoding="utf-8") as file:
-            dataset_info = yaml.load(file, Loader=yaml.FullLoader)
-
-        # Check if the file_shapes key is present in the dataset_info, otherwise redo the search
-        if "file_shapes" in dataset_info:
-            if verbose:
-                log.info(
-                    "Using pregenerated dataset info file: "
-                    f"{log.yellow(directory / dataset_info_filename)} ..."
-                )
-                log.info(f"...for reading file paths in {log.yellow(directory)}")
-            return dataset_info
-
-    if redo and verbose:
-        log.info(f"Overwriting dataset info file: {log.yellow(directory / dataset_info_filename)}")
-
-    # set default file type
-    if filetypes is None:
-        filetypes = _SUPPORTED_IMG_TYPES + _SUPPORTED_VID_TYPES + _SUPPORTED_ZEA_TYPES
-
-    file_paths = []
-
-    if isinstance(filetypes, str):
-        filetypes = [filetypes]
-
-    if hdf5_key_for_length is not None:
-        assert isinstance(hdf5_key_for_length, str), "hdf5_key_for_length must be a string"
-        assert set(filetypes).issubset({".hdf5", ".h5"}), (
-            "hdf5_key_for_length only works with when filetypes is set to "
-            f"`.hdf5` or `.h5`, got {filetypes}"
-        )
-
     # 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 Path(file).suffix in filetypes:
                 file_path = Path(dirpath) / file
-                file_path = file_path.relative_to(directory)
-                file_paths.append(str(file_path))
-
-    if hdf5_key_for_length is not None:
-        # using multiprocessing to speed up reading hdf5 files
-        # and getting the number of frames in each file
-        if verbose:
-            log.info("Getting number of frames in each hdf5 file...")
-
-        get_shape_partial = functools.partial(File.get_shape, key=hdf5_key_for_length)
-        # make sure to call search_file_tree from within a function
-        # or use if __name__ == "__main__":
-        # to avoid freezing the main process
-        absolute_file_paths = [directory / file for file in file_paths]
-        if parallel:
-            with multiprocessing.Pool() as pool:
-                file_shapes = list(
-                    tqdm.tqdm(
-                        pool.imap(
-                            get_shape_partial,
-                            absolute_file_paths,
-                        ),
-                        total=len(file_paths),
-                        desc="Getting number of frames in each hdf5 file",
-                        disable=not verbose,
-                    )
-                )
-        else:
-            file_shapes = []
-            for file_path in tqdm.tqdm(
-                absolute_file_paths,
-                desc="Getting number of frames in each hdf5 file",
-                disable=not verbose,
-            ):
-                file_shapes.append(File.get_shape(file_path, hdf5_key_for_length))
-
-    assert len(file_paths) > 0, f"No image files were found in: {directory}"
-    if verbose:
-        log.info(f"Found {len(file_paths)} image files in {log.yellow(directory)}")
-        log.info(f"Writing dataset info to {log.yellow(directory / dataset_info_filename)}")
+                if relative:
+                    file_path = file_path.relative_to(directory)
+                yield file_path
+
+
+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)
+
+    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
+
+    # 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])]
+
+    # convert bin index back to representative RGB (center of bin)
+    bins = np.array(
+        [((i // (bins_per * bins_per)), (i // bins_per) % bins_per, i % bins_per) for i in top_idx]
+    )
 
-    dataset_info = {"file_paths": file_paths, "total_num_files": len(file_paths)}
-    if len(file_shapes) > 0:
-        dataset_info["file_shapes"] = file_shapes
-        file_lengths = [shape[0] for shape in file_shapes]
-        dataset_info["file_lengths"] = file_lengths
-        dataset_info["total_num_frames"] = sum(file_lengths)
+    # expand bin centers back to 8-bit values
+    center = (
+        (bins * (1 << (8 - bits_per_channel)) + (1 << (7 - bits_per_channel))).clip(0, 255)
+    ).astype(np.uint8)
+    palette_colors = center.reshape(-1, 3)  # shape (k,3)
 
-    if write:
-        with open(directory / dataset_info_filename, "w", encoding="utf-8") as file:
-            yaml.dump(dataset_info, file)
+    # 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 dataset_info
+    return palette_img
 
 
 def matplotlib_figure_to_numpy(fig, **kwargs):

zea/metrics.py

@@ -7,12 +7,13 @@ import keras
 import numpy as np
 from keras import ops
 
-from zea import log, tensor_ops
+from zea import log
 from zea.backend import func_on_device
+from zea.func import tensor
+from zea.func.tensor import translate
 from zea.internal.registry import metrics_registry
 from zea.internal.utils import reduce_to_signature
 from zea.models.lpips import LPIPS
-from zea.tensor_ops import translate
 
 
 def get_metric(name, **kwargs):
@@ -197,7 +198,7 @@ def ssim(
 
         # Construct a 1D convolution.
         def filter_fn_1(z):
-            return tensor_ops.correlate(z, ops.flip(filt), mode="valid")
+            return tensor.correlate(z, ops.flip(filt), mode="valid")
 
         # Apply the vectorized filter along the y axis.
         def filter_fn_y(z):
@@ -300,7 +301,7 @@ def get_lpips(image_range, batch_size=None, clip=False):
 
         imgs = ops.stack([img1, img2], axis=-1)
         n_batch_dims = ops.ndim(img1) - 3
-        return tensor_ops.func_with_one_batch_dim(
+        return tensor.func_with_one_batch_dim(
             unstack_lpips, imgs, n_batch_dims, batch_size=batch_size
         )
 
@@ -372,7 +373,7 @@ class Metrics:
         # Because most metric functions do not support batching, we vmap over the batch axes.
         metric_fn = fun
         for ax in reversed(batch_axes):
-            metric_fn = tensor_ops.vmap(metric_fn, in_axes=ax, _use_torch_vmap=True)
+            metric_fn = tensor.vmap(metric_fn, in_axes=ax, _use_torch_vmap=True)
 
         out = func_on_device(metric_fn, device, y_true, y_pred)
 

zea/models/__init__.py

@@ -72,7 +72,7 @@ The following steps are recommended when adding a new model:
 
 1. Create a new module in the :mod:`zea.models` package for your model: ``zea.models.mymodel``.
 2. Add a model class that inherits from :class:`zea.models.base.Model`. For generative models, inherit from :class:`zea.models.generative.GenerativeModel` or :class:`zea.models.deepgenerative.DeepGenerativeModel` as appropriate. Make sure you implement the :meth:`call` method.
-3. Upload the pretrained model weights to `our Hugging Face <https://huggingface.co/zea>`_. Should be a ``config.json`` and a ``model.weights.h5`` file. See `Keras documentation <https://keras.io/guides/serialization_and_saving/>`_ how those can be saved from your model. Simply drag and drop the files to the Hugging Face website to upload them.
+3. Upload the pretrained model weights to `our Hugging Face <https://huggingface.co/zeahub>`_. Should be a ``config.json`` and a ``model.weights.h5`` file. See `Keras documentation <https://keras.io/guides/serialization_and_saving/>`_ how those can be saved from your model. Simply drag and drop the files to the Hugging Face website to upload them.
 
    .. tip::
       It is recommended to use the mentioned saving procedure. However, alternate saving methods are also possible, see the :class:`zea.models.echonet.EchoNet` module for an example. You do now have to implement a :meth:`custom_load_weights` method in your model class.

zea/models/diffusion.py

@@ -9,6 +9,10 @@ To try this model, simply load one of the available presets:
 
     >>> model = DiffusionModel.from_preset("diffusion-echonet-dynamic")  # doctest: +SKIP
 
+.. seealso::
+    A tutorial notebook where this model is used:
+    :doc:`../notebooks/models/diffusion_model_example`.
+
 """
 
 import abc
@@ -19,6 +23,7 @@ from keras import ops
 
 from zea.backend import _import_tf, jit
 from zea.backend.autograd import AutoGrad
+from zea.func.tensor import L2, fori_loop, split_seed
 from zea.internal.core import Object
 from zea.internal.operators import Operator
 from zea.internal.registry import diffusion_guidance_registry, model_registry, operator_registry
@@ -29,7 +34,6 @@ from zea.models.preset_utils import register_presets
 from zea.models.presets import diffusion_model_presets
 from zea.models.unet import get_time_conditional_unetwork
 from zea.models.utils import LossTrackerWrapper
-from zea.tensor_ops import L2, fori_loop, split_seed
 
 tf = _import_tf()
 
@@ -51,6 +55,9 @@ class DiffusionModel(DeepGenerativeModel):
         name="diffusion_model",
         guidance="dps",
         operator="inpainting",
+        ema_val=0.999,
+        min_t=0.0,
+        max_t=1.0,
         **kwargs,
     ):
         """Initialize a diffusion model.
@@ -58,17 +65,20 @@ class DiffusionModel(DeepGenerativeModel):
         Args:
             input_shape: Shape of the input data. Typically of the form
                 `(height, width, channels)` for images.
-            widths: List of filter widths for the UNet.
-            block_depth: Number of residual blocks in each UNet block.
-            timesteps: Number of diffusion timesteps.
-            beta_start: Initial noise schedule value.
-            beta_end: Final noise schedule value.
+            input_range: Range of the input data.
+            min_signal_rate: Minimum signal rate for the diffusion schedule.
+            max_signal_rate: Maximum signal rate for the diffusion schedule.
+            network_name: Name of the network architecture to use. Options are
+                "unet_time_conditional" or "dense_time_conditional".
+            network_kwargs: Additional keyword arguments for the network.
             name: Name of the model.
             guidance: Guidance method to use. Can be a string, or dict with
                 "name" and "params" keys. Additionally, can be a `DiffusionGuidance` object.
             operator: Operator to use. Can be a string, or dict with
                 "name" and "params" keys. Additionally, can be a `Operator` object.
-
+            ema_val: Exponential moving average value for the network weights.
+            min_t: Minimum diffusion time for sampling during training.
+            max_t: Maximum diffusion time for sampling during training.
             **kwargs: Additional arguments.
         """
         super().__init__(name=name, **kwargs)
@@ -79,10 +89,11 @@ class DiffusionModel(DeepGenerativeModel):
         self.max_signal_rate = max_signal_rate
         self.network_name = network_name
         self.network_kwargs = network_kwargs or {}
+        self.ema_val = ema_val
 
-        # reverse diffusion (i.e. sampling) goes from max_t to min_t
-        self.min_t = 0.0
-        self.max_t = 1.0
+        # reverse diffusion (i.e. sampling) goes from t = max_t to t = min_t
+        self.min_t = min_t
+        self.max_t = max_t
 
         if network_name == "unet_time_conditional":
             self.network = get_time_conditional_unetwork(
@@ -122,8 +133,11 @@ class DiffusionModel(DeepGenerativeModel):
                 "input_range": self.input_range,
                 "min_signal_rate": self.min_signal_rate,
                 "max_signal_rate": self.max_signal_rate,
+                "min_t": self.min_t,
+                "max_t": self.max_t,
                 "network_name": self.network_name,
                 "network_kwargs": self.network_kwargs,
+                "ema_val": self.ema_val,
             }
         )
         return config
@@ -316,8 +330,8 @@ class DiffusionModel(DeepGenerativeModel):
         # Sample uniform random diffusion times in [min_t, max_t]
         diffusion_times = keras.random.uniform(
             shape=[batch_size, *[1] * n_dims],
-            minval=self.min_signal_rate,
-            maxval=self.max_signal_rate,
+            minval=self.min_t,
+            maxval=self.max_t,
         )
         noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
 
@@ -337,6 +351,43 @@ class DiffusionModel(DeepGenerativeModel):
         self.noise_loss_tracker.update_state(noise_loss)
         self.image_loss_tracker.update_state(image_loss)
 
+        # track the exponential moving averages of weights.
+        # ema_network is used for inference / sampling
+        for weight, ema_weight in zip(self.network.weights, self.ema_network.weights):
+            ema_weight.assign(self.ema_val * ema_weight + (1 - self.ema_val) * weight)
+
+        return {m.name: m.result() for m in self.metrics}
+
+    def test_step(self, data):
+        """
+        Custom test step so we can call model.fit() on the diffusion model.
+        """
+        batch_size, *input_shape = ops.shape(data)
+        n_dims = len(input_shape)
+
+        noises = keras.random.normal(shape=ops.shape(data))
+
+        # sample uniform random diffusion times
+        diffusion_times = keras.random.uniform(
+            shape=[batch_size, *[1] * n_dims],
+            minval=self.min_t,
+            maxval=self.max_t,
+        )
+        noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)
+        # mix the images with noises accordingly
+        noisy_images = signal_rates * data + noise_rates * noises
+
+        # use the network to separate noisy images to their components
+        pred_noises, pred_images = self.denoise(
+            noisy_images, noise_rates, signal_rates, training=False
+        )
+
+        noise_loss = self.loss(noises, pred_noises)
+        image_loss = self.loss(data, pred_images)
+
+        self.noise_loss_tracker.update_state(noise_loss)
+        self.image_loss_tracker.update_state(image_loss)
+
         return {m.name: m.result() for m in self.metrics}
 
     def diffusion_schedule(self, diffusion_times):

zea/models/echonetlvh.py

@@ -25,12 +25,12 @@ To try this model, simply load one of the available presets:
 import numpy as np
 from keras import ops
 
+from zea.func.tensor import translate
 from zea.internal.registry import model_registry
 from zea.models.base import BaseModel
 from zea.models.deeplabv3 import DeeplabV3Plus
 from zea.models.preset_utils import register_presets
 from zea.models.presets import echonet_lvh_presets
-from zea.tensor_ops import translate
 
 
 @model_registry(name="echonetlvh")

zea/models/gmm.py

@@ -4,8 +4,8 @@ import keras
 import numpy as np
 from keras import ops
 
+from zea.func.tensor import linear_sum_assignment
 from zea.models.generative import GenerativeModel
-from zea.tensor_ops import linear_sum_assignment
 
 
 class GaussianMixtureModel(GenerativeModel):

zea/models/lv_segmentation.py

@@ -44,6 +44,8 @@ from zea.models.base import BaseModel
 from zea.models.preset_utils import get_preset_loader, register_presets
 from zea.models.presets import augmented_camus_seg_presets
 
+INFERENCE_SIZE = 256
+
 
 @model_registry(name="augmented_camus_seg")
 class AugmentedCamusSeg(BaseModel):