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

zea/models/carotid_segmenter.py

@@ -1,9 +1,24 @@
 """Carotid segmentation model.
 
-Original implementation of paper:
-    - "Unsupervised domain adaptation method for segmenting cross-sectional CCA images"
-    - https://doi.org/10.1016/j.cmpb.2022.107037
-    - Author: Luuk van Knippenberg
+To try this model, simply load one of the available presets:
+
+.. doctest::
+
+    >>> from zea.models.carotid_segmenter import CarotidSegmenter
+
+    >>> model = CarotidSegmenter.from_preset("carotid-segmenter")
+
+.. important::
+    This is a ``zea`` implementation of the model.
+    For the original paper see:
+
+    van Knippenberg, Luuk, et al.
+    "Unsupervised domain adaptation method for segmenting cross-sectional CCA images."
+    *https://doi.org/10.1016/j.cmpb.2022.107037*
+
+.. seealso::
+    A tutorial notebook where this model is used:
+    :doc:`../notebooks/models/carotid_segmentation_example`.
 """
 
 import keras

zea/models/diffusion.py

@@ -1,4 +1,15 @@
-"""Diffusion models"""
+"""
+Diffusion models for ultrasound image generation and posterior sampling.
+
+To try this model, simply load one of the available presets:
+
+.. doctest::
+
+    >>> from zea.models.diffusion import DiffusionModel
+
+    >>> model = DiffusionModel.from_preset("diffusion-echonet-dynamic")  # doctest: +SKIP
+
+"""
 
 import abc
 from typing import Literal
@@ -6,11 +17,12 @@ from typing import Literal
 import keras
 from keras import ops
 
-from zea.backend import _import_tf
+from zea.backend import _import_tf, jit
 from zea.backend.autograd import AutoGrad
 from zea.internal.core import Object
 from zea.internal.operators import Operator
 from zea.internal.registry import diffusion_guidance_registry, model_registry, operator_registry
+from zea.internal.utils import fn_requires_argument
 from zea.models.dense import get_time_conditional_dense_network
 from zea.models.generative import DeepGenerativeModel
 from zea.models.preset_utils import register_presets
@@ -18,7 +30,6 @@ 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
-from zea.utils import fn_requires_argument
 
 tf = _import_tf()
 
@@ -182,7 +193,7 @@ class DiffusionModel(DeepGenerativeModel):
             **kwargs: Additional arguments.
 
         Returns:
-            Generated samples.
+            Generated samples of shape `(n_samples, *input_shape)`.
         """
         seed, seed1 = split_seed(seed, 2)
 
@@ -233,7 +244,8 @@ class DiffusionModel(DeepGenerativeModel):
                 `(batch_size, n_samples, *input_shape)`.
 
         """
-        shape = ops.shape(measurements)
+        batch_size = ops.shape(measurements)[0]
+        shape = (batch_size, n_samples, *self.input_shape)
 
         def _tile_with_sample_dim(tensor):
             """Tile the tensor with an additional sample dimension."""
@@ -250,7 +262,7 @@ class DiffusionModel(DeepGenerativeModel):
         seed1, seed2 = split_seed(seed, 2)
 
         initial_noise = keras.random.normal(
-            shape=ops.shape(measurements),
+            shape=(batch_size * n_samples, *self.input_shape),
             seed=seed1,
         )
 
@@ -262,9 +274,9 @@ class DiffusionModel(DeepGenerativeModel):
             initial_step=initial_step,
             seed=seed2,
             **kwargs,
-        )
-        # returns: (batch_size, n_samples, *input_shape)
-        return ops.reshape(out, (shape[0], n_samples, *shape[1:]))
+        )  # ( batch_size * n_samples, *self.input_shape)
+
+        return ops.reshape(out, shape)  # (batch_size, n_samples, *input_shape)
 
     def log_likelihood(self, data, **kwargs):
         """Approximate log-likelihood of the data under the model.
@@ -776,7 +788,12 @@ register_presets(diffusion_model_presets, DiffusionModel)
 class DiffusionGuidance(abc.ABC, Object):
     """Base class for diffusion guidance methods."""
 
-    def __init__(self, diffusion_model, operator, disable_jit=False):
+    def __init__(
+        self,
+        diffusion_model: DiffusionModel,
+        operator: Operator,
+        disable_jit: bool = False,
+    ):
         """Initialize the diffusion guidance.
 
         Args:
@@ -828,8 +845,7 @@ class DPS(DiffusionGuidance):
 
         Args:
             noisy_images: Noisy images.
