zea 0.0.6__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

zea/metrics.py

@@ -10,8 +10,9 @@ from keras import ops
 from zea import log, tensor_ops
 from zea.backend import func_on_device
 from zea.internal.registry import metrics_registry
+from zea.internal.utils import reduce_to_signature
 from zea.models.lpips import LPIPS
-from zea.utils import reduce_to_signature, translate
+from zea.tensor_ops import translate
 
 
 def get_metric(name, **kwargs):
@@ -313,11 +314,18 @@ class Metrics:
     if specified.
 
     Example:
-        .. code-block:: python
-
-            metrics = zea.metrics.Metrics(["psnr", "lpips"], image_range=[0, 255])
-            result = metrics(y_true, y_pred)
-            print(result)  # {"psnr": 30.5, "lpips": 0.15}
+        .. doctest::
+
+            >>> from zea import metrics
+            >>> import numpy as np
+
+            >>> metrics = metrics.Metrics(["psnr", "lpips"], image_range=[0, 255])
+            >>> y_true = np.random.rand(4, 128, 128, 1)
+            >>> y_pred = np.random.rand(4, 128, 128, 1)
+            >>> result = metrics(y_true, y_pred)
+            >>> result = {k: float(v) for k, v in result.items()}
+            >>> print(result)  # doctest: +ELLIPSIS
+            {'psnr': ..., 'lpips': ...}
     """
 
     def __init__(
@@ -364,7 +372,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)
+            metric_fn = tensor_ops.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

@@ -2,31 +2,37 @@
 
 ``zea`` contains a collection of models for various tasks, all located in the :mod:`zea.models` package.
 
-Currently, the following models are available (all inherited from :class:`zea.models.BaseModel`):
+See the following dropdown for a list of available models:
 
-- :class:`zea.models.echonet.EchoNetDynamic`: A model for left ventricle segmentation.
-- :class:`zea.models.carotid_segmenter.CarotidSegmenter`: A model for carotid artery segmentation.
-- :class:`zea.models.echonetlvh.EchoNetLVH`: A model for left ventricle hypertrophy segmentation.
-- :class:`zea.models.unet.UNet`: A simple U-Net implementation.
-- :class:`zea.models.lpips.LPIPS`: A model implementing the perceptual similarity metric.
-- :class:`zea.models.taesd.TinyAutoencoder`: A tiny autoencoder model for image compression.
+.. dropdown:: **Available models**
+
+    - :class:`zea.models.echonet.EchoNetDynamic`: A model for left ventricle segmentation.
+    - :class:`zea.models.carotid_segmenter.CarotidSegmenter`: A model for carotid artery segmentation.
+    - :class:`zea.models.echonetlvh.EchoNetLVH`: A model for left ventricle hypertrophy segmentation.
+    - :class:`zea.models.unet.UNet`: A simple U-Net implementation.
+    - :class:`zea.models.lpips.LPIPS`: A model implementing the perceptual similarity metric.
+    - :class:`zea.models.taesd.TinyAutoencoder`: A tiny autoencoder model for image compression.
+    - :class:`zea.models.regional_quality.MobileNetv2RegionalQuality`: A scoring model for myocardial regions in apical views.
+    - :class:`zea.models.lv_segmentation.AugmentedCamusSeg`: A nnU-Net based left ventricle and myocardium segmentation model.
 
 Presets for these models can be found in :mod:`zea.models.presets`.
 
 To use these models, you can import them directly from the :mod:`zea.models` module and load the pretrained weights using the :meth:`from_preset` method. For example:
 
-.. code-block:: python
+.. doctest::
 
-    from zea.models.unet import UNet
+    >>> from zea.models.unet import UNet
 
-    model = UNet.from_preset("unet-echonet-inpainter")
+    >>> model = UNet.from_preset("unet-echonet-inpainter")
 
 You can list all available presets using the :attr:`presets` attribute:
 
-.. code-block:: python
+.. doctest::
 
-    presets = list(UNet.presets.keys())
-    print(f"Available built-in zea presets for UNet: {presets}")
+    >>> from zea.models.unet import UNet
+    >>> presets = list(UNet.presets.keys())
+    >>> print(f"Available built-in zea presets for UNet: {presets}")
+    Available built-in zea presets for UNet: ['unet-echonet-inpainter']
 
 
 zea generative models
@@ -40,19 +46,21 @@ Typically, these models have some additional methods, such as:
 - :meth:`posterior_sample` for drawing samples from the posterior given measurements
 - :meth:`log_density` for computing the log-probability of data under the model
 
-The following generative models are currently available:
+See the following dropdown for a list of available *generative* models:
+
+.. dropdown:: **Available models**
 
-- :class:`zea.models.diffusion.DiffusionModel`: A deep generative diffusion model for ultrasound image generation.
-- :class:`zea.models.gmm.GaussianMixtureModel`: A Gaussian Mixture Model.
+    - :class:`zea.models.diffusion.DiffusionModel`: A deep generative diffusion model for ultrasound image generation.
+    - :class:`zea.models.gmm.GaussianMixtureModel`: A Gaussian Mixture Model.
 
 An example of how to use the :class:`zea.models.diffusion.DiffusionModel` is shown below:
 
-.. code-block:: python
+.. doctest::
 
-    from zea.models.diffusion import DiffusionModel
+    >>> from zea.models.diffusion import DiffusionModel
 
-    model = DiffusionModel.from_preset("diffusion-echonet-dynamic")
-    samples = model.sample(n_samples=4)
+    >>> model = DiffusionModel.from_preset("diffusion-echonet-dynamic")  # doctest: +SKIP
+    >>> samples = model.sample(n_samples=4)  # doctest: +SKIP
 
 
 Contributing and adding new models
@@ -84,7 +92,9 @@ from . import (
     gmm,
     layers,
     lpips,
+    lv_segmentation,
     presets,
+    regional_quality,
     taesd,
     unet,
     utils,

zea/models/base.py

@@ -8,6 +8,7 @@ import importlib
 import keras
 from keras.src.saving.serialization_lib import record_object_after_deserialization
 
+from zea import log
 from zea.internal.core import classproperty
 from zea.models.preset_utils import builtin_presets, get_preset_loader, get_preset_saver
 
@@ -77,7 +78,7 @@ class BaseModel(keras.models.Model):
                 initialized.
             **kwargs: Additional keyword arguments.
 
-        Examples:
+        Example:
             .. code-block:: python
 
                 # Load a Gemma backbone with pre-trained weights.
@@ -96,14 +97,29 @@ class BaseModel(keras.models.Model):
 
         """
         loader = get_preset_loader(preset)
