zea 0.0.5__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:
@@ -823,12 +840,12 @@ class DPS(DiffusionGuidance):
         omega,
         **kwargs,
     ):
-        """Compute measurement error for diffusion posterior sampling.
+        """
+        Compute measurement error for diffusion posterior sampling.
 
         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.
@@ -849,20 +866,164 @@ class DPS(DiffusionGuidance):
         return measurement_error, (pred_noises, pred_images)
 
     def __call__(self, noisy_images, **kwargs):
-        """Call the gradient function.
-
-        Returns a function with the following signature:
-            (
-                noisy_images,
-                measurement,
-                operator,
-                noise_rates,
-                signal_rates,
-                omega,
-                **operator_kwargs,
-            ) -> gradients, (error, (pred_noises, pred_images))
+        """
+        Call the gradient function.
 
-        where operator_kwargs are the kwargs for the operator.
+        Args:
+            noisy_images: Noisy images.
+            measurement: Target measurement.
+            operator: Forward operator.
+            noise_rates: Current noise rates.
+            signal_rates: Current signal rates.
+            omega: Weight for the measurement error.
+            **kwargs: Additional arguments for the operator.
 
+        Returns:
+            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):
@@ -37,17 +61,6 @@ class EchoNetLVH(BaseModel):
         """
         super().__init__(**kwargs)
 
-        # Scan conversion constants for echonet processing
-        self.rho_range = (0, 224)  # Radial distance range in pixels
-        self.theta_range = (np.deg2rad(-45), np.deg2rad(45))  # Angular range in radians
-        self.fill_value = -1.0  # Fill value for scan conversion
-        self.resolution = 1.0  # mm per pixel resolution
-
-        # Network input/output dimensions
-        self.n_rho = 224
-        self.n_theta = 224
-        self.output_shape = (224, 224, 4)
-
         # Pre-computed coordinate grid for efficient processing
         self.coordinate_grid = ops.stack(
             ops.cast(ops.convert_to_tensor(np.indices((224, 224))), "float32"), axis=-1

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

@@ -0,0 +1,96 @@
+"""
+nnU-Net segmentation model trained on the augmented CAMUS dataset.
+
+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.
+
+.. 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:
+
+    .. code-block:: bash
+
+        pip install onnxruntime
+
+"""  # noqa: E501
+
+from keras import ops
+
+from zea.internal.registry import model_registry
+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
+
+
+@model_registry(name="augmented_camus_seg")
+class AugmentedCamusSeg(BaseModel):
+    """
+    nnU-Net based left ventricle and myocardium segmentation model.
+
+    - Trained on the augmented CAMUS dataset.
+    - This class loads an ONNX model and provides inference for cardiac ultrasound segmentation tasks.
+
+    """  # noqa: E501
+
+    def call(self, inputs):
+        """
+        Run inference on the input data using the loaded ONNX model.
+
+        Args:
+            inputs (np.ndarray): Input image or batch of images for segmentation.
+                Shape: [batch, 1, 256, 256]
+                Range: Any numeric range; normalized internally.
+
+        Returns:
+            np.ndarray: Segmentation mask(s) for left ventricle and myocardium.
+                Shape: [batch, 3, 256, 256]  (logits for background, LV, myocardium)
+
+        Raises:
+            ValueError: If model weights are not loaded.
+        """
+        if not hasattr(self, "onnx_sess"):
+            raise ValueError("Model weights not loaded. Please call custom_load_weights() first.")
+        input_name = self.onnx_sess.get_inputs()[0].name
+        output_name = self.onnx_sess.get_outputs()[0].name
+        inputs = ops.convert_to_numpy(inputs).astype("float32")
+        output = self.onnx_sess.run([output_name], {input_name: inputs})[0]
+        return output
+
+    def custom_load_weights(self, preset, **kwargs):
+        """Load the ONNX weights for the segmentation model."""
+        try:
+            import onnxruntime
+        except ImportError:
+            raise ImportError(
+                "onnxruntime is not installed. Please run "
+                "`pip install onnxruntime` to use this model."
+            )
+        loader = get_preset_loader(preset)
+        filename = loader.get_file("model.onnx")
+        self.onnx_sess = onnxruntime.InferenceSession(filename)
+
+
+register_presets(augmented_camus_seg_presets, AugmentedCamusSeg)

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/presets.py

@@ -47,6 +47,34 @@ echonet_dynamic_presets = {
     },
 }
 
+augmented_camus_seg_presets = {
+    "augmented_camus_seg": {
+        "metadata": {
+            "description": (
+                "Augmented CAMUS segmentation model for cardiac ultrasound segmentation. "
+                "Original paper and code: https://arxiv.org/abs/2502.20100"
+            ),
+            "params": 33468899,
+            "path": "lv_segmentation",
+        },
+        "hf_handle": "hf://zeahub/augmented-camus-segmentation",
+    },
+}
+
+regional_quality_presets = {
+    "mobilenetv2_regional_quality": {
+        "metadata": {
+            "description": (
+                "MobileNetV2-based regional myocardial image quality scoring model. "
+                "Original GitHub repository and code: https://github.com/GillesVanDeVyver/arqee"
+            ),
+            "params": 2217064,
+            "path": "regional_quality",
+        },
+        "hf_handle": "hf://zeahub/mobilenetv2-regional-quality",
+    }
+}
+
 echonet_lvh_presets = {
     "echonetlvh": {
         "metadata": {
@@ -97,6 +125,14 @@ diffusion_model_presets = {
         },
         "hf_handle": "hf://zeahub/diffusion-echonet-dynamic",
     },
+    "diffusion-echonetlvh-3-frame": {
+        "metadata": {
+            "description": ("3-frame diffusion model trained on EchoNetLVH dataset."),
+            "params": 0,
+            "path": "diffusion",
+        },
+        "hf_handle": "hf://zeahub/diffusion-echonetlvh",
+    },
 }
 
 carotid_segmenter_presets = {