-            measurement: Target measurement.
-            operator: Forward operator.
+            measurements: Target measurement.
             noise_rates: Current noise rates.
             signal_rates: Current signal rates.
             omega: Weight for the measurement error.
@@ -866,3 +882,148 @@ class DPS(DiffusionGuidance):
             Tuple of (gradients, (measurement_error, (pred_noises, pred_images)))
         """
         return self.gradient_fn(noisy_images, **kwargs)
+
+
+@diffusion_guidance_registry(name="dds")
+class DDS(DiffusionGuidance):
+    """
+    Decomposed Diffusion Sampling guidance.
+
+    Reference paper: https://arxiv.org/pdf/2303.05754
+    """
+
+    def setup(self):
+        """Setup DDS guidance function."""
+        if not self.disable_jit:
+            self.call = jit(self.call)
+
+    def Acg(self, x, **op_kwargs):
+        # we transform the operator from A(x) to A.T(A(x)) to get the normal equations,
+        # so that it is suitable for conjugate gradient. (symmetric, positive definite)
+        # Normal equations: A^T y = A^T A x
+        return self.operator.transpose(self.operator.forward(x, **op_kwargs), **op_kwargs)
+
+    def conjugate_gradient_inner_loop(self, i, loop_state, eps=1e-5):
+        """
+        A single iteration of the conjugate gradient method.
+        This involves minimizing the error of x along the current search
+        vector p, and then choosing the next search vector.
+
+        Reference code from: https://github.com/svi-diffusion/
+        """
+        p, rs_old, r, x, eps, op_kwargs = loop_state
+
+        # compute alpha
+        Ap = self.Acg(p, **op_kwargs)  # transform search vector p by A
+        a = rs_old / ops.sum(p * Ap)  # minimize f along the line p
+
+        x_new = x + a * p  # set new x at the minimum of f along line p
+        r_new = r - a * Ap  # shortcut to compute next residual
+
+        # compute Gram-Schmidt coefficient beta to choose next search vector
+        # so that p_new is A-orthogonal to p_current.
+        rs_new = ops.sum(r_new * r_new)
+        p_new = r_new + (rs_new / rs_old) * p
+
+        # this is like a jittable 'break' -- if the residual
+        # is less than eps, then we just return the old
+        # loop state rather than the updated one.
+        next_loop_state = ops.cond(
+            ops.abs(ops.sqrt(rs_old)) < eps,
+            lambda: (p, rs_old, r, x, eps, op_kwargs),
+            lambda: (p_new, rs_new, r_new, x_new, eps, op_kwargs),
+        )
+
+        return next_loop_state
+
+    def call(
+        self,
+        noisy_images,
+        measurements,
+        noise_rates,
+        signal_rates,
+        n_inner,
+        eps,
+        verbose,
+        **op_kwargs,
+    ):
+        """
+        Call the DDS guidance function
+
+        Args:
+            noisy_images: Noisy images.
+            measurement: Target measurement.
+            noise_rates: Current noise rates.
+            signal_rates: Current signal rates.
+            n_inner: Number of conjugate gradient steps.
+            verbose: Whether to calculate error.
+
+        Returns:
+            Tuple of (gradients, (measurement_error, (pred_noises, pred_images)))
+        """
+        pred_noises, pred_images = self.diffusion_model.denoise(
+            noisy_images,
+            noise_rates,
+            signal_rates,
+            training=False,
+        )
+        measurements_cg = self.operator.transpose(measurements, **op_kwargs)
+        r = measurements_cg - self.Acg(pred_images, **op_kwargs)  # residual
+        p = ops.copy(r)  # initial search vector = residual
+        rs_old = ops.sum(r * r)  # residual dot product
+        _, _, _, pred_images_updated_cg, _, _ = fori_loop(
+            0,
+            n_inner,
+            self.conjugate_gradient_inner_loop,
+            (p, rs_old, r, pred_images, eps, op_kwargs),
+        )
+
+        # Not strictly necessary, just for debugging
+        error = ops.cond(
+            verbose,
+            lambda: L2(measurements - self.operator.forward(pred_images_updated_cg, **op_kwargs)),
+            lambda: 0.0,
+        )
+
+        pred_images = pred_images_updated_cg
+        # we have already performed the guidance steps in self.conjugate_gradient_method, so
+        # we can set these gradients to zero.
+        gradients = ops.zeros_like(pred_images)
+        return gradients, (error, (pred_noises, pred_images))
+
+    def __call__(
+        self,
+        noisy_images,
+        measurements,
+        noise_rates,
+        signal_rates,
+        n_inner=5,
+        eps=1e-5,
+        verbose=False,
+        **op_kwargs,
+    ):
+        """
+        Call the DDS guidance function
+
+        Args:
+            noisy_images: Noisy images.
+            measurement: Target measurement.
+            noise_rates: Current noise rates.
+            signal_rates: Current signal rates.
+            n_inner: Number of conjugate gradient steps.
+            verbose: Whether to calculate error.
+            **kwargs: Additional arguments for the operator.
+
+        Returns:
+            Tuple of (gradients, (measurement_error, (pred_noises, pred_images)))
+        """
+        return self.call(
+            noisy_images,
+            measurements,
+            noise_rates,
+            signal_rates,
+            n_inner,
+            eps,
+            verbose,
+            **op_kwargs,
+        )

zea/models/echonet.py

@@ -1,6 +1,25 @@
-"""Echonet-Dynamic segmentation model for cardiac ultrasound segmentation.
-Link below does not work it seems, this is slightly different but does have some info:
-https://github.com/bryanhe/dynamic
+"""
+Echonet-Dynamic segmentation model for cardiac ultrasound segmentation.
+
+To try this model, simply load one of the available presets:
+
+.. doctest::
+
+    >>> from zea.models.echonet import EchoNetDynamic
+
+    >>> model = EchoNetDynamic.from_preset("echonet-dynamic")  # doctest: +SKIP
+
+.. important::
+    This is a ``zea`` implementation of the model.
+    For the original paper and code, see `here <https://echonet.github.io/dynamic/>`_.
+
+    Ouyang, David, et al. "Video-based AI for beat-to-beat assessment of cardiac function."
+    *Nature 580.7802 (2020): 252-256*
+
+.. seealso::
+    A tutorial notebook where this model is used:
+    :doc:`../notebooks/models/left_ventricle_segmentation_example`.
+
 """
 
 from pathlib import Path
