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

zea/models/regional_quality.py

@@ -0,0 +1,142 @@
+"""
+MobileNetv2 based image quality model for myocardial regions in apical views.
+
+To try this model, simply load one of the available presets:
+
+.. 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.
+
+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
+
+"""  # noqa: E501
+
+import numpy as np
+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 regional_quality_presets
+
+# Visualization colors and helper for regional quality (arqee-inspired)
+QUALITY_COLORS = np.array(
+    [
+        [0.929, 0.106, 0.141],  # not visible, red
+        [0.957, 0.396, 0.137],  # poor, orange
+        [1, 0.984, 0.090],  # ok, yellow
+        [0.553, 0.776, 0.098],  # good, light green
+        [0.09, 0.407, 0.216],  # excellent, dark green
+    ]
+)
+REGION_LABELS = [
+    "basal_left",
+    "mid_left",
+    "apical_left",
+    "apical_right",
+    "mid_right",
+    "basal_right",
+    "annulus_left",
+    "annulus_right",
+]
+QUALITY_CLASSES = ["not visible", "poor", "ok", "good", "excellent"]
+
+
+@model_registry(name="mobilenetv2_regional_quality")
+class MobileNetv2RegionalQuality(BaseModel):
+    """
+    MobileNetV2 based regional image quality scoring model for myocardial regions in apical views.
+
+    This class loads an ONNX model and provides inference for regional image quality scoring tasks.
+    """
+
+    def preprocess_input(self, inputs):
+        """
+        Normalize input image(s) to [0, 255] range.
+
+        Args:
+            inputs (np.ndarray): Input image(s), any numeric range.
+
+        Returns:
+            np.ndarray: Normalized image(s) in [0, 255] range.
+        """
+        inputs = ops.convert_to_numpy(inputs).astype("float32")
+        max_val = np.max(inputs)
+        min_val = np.min(inputs)
+        denom = max_val - min_val
+        if denom > 0.0:
+            inputs = (inputs - min_val) / denom * 255.0
+        else:
+            inputs = np.zeros_like(inputs, dtype=np.float32)
+        return inputs
+
+    def call(self, inputs):
+        """
+        Predict regional image quality scores for input image(s).
+
+        Args:
+            inputs (np.ndarray): Input image or batch of images.
+            Shape: [batch, 1, 256, 256]
+
+        Returns:
+            np.ndarray: Regional quality scores.
+                Shape is [batch, 8] with regions in order:
+                basal_left, mid_left, apical_left, apical_right,
+                mid_right, basal_right, annulus_left, annulus_right
+        """
+        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 = self.preprocess_input(inputs)
+
+        output = self.onnx_sess.run([output_name], {input_name: inputs})[0]
+        slope = self.slope_intercept[0]
+        intercept = self.slope_intercept[1]
+        output_debiased = (output - intercept) / slope
+        return output_debiased
+
+    def custom_load_weights(self, preset, **kwargs):
+        """Load ONNX model weights and bias correction for regional image quality scoring."""
+        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)
+        filename = loader.get_file("slope_intercept_bias_correction.npy")
+        self.slope_intercept = np.load(filename)
+
+
+register_presets(regional_quality_presets, MobileNetv2RegionalQuality)

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