-        model_cls = loader.check_model_class()
-        if not issubclass(model_cls, cls):
-            raise ValueError(
-                f"Saved preset has type `{model_cls.__name__}` which is not "
-                f"a subclass of calling class `{cls.__name__}`. Call "
-                f"`from_preset` directly on `{model_cls.__name__}` instead."
-            )
-        return loader.load_model(model_cls, load_weights, **kwargs)
+        loader_cls = loader.check_model_class()
+        if cls != loader_cls:
+            full_cls_name = f"{cls.__module__}.{cls.__name__}"
+            full_loader_cls_name = f"{loader_cls.__module__}.{loader_cls.__name__}"
+            if issubclass(cls, loader_cls):
+                log.warning(
+                    f"The preset '{preset}' is for model class '{full_loader_cls_name}', but you "
+                    f"are calling from a subclass '{full_cls_name}', so the returned object will "
+                    f"be of type '{full_cls_name}'."
+                )
+            elif issubclass(loader_cls, cls):
+                log.warning(
+                    f"The preset '{preset}' is for model class '{full_loader_cls_name}', "
+                    f"which is a subclass of the calling class '{full_cls_name}', "
+                    f"so the returned object will be of type '{full_cls_name}'."
+                )
+            else:
+                raise ValueError(
+                    f"The preset '{preset}' is for model class '{full_loader_cls_name}', "
+                    f"which is not compatible with the calling class '{full_cls_name}'. "
+                    f"Please call '{full_loader_cls_name}.from_preset()' instead."
+                )
+        return loader.load_model(cls, load_weights, **kwargs)
 
     def save_to_preset(self, preset_dir):
         """Save backbone to a preset directory.
@@ -115,7 +131,7 @@ class BaseModel(keras.models.Model):
         saver.save_model(self)
 
 
-def deserialize_zea_object(config):
+def deserialize_zea_object(config, cls=None):
     """Retrieve the object by deserializing the config dict.
 
     Need to borrow this function from keras and customize a bit to allow
@@ -132,10 +148,10 @@ def deserialize_zea_object(config):
     class_name = config["class_name"]
     inner_config = config["config"] or {}
 
-    module = config.get("module", None)
-    registered_name = config.get("registered_name", class_name)
-
-    cls = _retrieve_class(module, registered_name, config)
+    if cls is None:
+        module = config.get("module", None)
+        registered_name = config.get("registered_name", class_name)
+        cls = _retrieve_class(module, registered_name, config)
 
     if not hasattr(cls, "from_config"):
         raise TypeError(