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

zea/models/diffusion.py

@@ -728,15 +728,15 @@ class DiffusionModel(DeepGenerativeModel):
             )
         return next_noisy_images
 
-    def start_track_progress(self, diffusion_steps):
+    def start_track_progress(self, diffusion_steps, initial_step=0):
         """Initialize the progress tracking for the diffusion process.
-
         For diffusion animation we keep track of the diffusion progress.
         For large number of steps, we do not store all the images due to memory constraints.
         """
         self.track_progress = []
-        if diffusion_steps > 50:
-            self.track_progress_interval = diffusion_steps // 50
+        remaining = max(1, diffusion_steps - int(initial_step))
+        if remaining > 50:
+            self.track_progress_interval = remaining // 50
         else:
             self.track_progress_interval = 1
 
@@ -823,7 +823,8 @@ class DPS(DiffusionGuidance):
         omega,
         **kwargs,
     ):
-        """Compute measurement error for diffusion posterior sampling.
+        """
+        Compute measurement error for diffusion posterior sampling.
 
         Args:
             noisy_images: Noisy images.
@@ -849,20 +850,19 @@ 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)

zea/models/echonetlvh.py

@@ -0,0 +1,279 @@
+"""EchoNetLVH model for segmentation of PLAX view cardiac ultrasound. For more details see https://echonet.github.io/lvh/index.html."""
+
+import numpy as np
+from keras import ops
+
+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.utils import translate
+
+
+@model_registry(name="echonetlvh")
+class EchoNetLVH(BaseModel):
+    """
+    EchoNet Left Ventricular Hypertrophy (LVH) model for echocardiogram analysis.
+
+    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
+
+    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
+    """
+
+    def __init__(self, **kwargs):
+        """
+        Initialize the EchoNetLVH model.
+
+        Args:
+            **kwargs: Additional keyword arguments passed to BaseModel
+        """
+        super().__init__(**kwargs)
+
+        # 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
+        )
+
+        # Initialize the underlying segmentation network
+        self.network = DeeplabV3Plus(image_shape=(224, 224, 3), num_classes=4)
+
+    def call(self, inputs):
+        """
+        Forward pass of the model.
+
+        Args:
+            inputs (Tensor): Input images of shape [B, H, W, C]. They should
+                be scan converted, with pixel values in range [0, 255].
+
+        Returns:
+            Tensor: Logits of shape [B, H, W, 4] with 4 channels for each landmark
+        """
+        assert len(ops.shape(inputs)) == 4
+
+        # Store original dimensions for output resizing
+        original_size = ops.shape(inputs)[1:3]
+
+        # Resize to network input size
+        inputs_resized = ops.image.resize(inputs, size=(224, 224))
+
+        # Get network predictions
+        logits = self.network(inputs_resized)
+
+        # Resize logits back to original input dimensions
+        logits_output = ops.image.resize(logits, original_size)
+        return logits_output
+
+    def extract_key_points_as_indices(self, logits):
+        """
+        Extract key point coordinates from logits using center-of-mass calculation.
+
+        Args:
+            logits (Tensor): Model output logits of shape [B, H, W, 4]
+
+        Returns:
+            Tensor: Key point coordinates of shape [B, 4, 2] where each point
+                   is in (x, y) format
+        """
+        # Create coordinate grid for the current logit dimensions
+        input_shape = ops.shape(logits)[1:3]
+        input_space_coordinate_grid = ops.stack(
+            ops.cast(ops.convert_to_tensor(np.indices(input_shape)), "float32"), axis=-1
+        )
+
+        # Transpose logits to [B, 4, H, W] for vectorized processing
+        logits_batchified = ops.transpose(logits, (0, 3, 1, 2))
+
+        # Extract expected coordinates for each channel
+        return ops.flip(
+            ops.vectorized_map(
+                lambda logit: self.expected_coordinate(logit, input_space_coordinate_grid),
+                logits_batchified,
+            ),
+            axis=-1,  # Flip to convert from (y, x) to (x, y)
+        )
+
+    def expected_coordinate(self, mask, coordinate_grid=None):
+        """
+        Compute the expected coordinate (center-of-mass) of a heatmap.
+
+        This implements a differentiable version of taking the max of a heatmap
+        by computing the weighted average of coordinates.
+
+        Reference: https://arxiv.org/pdf/1711.08229
+
+        Args:
+            mask (Tensor): Heatmap of shape [B, H, W]
+            coordinate_grid (Tensor, optional): Grid of coordinates. If None,
+                                              uses self.coordinate_grid
+
+        Returns:
+            Tensor: Expected coordinates of shape [B, 2] in (x, y) format
+        """
+        if coordinate_grid is None:
+            coordinate_grid = self.coordinate_grid
+
+        # Ensure mask values are non-negative and normalized
+        mask_clipped = ops.clip(mask, 0, None)
+        mask_normed = mask_clipped / ops.max(mask_clipped)
+
+        def safe_normalize(m):
+            mask_sum = ops.sum(m)
+            return ops.where(mask_sum > 0, m / mask_sum, m)
+
+        coordinate_probabilities = ops.map(safe_normalize, mask_normed)
+
+        # Add dimension for broadcasting with coordinate grid
+        coordinate_probabilities = ops.expand_dims(coordinate_probabilities, axis=-1)
+
+        # Compute weighted average of coordinates
+        expected_coordinate = ops.sum(
+            ops.expand_dims(coordinate_grid, axis=0) * coordinate_probabilities,
+            axis=(1, 2),
+        )
+
+        # Flip to convert from (y, x) to (x, y) format for euclidean distance calculation
+        return ops.flip(expected_coordinate, axis=-1)
+
+    def overlay_labels_on_image(self, image, label, alpha=0.5):
+        """
+        Overlay predicted heatmaps and connecting lines on the input image.
+
+        Args:
+            image (Tensor): Input image of shape [H, W] or [H, W, C]
+            label (Tensor): Predicted logits of shape [H, W, 4]
+            alpha (float): Blending factor for overlay (0=transparent, 1=opaque)
+
+        Returns:
+            ndarray: Image with overlaid heatmaps and measurements of shape [H, W, 3]
+        """
+        try:
+            import cv2
+
+        except ImportError as exc:
+            raise ImportError(
+                "OpenCV is required for `EchoNetLVH.overlay_labels_on_image`. "
+                "Please install it with 'pip install opencv-python' or "
+                "'pip install opencv-python-headless'."
+            ) from exc
+
+        # Color scheme for each landmark
+        overlay_colors = np.array(
+            [
+                [1, 1, 0],  # Yellow (LVPWd_X1)
+                [1, 0, 1],  # Magenta (LVPWd_X2)
+                [0, 1, 1],  # Cyan (IVSd_X1)
+                [0, 1, 0],  # Green (IVSd_X2)
+            ],
+        )
+
+        # Convert to numpy and ensure RGB format
+        image = ops.convert_to_numpy(image)
+        label = ops.convert_to_numpy(label)
+
+        if image.ndim == 2:
+            image = np.stack([image] * 3, axis=-1)
+        elif image.shape[-1] == 1:
+            image = np.repeat(image, 3, axis=-1)
+        else:
+            image = image.copy()
+
+        # Normalize each channel to [0, 1] for proper visualization
+        label = np.clip(label, 0, None)
+        for ch in range(label.shape[-1]):
+            max_val = np.max(label[..., ch])
+            if max_val > 0:
+                label[..., ch] = label[..., ch] / max_val
+
+        # Initialize overlay and tracking variables
+        overlay = np.zeros_like(image, dtype=np.float32)
+        centers = []
+
+        # Process each landmark channel
+        for ch in range(4):
+            # Square the mask to enhance peak responses
+            mask = label[..., ch] ** 2
+            color = overlay_colors[ch]
+
+            # Find center of mass for this channel
+            center_coords = self.expected_coordinate(ops.expand_dims(mask, axis=0))
+            center_x = ops.convert_to_numpy(center_coords[0, 0])
+            center_y = ops.convert_to_numpy(center_coords[0, 1])
+
+            # Bounds check before conversion to int
+            if 0 <= center_x < image.shape[1] and 0 <= center_y < image.shape[0]:
+                center = (int(center_x), int(center_y))
+            else:
+                center = None
+
+            if center is not None:
+                # Blend heatmap with overlay
+                mask_alpha = mask * alpha
+                for c in range(3):
+                    overlay[..., c] += mask_alpha * color[c]
+            centers.append(center)
+
+        # Draw connecting lines between consecutive landmarks
+        for i in range(3):
+            pt1, pt2 = centers[i], centers[i + 1]
+            if pt1 is not None and pt2 is not None:
+                color = tuple(int(x) for x in overlay_colors[i])
+
+                # Create line mask
+                line_mask = np.zeros(image.shape[:2], dtype=np.uint8)
+                cv2.line(line_mask, pt1, pt2, color=1, thickness=2)
+
+                # Apply line to overlay
+                for c in range(3):
+                    overlay[..., c][line_mask.astype(bool)] = color[c] * alpha
+
+        # Blend overlay with original image
+        overlay = np.clip(overlay, 0, 1)
+        out = image.astype(np.float32)
+        blend_mask = np.any(overlay > 0.02, axis=-1)
+        out[blend_mask] = (1 - alpha) * out[blend_mask] + overlay[blend_mask]
+
+        return np.clip(out, 0, 1)
+
+    def visualize_logits(self, images, logits):
+        """
+        Create visualization of model predictions overlaid on input images.
+
+        Args:
+            images (Tensor): Input images of shape [B, H, W, C]
+            logits (Tensor): Model predictions of shape [B, H, W, 4]
+
+        Returns:
+            Tensor: Images with overlaid predictions of shape [B, H, W, 3]
+        """
+        # Store original dimensions for final output
+        original_size = ops.shape(images)[1:3]
+
+        # Resize to standard processing size
+        images_resized = ops.image.resize(images, size=(224, 224), interpolation="nearest")
+        logits_resized = ops.image.resize(logits, size=(224, 224), interpolation="nearest")
+
+        # Normalize images to [0, 1] range
+        images_clipped = ops.clip(images_resized, 0, 255)
+        images = translate(images_clipped, range_from=(0, 255), range_to=(0, 1))
+
+        # Generate overlays for each image in the batch
+        images_with_overlay = []
+        for img, logit_heatmap in zip(images, logits_resized):
+            overlay = self.overlay_labels_on_image(img, logit_heatmap)
+            images_with_overlay.append(overlay)
+
+        # Stack results and resize back to original dimensions
+        images_with_overlay = np.stack(images_with_overlay, axis=0)
+        return ops.image.resize(images_with_overlay, original_size)
+
+
+# Register model presets
+register_presets(echonet_lvh_presets, EchoNetLVH)

zea/models/lv_segmentation.py

@@ -0,0 +1,79 @@
+"""
+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).
+
+GitHub original repo: https://github.com/GillesVanDeVyver/EchoGAINS
+
+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.
+
+Note:
+-----
+To use this model, you must install the `onnxruntime` Python package:
+
+    pip install onnxruntime
+
+This is required for ONNX model inference.
+"""  # 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/presets.py

@@ -47,6 +47,48 @@ 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": {
+            "description": (
+                "EchoNetLVH segmentation model for PLAX-view cardiac ultrasound segmentation. "
+                "Trained on images of size (224, 224)."
+            ),
+            "params": 0,
+            "path": "echonetlvh",
+        },
+        "hf_handle": "hf://zeahub/echonetlvh",
+    },
+}
+
 lpips_presets = {
     "lpips": {
         "metadata": {
@@ -83,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 = {

zea/models/regional_quality.py

@@ -0,0 +1,122 @@
+"""
+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.
+
+GitHub original repo: https://github.com/GillesVanDeVyver/arqee
+
+The model is originally a PyTorch model converted to ONNX. 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:
+
+    pip install onnxruntime
+
+This is required for ONNX model inference.
+"""  # 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)