zea 0.0.5__tar.gz → 0.0.6__tar.gz

{zea-0.0.5 → zea-0.0.6}/PKG-INFO

@@ -1,7 +1,8 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: zea
-Version: 0.0.5
+Version: 0.0.6
 Summary: A Toolbox for Cognitive Ultrasound Imaging. Provides a set of tools for processing of ultrasound data, all built in your favorite machine learning framework.
+License-File: LICENSE
 Keywords: ultrasound,machine learning,beamforming
 Author: Tristan Stevens
 Author-email: t.s.w.stevens@tue.nl
@@ -21,6 +22,7 @@ Provides-Extra: display
 Provides-Extra: display-headless
 Provides-Extra: docs
 Provides-Extra: jax
+Provides-Extra: models
 Provides-Extra: tests
 Requires-Dist: IPython ; extra == "dev"
 Requires-Dist: IPython ; extra == "docs"
@@ -49,6 +51,8 @@ Requires-Dist: myst-parser ; extra == "docs"
 Requires-Dist: nbsphinx ; extra == "dev"
 Requires-Dist: nbsphinx ; extra == "docs"
 Requires-Dist: numpy (>=1.24)
+Requires-Dist: onnxruntime (>=1.15) ; extra == "dev"
+Requires-Dist: onnxruntime (>=1.15) ; extra == "models"
 Requires-Dist: opencv-python (>=4) ; extra == "display"
 Requires-Dist: opencv-python-headless (>=4) ; extra == "dev"
 Requires-Dist: opencv-python-headless (>=4) ; extra == "display-headless"