@@ -33,11 +52,6 @@ tf = _import_tf()
 class EchoNetDynamic(BaseModel):
     """EchoNet-Dynamic segmentation model for cardiac ultrasound segmentation.
 
-    Original paper and code: https://echonet.github.io/dynamic/
-
-    This class extracts useful parts of the original code and wraps it in a
-    easy to use class.
-
     Preprocessing should normalize the input images with mean and standard deviation.
 
     """

zea/models/echonetlvh.py

@@ -1,4 +1,26 @@
-"""EchoNetLVH model for segmentation of PLAX view cardiac ultrasound. For more details see https://echonet.github.io/lvh/index.html."""
+"""EchoNetLVH model for segmentation of PLAX view cardiac ultrasound.
+
+To try this model, simply load one of the available presets:
+
+.. doctest::
+
+    >>> from zea.models.echonetlvh import EchoNetLVH
+
+    >>> model = EchoNetLVH.from_preset("echonetlvh")
+
+.. important::
+    This is a ``zea`` implementation of the model.
+    For the original paper and code, see `here <https://echonet.github.io/lvh/>`_.
+
+    Duffy, Grant, et al.
+    "High-throughput precision phenotyping of left ventricular hypertrophy with cardiovascular deep learning."
+    *JAMA cardiology 7.4 (2022): 386-395*
+
+.. seealso::
+    A tutorial notebook where this model is used:
+    :doc:`../notebooks/agent/task_based_perception_action_loop`.
+
+"""  # noqa: E501
 
 import numpy as np
 from keras import ops
@@ -8,7 +30,7 @@ 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.utils import translate
+from zea.tensor_ops import translate
 
 
 @model_registry(name="echonetlvh")
@@ -18,14 +40,16 @@ class EchoNetLVH(BaseModel):
 
     This model performs semantic segmentation on echocardiogram images to identify
     key anatomical landmarks for measuring left ventricular wall thickness:
-    - LVPWd_1: Left Ventricular Posterior Wall point 1
-    - LVPWd_2: Left Ventricular Posterior Wall point 2
-    - IVSd_1: Interventricular Septum point 1
-    - IVSd_2: Interventricular Septum point 2
+
+    - **LVPWd_1**: Left Ventricular Posterior Wall point 1
+    - **LVPWd_2**: Left Ventricular Posterior Wall point 2
+    - **IVSd_1**: Interventricular Septum point 1
+    - **IVSd_2**: Interventricular Septum point 2
 
     The model outputs 4-channel logits corresponding to heatmaps for each landmark.
 
-    For more information, see the original project page at https://echonet.github.io/lvh/index.html
+    For more information, see the original project page:
+    https://echonet.github.io/lvh/
     """
 
     def __init__(self, **kwargs):

zea/models/lpips.py

@@ -1,7 +1,24 @@
 """LPIPS model for perceptual similarity.
 
-See original code: https://github.com/richzhang/PerceptualSimilarity
-As well as the paper: https://arxiv.org/abs/1801.03924
+To try this model, simply load one of the available presets:
+
+.. doctest::
+
+    >>> from zea.models.lpips import LPIPS
+
+    >>> model = LPIPS.from_preset("lpips")
+
+.. important::
+    This is a ``zea`` implementation of the model.
+    For the original paper and code, see `here <https://github.com/richzhang/PerceptualSimilarity>`_.
+
+    Zhang, Richard, et al.
+    "The Unreasonable Effectiveness of Deep Features as a Perceptual Metric."
+    *https://arxiv.org/abs/1801.03924*
+
+.. seealso::
+    A tutorial notebook where this model is used:
+    :doc:`../notebooks/metrics/lpips_example`.
 
 """
 

zea/models/lv_segmentation.py

@@ -1,23 +1,40 @@
 """
-The model is the nnU-Net model trained on the augmented CAMUS dataset from the following publication:
-Van De Vyver, Gilles, et al.
-"Generative augmentations for improved cardiac ultrasound segmentation using diffusion models."
-arXiv preprint arXiv:2502.20100 (2025).
+nnU-Net segmentation model trained on the augmented CAMUS dataset.
 
-GitHub original repo: https://github.com/GillesVanDeVyver/EchoGAINS
+To try this model, simply load one of the available presets:
+
+.. doctest::
+
+    >>> from zea.models.lv_segmentation import AugmentedCamusSeg
+
+    >>> model = AugmentedCamusSeg.from_preset("augmented_camus_seg")
+
+The model segments both the left ventricle and myocardium.
 
 At the time of writing (17 September 2025) and to the best of our knowledge,
 it is the state-of-the-art model for left ventricle segmentation on the CAMUS dataset.
 
-The model is originally a PyTorch model converted to ONNX. The model segments the left ventricle and myocardium.
+.. important::
+    This is a ``zea`` implementation of the model.
+    For the original paper and code, see `here <https://github.com/GillesVanDeVyver/EchoGAINS>`_.
+
+    Van De Vyver, Gilles, et al.
+    "Generative augmentations for improved cardiac ultrasound segmentation using diffusion models."
+    *https://arxiv.org/abs/2502.20100*
+
+.. seealso::
+    A tutorial notebook where this model is used:
+    :doc:`../notebooks/models/left_ventricle_segmentation_example`.
+
+.. note::
+    The model is originally a PyTorch model converted to ONNX. To use this model, you must have `onnxruntime` installed. This is required for ONNX model inference.
+
+    You can install it using pip:
 
-Note:
------
-To use this model, you must install the `onnxruntime` Python package:
+    .. code-block:: bash
 
-    pip install onnxruntime
+        pip install onnxruntime
 
-This is required for ONNX model inference.
 """  # noqa: E501
 
 from keras import ops

zea/models/preset_utils.py

@@ -136,7 +136,7 @@ def load_json(preset, config_file=CONFIG_FILE):
     return config
 
 
-def load_serialized_object(config, **kwargs):
+def load_serialized_object(config, cls, **kwargs):
     """Load a serialized Keras object from a config."""
     # `dtype` in config might be a serialized `DTypePolicy` or `DTypePolicyMap`.
     # Ensure that `dtype` is properly configured.
@@ -145,7 +145,7 @@ def load_serialized_object(config, **kwargs):
 
     config["config"] = {**config["config"], **kwargs}
     # return keras.saving.deserialize_keras_object(config)
-    return zea.models.base.deserialize_zea_object(config)
+    return zea.models.base.deserialize_zea_object(config, cls)
 
 
 def check_config_class(config):
@@ -275,7 +275,7 @@ class KerasPresetLoader(PresetLoader):
 
     def load_model(self, cls, load_weights, **kwargs):
         """Load a model from a serialized Keras config."""
-        model = load_serialized_object(self.config, **kwargs)
+        model = load_serialized_object(self.config, cls=cls, **kwargs)
 
         if not load_weights:
             return model
@@ -306,7 +306,7 @@ class KerasPresetLoader(PresetLoader):
     def load_image_converter(self, cls, **kwargs):
         """Load an image converter from the preset."""
         converter_config = load_json(self.preset, IMAGE_CONVERTER_CONFIG_FILE)
-        return load_serialized_object(converter_config, **kwargs)
+        return load_serialized_object(converter_config, cls, **kwargs)
 
     def get_file(self, path):
         """Get a file from the preset."""
@@ -322,7 +322,7 @@ class KerasPresetLoader(PresetLoader):
         if not issubclass(check_config_class(preprocessor_json), cls):
             return super().load_preprocessor(cls, **kwargs)
         # We found a `preprocessing.json` with a complete config for our class.
-        preprocessor = load_serialized_object(preprocessor_json, **kwargs)
+        preprocessor = load_serialized_object(preprocessor_json, cls, **kwargs)
         if hasattr(preprocessor, "load_preset_assets"):
             preprocessor.load_preset_assets(self.preset)
         return preprocessor

zea/models/regional_quality.py

@@ -1,21 +1,41 @@
 """