{zea-0.0.5 → zea-0.0.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "zea"
-version = "0.0.5"
+version = "0.0.6"
 description = "A Toolbox for Cognitive Ultrasound Imaging. Provides a set of tools for processing of ultrasound data, all built in your favorite machine learning framework."
 authors = [
     { name = "Tristan Stevens", email = "t.s.w.stevens@tue.nl" },
@@ -74,6 +74,8 @@ dev = [
     "IPython",
     # display
     "opencv-python-headless>=4",
+    # models
+    "onnxruntime>=1.15",
 ]
 
 tests = [
@@ -110,6 +112,11 @@ display = [
 display-headless = [
     "opencv-python-headless>=4",
 ]
+# most models don't needs these optional dependencies
+# but some do, so we list them here
+models = [
+    "onnxruntime>=1.15",
+]
 # these are just here for .readthedocs.yaml
 # please for proper install (with GPU support)
 # install manually

{zea-0.0.5 → zea-0.0.6}/zea/__init__.py

@@ -7,7 +7,7 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.5"
+__version__ = "0.0.6"
 
 
 def _bootstrap_backend():

{zea-0.0.5 → zea-0.0.6}/zea/agent/selection.py

@@ -11,11 +11,14 @@ For a comprehensive example usage, see: :doc:`../notebooks/agent/agent_example`
 All strategies are stateless, meaning that they do not maintain any internal state.
 """
 
+from typing import Callable
+
 import keras
 from keras import ops
 
 from zea import tensor_ops
 from zea.agent import masks
+from zea.backend.autograd import AutoGrad
 from zea.internal.registry import action_selection_registry
 
 
@@ -493,3 +496,166 @@ class CovarianceSamplingLines(LinesActionModel):
         best_mask = ops.squeeze(best_mask, axis=0)
 
         return best_mask, self.lines_to_im_size(best_mask)
+
+
+class TaskBasedLines(GreedyEntropy):
+    """Task-based line selection for maximizing information gain.
+
+    This action selection strategy chooses lines to maximize information gain with respect
+    to a downstream task outcome. It uses gradient-based saliency to identify which image
+    regions contribute most to task uncertainty, then selects lines accordingly.
+    """
+
+    def __init__(
+        self,
+        n_actions: int,
+        n_possible_actions: int,
+        img_width: int,
+        img_height: int,
+        downstream_task_function: Callable,
+        mean: float = 0,
+        std_dev: float = 1,
+        num_lines_to_update: int = 5,
+        **kwargs,
+    ):
+        """Initialize the TaskBasedLines action selection model.
+
+        Args:
+            n_actions (int): The number of actions the agent can take.
+            n_possible_actions (int): The number of possible actions (line positions).
+            img_width (int): The width of the input image.
+            img_height (int): The height of the input image.
+            downstream_task_function (Callable): A differentiable function that takes a
+                batch of inputs and produces scalar outputs. This represents the downstream
+                task for which information gain should be maximized.
+            mean (float, optional): The mean of the RBF used for reweighting. Defaults to 0.
+            std_dev (float, optional): The standard deviation of the RBF used for reweighting.
+                Defaults to 1.
+            num_lines_to_update (int, optional): The number of lines around the selected line
+                to update during reweighting. Must be odd. Defaults to 5.
+            **kwargs: Additional keyword arguments passed to the parent class.
+        """
+        super().__init__(
+            n_actions,
+            n_possible_actions,
+            img_width,
+            img_height,
+            mean,
+            std_dev,
+            num_lines_to_update,
+        )
+        self.downstream_task_function = downstream_task_function
+
+    def compute_output_and_saliency_propagation(self, particles):
+        """Compute saliency-weighted posterior variance for task-based selection.
+
+        This method computes how much each pixel contributes to the variance of the
+        downstream task output. It uses automatic differentiation to compute gradients
+        of the task function with respect to each particle, then weights the posterior
+        variance by the squared mean gradient.
+
+        Args:
+            particles (Tensor): Particles of shape (batch_size, n_particles, height, width)
+                representing the posterior distribution over images.
+
+        Returns:
+            Tensor: Pixelwise contribution to downstream task variance,
+                of shape (batch_size, height, width). Higher values indicate pixels
+                that contribute more to task uncertainty.
+        """
+        autograd = AutoGrad()
+
+        autograd.set_function(self.downstream_task_function)
+        downstream_grad_and_value_fn = autograd.get_gradient_and_value_jit_fn()
+        jacobian, _ = ops.vectorized_map(
+            lambda p: ops.vectorized_map(
+                downstream_grad_and_value_fn,
+                p,
+            ),
+            particles,
+        )
+
+        posterior_variance = ops.var(particles, axis=1)
+        mean_jacobian = ops.mean(jacobian, axis=1)
+        return posterior_variance * (mean_jacobian**2)
+
+    def sum_neighbouring_columns_into_n_possible_actions(self, full_linewise_salience):
+        """Aggregate column-wise saliency into line-wise saliency scores.
+
+        This method groups neighboring columns together to create saliency scores
+        for each possible line action. Since each line action may correspond to
+        multiple image columns, this aggregation is necessary to match the action space.
+
+        Args:
+            full_linewise_salience (Tensor): Saliency values for each column,
+                of shape (batch_size, full_image_width).
+
+        Returns:
+            Tensor: Aggregated saliency scores for each possible action,
+                of shape (batch_size, n_possible_actions).
+
+        Raises:
+            AssertionError: If the image width is not evenly divisible by n_possible_actions.
+        """
+        batch_size = ops.shape(full_linewise_salience)[0]
+        full_image_width = ops.shape(full_linewise_salience)[1]
+        assert full_image_width % self.n_possible_actions == 0, (
+            "n_possible_actions must divide evenly into image width"
+        )
+        cols_per_action = full_image_width // self.n_possible_actions
+        stacked_linewise_salience = ops.reshape(
+            full_linewise_salience,
+            (batch_size, self.n_possible_actions, cols_per_action),
+        )
+        return ops.sum(stacked_linewise_salience, axis=2)
+
+    def sample(self, particles):
+        """Sample actions using task-based information gain maximization.
+
+        This method computes which lines would provide the most information about
+        the downstream task by:
+        1. Computing pixelwise contribution to task variance using gradients
+        2. Aggregating contributions into line-wise scores
+        3. Greedily selecting lines with highest contribution scores
+        4. Reweighting scores around selected lines (inherited from GreedyEntropy)
+
+        Args:
+            particles (Tensor): Particles representing the posterior distribution,
+                of shape (batch_size, n_particles, height, width).
+
+        Returns:
+            Tuple[Tensor, Tensor, Tensor]:
+                - selected_lines_k_hot: Selected lines as k-hot vectors,
+                  shaped (batch_size, n_possible_actions)
+                - masks: Binary masks of shape (batch_size, img_height, img_width)
+                - pixelwise_contribution_to_var_dst: Pixelwise contribution to downstream
+                  task variance, of shape (batch_size, height, width)
+
+        Note:
+            Unlike the parent GreedyEntropy class, this method returns an additional
+            tensor containing the pixelwise contribution scores for analysis.
+        """
+        pixelwise_contribution_to_var_dst = self.compute_output_and_saliency_propagation(particles)
+        linewise_contribution_to_var_dst = ops.sum(pixelwise_contribution_to_var_dst, axis=1)
+        actionwise_contribution_to_var_dst = self.sum_neighbouring_columns_into_n_possible_actions(
+            linewise_contribution_to_var_dst
+        )
+
+        # Greedily select best line, reweight entropies, and repeat
+        all_selected_lines = []
+        for _ in range(self.n_actions):
+            max_contribution_line, actionwise_contribution_to_var_dst = ops.vectorized_map(
+                self.select_line_and_reweight_entropy,
+                actionwise_contribution_to_var_dst,
+            )
+            all_selected_lines.append(max_contribution_line)
+
+        selected_lines_k_hot = ops.any(
+            ops.one_hot(all_selected_lines, self.n_possible_actions, dtype=masks._DEFAULT_DTYPE),
+            axis=0,
+        )
+        return (
+            selected_lines_k_hot,
+            self.lines_to_im_size(selected_lines_k_hot),
+            pixelwise_contribution_to_var_dst,
+        )

{zea-0.0.5 → zea-0.0.6}/zea/backend/__init__.py

@@ -25,6 +25,8 @@ Key Features
 
 """
 
+from contextlib import nullcontext
+
 import keras
 
 from zea import log
@@ -114,3 +116,90 @@ def _jit_compile(func, jax=True, tensorflow=True, **kwargs):
         log.warning("Initialize zea.Pipeline with jit_options=None to suppress this warning.")
         log.warning("Falling back to non-compiled mode.")
         return func
+
+
+class on_device:
+    """Context manager to set the device regardless of backend.
+
+    For the `torch` backend, you need to manually move the model and data to the device before
+    using this context manager.
+
+    Args:
+        device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``.
+
+    Example:
+        .. code-block:: python
+
+            with zea.backend.on_device("gpu:3"):
+                pipeline = zea.Pipeline([zea.keras_ops.Abs()])
+                output = pipeline(data=keras.random.normal((10, 10)))  # output is on "cuda:3"
+    """
+
+    def __init__(self, device: str):
+        self.device = self.get_device(device)
+        self._context = self.get_context(self.device)
+
+    def get_context(self, device):
+        if device is None:
+            return nullcontext()
+
+        if keras.backend.backend() == "tensorflow":
+            import tensorflow as tf
+
+            return tf.device(device)
+
+        if keras.backend.backend() == "jax":
+            import jax
+
+            return jax.default_device(device)
+        if keras.backend.backend() == "torch":
+            import torch
+
+            return torch.device(device)
+
+        return nullcontext()
+
+    def get_device(self, device: str):
+        if device is None:
+            return None
+
+        device = device.lower()
+
+        if keras.backend.backend() == "tensorflow":
+            return device.replace("cuda", "gpu")
+
+        if keras.backend.backend() == "jax":
+            from zea.backend.jax import str_to_jax_device
+
+            device = device.replace("cuda", "gpu")
+            return str_to_jax_device(device)
+
+        if keras.backend.backend() == "torch":
+            return device.replace("gpu", "cuda")
+
+    def __enter__(self):
+        self._context.__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._context.__exit__(exc_type, exc_val, exc_tb)
+
+
+if keras.backend.backend() in ["tensorflow", "jax", "numpy"]:
+
+    def func_on_device(func, device, *args, **kwargs):
+        """Moves all tensor arguments of a function to a specified device before calling it.
+
+        Args:
+            func (callable): Function to be called.
+            device (str): Device to move tensors to.
+            *args: Positional arguments to be passed to the function.
+            **kwargs: Keyword arguments to be passed to the function.
+        Returns:
+            The output of the function.
+        """
+        with on_device(device):
+            return func(*args, **kwargs)
+elif keras.backend.backend() == "torch":
+    from zea.backend.torch import func_on_device
+else:
+    raise ValueError(f"Unsupported backend: {keras.backend.backend()}")

zea-0.0.6/zea/backend/jax/__init__.py

@@ -0,0 +1,33 @@
+"""Jax utilities for zea."""
+
+import jax
+
+
+def str_to_jax_device(device):
+    """Convert a device string to a JAX device.
+    Args:
+        device (str): Device string, e.g. ``'gpu:0'``, or ``'cpu:0'``.
+    Returns:
+        jax.Device: The corresponding JAX device.
+    """
+
+    if not isinstance(device, str):
+        raise ValueError(f"Device must be a string, got {type(device)}")
+
+    device = device.lower().replace("cuda", "gpu")
+
+    device = device.split(":")
+    if len(device) == 2:
+        device_type, device_number = device
+        device_number = int(device_number)
+    else:
+        # if no device number is specified, use the first device
+        device_type = device[0]
+        device_number = 0
+
+    available = jax.devices(device_type)
+    if len(available) == 0:
+        raise ValueError(f"No JAX devices available for type '{device_type}'.")
+    if device_number < 0 or device_number >= len(available):
+        raise ValueError(f"Device '{device}' is not available; JAX devices found: {available}")
+    return available[device_number]

zea-0.0.6/zea/backend/tensorflow/__init__.py

@@ -0,0 +1,17 @@
+"""Tensorflow Ultrasound Beamforming Library.
+
+Initialize modules for registries.
+"""
+
+import sys
+from pathlib import PosixPath
+
+import numpy as np
+
+# Convert PosixPath objects to strings in sys.path
+# this is necessary due to weird TF bug when importing
+sys.path = [str(p) if isinstance(p, PosixPath) else p for p in sys.path]
+
+import tensorflow as tf  # noqa: E402
+
+from .dataloader import make_dataloader  # noqa: E402

zea-0.0.6/zea/backend/torch/__init__.py

@@ -0,0 +1,39 @@
+"""Pytorch Ultrasound Beamforming Library.
+
+Initialize modules for registries.
+"""
+
+import torch
+
+
+def func_on_device(func, device, *args, **kwargs):
+    """Moves all tensor arguments of a function to a specified device before calling it.
+
+    Args:
+        func (callable): Function to be called.
+        device (str or torch.device): Device to move tensors to.
+        *args: Positional arguments to be passed to the function.
+        **kwargs: Keyword arguments to be passed to the function.
+    Returns:
+        The output of the function.
+    """
+    if device is None:
+        return func(*args, **kwargs)
+
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    def move_to_device(x):
+        if isinstance(x, torch.Tensor):
+            return x.to(device)
+        elif isinstance(x, (list, tuple)):
+            return type(x)(move_to_device(i) for i in x)
+        elif isinstance(x, dict):
+            return {k: move_to_device(v) for k, v in x.items()}
+        else:
+            return x
+
+    args = move_to_device(args)
+    kwargs = move_to_device(kwargs)
+
+    return func(*args, **kwargs)

{zea-0.0.5 → zea-0.0.6}/zea/data/layers.py

@@ -25,7 +25,7 @@ class Resizer(TFDataLayer):
     Resize layer for resizing images. Can deal with N-dimensional images.
     Can do resize, center_crop, random_crop and crop_or_pad.
 
-    Can be used in tf.data pipelines.
+    Can be used in `tf.data` pipelines.
     """
 
     def __init__(
@@ -36,7 +36,6 @@ class Resizer(TFDataLayer):
         seed: int | None = None,
         **resize_kwargs,
     ):
-        # noqa: E501
         """
         Initializes the data loader with the specified parameters.
 
@@ -47,7 +46,7 @@ class Resizer(TFDataLayer):
                 ['random_crop'](https://keras.io/api/layers/preprocessing_layers/image_augmentation/random_crop/),
                 ['resize'](https://keras.io/api/layers/preprocessing_layers/image_preprocessing/resizing/),
                 'crop_or_pad': resizes an image to a target width and height by either centrally
-                    cropping the image, padding it evenly with zeros or a combination of both.
+                cropping the image, padding it evenly with zeros or a combination of both.
             resize_axes (tuple | None, optional): The axes along which to resize.
                 Must be of length 2. Defaults to None. In that case, can only process
                 default tensors of shape (batch, height, width, channels), where the

{zea-0.0.5 → zea-0.0.6}/zea/internal/registry.py

@@ -200,7 +200,7 @@ tf_beamformer_registry = RegisterDecorator(items_to_register=["name", "framework
 
 torch_beamformer_registry = RegisterDecorator(items_to_register=["name", "framework"])
 
-metrics_registry = RegisterDecorator(items_to_register=["name", "framework", "supervised"])
+metrics_registry = RegisterDecorator(items_to_register=["name", "paired"])
 
 checks_registry = RegisterDecorator(items_to_register=["data_type"])
 ops_registry = RegisterDecorator(items_to_register=["name"])