-The model is the movilenetV2 based image quality model from:
-Van De Vyver, et al. "Regional Image Quality Scoring for 2-D Echocardiography Using Deep Learning."
-Ultrasound in Medicine & Biology 51.4 (2025): 638-649.
+MobileNetv2 based image quality model for myocardial regions in apical views.
 
-GitHub original repo: https://github.com/GillesVanDeVyver/arqee
+To try this model, simply load one of the available presets:
 
-The model is originally a PyTorch model converted to ONNX. The model predicts the regional image quality of
+.. doctest::
+
+    >>> from zea.models.regional_quality import MobileNetv2RegionalQuality
+
+    >>> model = MobileNetv2RegionalQuality.from_preset("mobilenetv2_regional_quality")
+
+The model predicts the regional image quality of
 the myocardial regions in apical views. It can also be used to get the overall image quality by averaging the
 regional scores.
 
-Note:
------
-To use this model, you must install the `onnxruntime` Python package:
+At the time of writing (17 September 2025) and to the best of our knowledge,
+it is the state-of-the-art model for left ventricle segmentation on the CAMUS dataset.
+
+.. important::
+    This is a ``zea`` implementation of the model.
+    For the original paper and code, see `here <https://github.com/GillesVanDeVyver/arqee>`_.
+
+    Van De Vyver, et al. "Regional Image Quality Scoring for 2-D Echocardiography Using Deep Learning."
+    *Ultrasound in Medicine & Biology 51.4 (2025): 638-649*
+
+.. seealso::
+    A tutorial notebook where this model is used:
+    :doc:`../notebooks/metrics/myocardial_quality_example`.
+
+.. note::
+    The model is originally a PyTorch model converted to ONNX. To use this model, you must have `onnxruntime` installed. This is required for ONNX model inference.
+
+    You can install it using pip:
+
+    .. code-block:: bash
 
-    pip install onnxruntime
+        pip install onnxruntime
 
-This is required for ONNX model inference.
 """  # noqa: E501
 
 import numpy as np

zea/models/taesd.py

@@ -1,9 +1,19 @@
-"""Tiny Autoencoder (TAESD) model converted to Tensorflow.
+"""
+Tiny Autoencoder (TAESD) model.
+
+.. doctest::
+
+    >>> from zea.models.taesd import TinyAutoencoder
+
+    >>> model = TinyAutoencoder.from_preset("taesdxl")  # doctest: +SKIP
 
-For the original implementation, see the `TAESD repository <https://github.com/madebyollin/taesd>`_.
+.. important::
+    This is a ``zea`` implementation of the model.
+    For the original code, see `here <https://github.com/madebyollin/taesd>`_.
 
-You can see an example of how to use this model in the example notebook:
-:doc:`../notebooks/models/taesd_autoencoder_example`.
+.. seealso::
+    A tutorial notebook where this model is used:
+    :doc:`../notebooks/models/taesd_autoencoder_example`.
 
 """
 
@@ -23,7 +33,13 @@ tf = _import_tf()
 
 @model_registry(name="taesdxl")
 class TinyAutoencoder(BaseModel):
-    """[TAESD](https://github.com/madebyollin/taesd) model in TensorFlow."""
+    """Tiny Autoencoder model.
+
+    .. note::
+
+        This model currently only supports TensorFlow and Jax backends.
+
+    """
 
     def __init__(self, **kwargs):
         """

zea/models/unet.py

@@ -1,4 +1,18 @@
-"""UNet models and architectures"""
+"""UNet models and architectures.
+
+To try this model, simply load one of the available presets:
+
+.. doctest::
+
+    >>> from zea.models.unet import UNet
+
+    >>> model = UNet.from_preset("unet-echonet-inpainter")
+
+.. seealso::
+    A tutorial notebook where this model is used:
+    :doc:`../notebooks/models/unet_example`.
+
+"""
 
 import keras
 from keras import layers