zea 0.0.0__py3-none-any.whl → 0.0.2__py3-none-any.whl

zea/__init__.py

@@ -0,0 +1,74 @@
+"""``zea``: *A Toolbox for Cognitive Ultrasound Imaging.*"""
+
+import importlib.util
+import os
+
+from . import log
+
+# dynamically add __version__ attribute (see pyproject.toml)
+# __version__ = __import__("importlib.metadata").metadata.version(__package__)
+__version__ = "0.0.2"
+
+
+def setup():
+    """Setup function to initialize the zea package."""
+
+    def _check_backend_installed():
+        """Assert that at least one ML backend (torch, tensorflow, jax) is installed.
+        If not, raise an AssertionError with a helpful install message.
+        """
+
+        ml_backends = ["torch", "tensorflow", "jax"]
+        for backend in ml_backends:
+            if importlib.util.find_spec(backend) is not None:
+                return
+
+        backend_env = os.environ.get("KERAS_BACKEND", "numpy")
+        install_guide_urls = {
+            "torch": "https://pytorch.org/get-started/locally/",
+            "tensorflow": "https://www.tensorflow.org/install",
+            "jax": "https://docs.jax.dev/en/latest/installation.html",
+        }
+        guide_url = install_guide_urls.get(backend_env, "https://keras.io/getting_started/")
+        raise ImportError(
+            "No ML backend (torch, tensorflow, jax) installed in current environment. "
+            f"Please install at least one ML backend before importing {__package__} or "
+            f"any other library. Current KERAS_BACKEND is set to '{backend_env}', "
+            f"please install it first, see: {guide_url}. One simple alternative is to "
+            f"install with default backend: `pip install {__package__}[jax]`."
+        )
+
+    _check_backend_installed()
+
+    import keras
+
+    log.info(f"Using backend {keras.backend.backend()!r}")
+
+
+# call and clean up namespace
+setup()
+del setup
+
+from . import (
+    agent,
+    beamform,
+    data,
+    display,
+    io_lib,
+    metrics,
+    models,
+    simulator,
+    tensor_ops,
+    utils,
+    visualize,
+)
+from .config import Config
+from .data.datasets import Dataset, Folder
+from .data.file import File, load_file
+from .datapaths import set_data_paths
+from .interface import Interface
+from .internal.device import init_device
+from .internal.setup_zea import set_backend, setup, setup_config
+from .ops import Pipeline
+from .probes import Probe
+from .scan import Scan

zea/__main__.py

@@ -0,0 +1,82 @@
+"""Main entry point for zea
+
+Run as `zea --config path/to/config.yaml` to start the zea interface.
+Or do not pass a config file to open a file dialog to choose a config file.
+
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+from zea import log
+from zea.visualize import set_mpl_style
+
+
+def get_args():
+    """Command line argument parser"""
+    parser = argparse.ArgumentParser(description="Process ultrasound data.")
+    parser.add_argument("-c", "--config", type=str, default=None, help="path to config file.")
+    parser.add_argument(
+        "-t",
+        "--task",
+        default="view",
+        choices=["view"],
+        type=str,
+        help="which task to run",
+    )
+    parser.add_argument(
+        "--backend",
+        default=None,
+        type=str,
+        help=(
+            "Keras backend to use. Default is the one set by the environment "
+            "variable KERAS_BACKEND."
+        ),
+    )
+    parser.add_argument(
+        "--skip_validate_file",
+        default=False,
+        action="store_true",
+        help="Skip zea file integrity checks. Use with caution.",
+    )
+    parser.add_argument("--gui", default=False, action=argparse.BooleanOptionalAction)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    """main entrypoint for zea"""
+    args = get_args()
+
+    set_mpl_style()
+
+    if args.backend:
+        from zea.internal.setup_zea import set_backend
+
+        set_backend(args.backend)
+
+    wd = Path(__file__).parent.resolve()
+    sys.path.append(str(wd))
+
+    import keras
+
+    from zea.interface import Interface
+    from zea.internal.setup_zea import setup
+
+    config = setup(args.config)
+
+    if args.task == "view":
+        cli = Interface(
+            config,
+            validate_file=not args.skip_validate_file,
+        )
+
+        log.info(f"Using {keras.backend.backend()} backend")
+        cli.run(plot=True)
+    else:
+        raise ValueError(f"Unknown task {args.task}, see `zea --help` for available tasks.")
+
+
+if __name__ == "__main__":
+    main()

zea/agent/__init__.py

@@ -0,0 +1,27 @@
+"""Agent subpackage for closing action-perception loop in ultrasound imaging.
+
+The `agent` subpackage provides tools and utilities for agent-based algorithms within the ``zea`` framework, including mask generation and action selection strategies. See :mod:`zea.agent.masks` and :mod:`zea.agent.selection` for key functions implementing intelligent focused transmit selection, such as the :class:`zea.agent.selection.GreedyEntropy` algorithm.
+
+For a practical example, see :doc:`../notebooks/agent/agent_example`.
+
+Example usage
+^^^^^^^^^^^^^
+
+.. code-block:: python
+
+    import zea
+    import numpy as np
+
+    agent = zea.agent.selection.GreedyEntropy(
+        n_actions=7,
+        n_possible_actions=112,
+        img_width=112,
+        img_height=112,
+    )
+
+    # (batch, samples, height, width)
+    particles = np.random.rand(1, 10, 112, 112)
+    lines, mask = agent.sample(particles)
+"""
+
+from . import masks, selection

zea/agent/gumbel.py

@@ -0,0 +1,107 @@
+"""Gumbel-Softmax trick implemented with the multi-backend ``keras.ops``."""
+
+import keras
+import numpy as np
+from keras import ops
+
+if keras.backend.backend() != "jax":
+    # This allows tensorflow tracing
+    prod = ops.prod
+else:
+    # Jax does not allow shapes to be tensors
+    prod = np.prod
+
+
+class SubsetOperator:
+    """SubsetOperator applies the Gumbel-Softmax trick for continuous top-k selection.
+
+    Args:
+        k (int): The number of elements to select.
+        tau (float, optional): The temperature parameter for Gumbel-Softmax. Defaults to 1.0.
+        hard (bool, optional): Whether to use straight-through Gumbel-Softmax. Defaults to False.
+
+    Sources:
+        - `Reparameterizable Subset Sampling via Continuous Relaxations <https://github.com/ermongroup/subsets>`_
+        - `Sampling Subsets with Gumbel-Top Relaxations <https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/DL2/sampling/subsets.html>`_
+    """  # noqa: E501
+
+    def __init__(self, k, tau=1.0, hard=False, n_value_dims=1):
+        self.k = k
+        self.tau = tau
+        self.hard = hard
+        self.EPSILON = np.finfo(np.float32).tiny
+        self.n_value_dims = n_value_dims  # for a image mask: n_value_dims=2
+
+    def gumbel_sample(self, shape):
+        """Samples from Gumbel(0,1) distribution"""
+        uniform = keras.random.uniform(shape, minval=0, maxval=1)
+        return -ops.log(-ops.log(uniform + self.EPSILON) + self.EPSILON)
+
+    def __call__(self, scores):
+        # Gumbel-Softmax trick to make the sampling differentiable
+        gumbel_noise = self.gumbel_sample(ops.shape(scores))
+        scores = scores + gumbel_noise
+
+        # Continuous top-k selection
+        khot = ops.zeros_like(scores)
+        onehot_approx = ops.zeros_like(scores)
+
+        for _ in range(self.k):
+            khot_mask = ops.max(1.0 - onehot_approx, self.EPSILON)
+            scores = scores + ops.log(khot_mask)
+            onehot_approx = ops.softmax(scores / self.tau, axis=1)
+            khot = khot + onehot_approx
+
+        # Optionally convert soft selection to hard selection using straight-through estimator
+        if self.hard:
+            res = hard_straight_through(khot, self.k, self.n_value_dims)
+        else:
+            res = khot
+
+        return res
+
+
+def hard_straight_through(khot_orig, k, n_value_dims=1):
+    """Applies the hard straight-through estimator to the given k-hot encoded tensor.
+
+    Args:
+        khot_orig (Tensor): The original k-hot encoded tensor.
+        k (int): The number of top elements to select.
+        n_value_dims (int, optional): The number of value dimensions in the input tensor.
+            Defaults to 1. E.g. for a 2D image mask, `n_value_dims=2`.
+    Returns:
+        Tensor: The tensor after applying the hard straight-through estimator,
+            with the same shape as `khot_orig`.
+    """
+
+    # Extract the batch dimensions and the value dimensions
+    original_shape = ops.shape(khot_orig)
+    value_dims = original_shape[-n_value_dims:]
+
+    # Flatten the input tensor along the value dimensions
+    khot = ops.reshape(khot_orig, (-1, prod(value_dims)))
+
+    # Get the top-k indices
+    indices = ops.top_k(khot, k)[1]
+
+    # Reshape the indices for use with ops.scatter
+    scatter_indices = ops.stack(
+        [
+            ops.repeat(ops.arange(ops.shape(khot)[0]), k),
+            ops.reshape(indices, (-1,)),
+        ],
+        axis=-1,
+    )
+
+    # Create the hard k-hot tensor
+    khot_hard = ops.scatter(
+        scatter_indices,
+        ops.ones(prod(ops.shape(indices)), "float32"),
+        ops.shape(khot),
+    )
+
+    # Straight-through estimator
+    out = khot_hard - ops.stop_gradient(khot) + khot
+
+    # Reshape to the original shape
+    return ops.reshape(out, original_shape)

zea/agent/masks.py

@@ -0,0 +1,192 @@
+"""
+Mask generation utilities.
+
+These masks are used as a measurement operator for focused scan-line subsampling.
+"""
+
+from typing import List
+
+import keras
+from keras import ops
+
+from zea import tensor_ops
+from zea.agent.gumbel import hard_straight_through
+
+_DEFAULT_DTYPE = "bool"
+
+
+def indices_to_k_hot(
+    indices: List[int],
+    n_possible_actions: int,
+    dtype=_DEFAULT_DTYPE,
+):
+    """Convert a list of indices to a k-hot encoded vector.
+
+    A k-hot encoded vector is suitable during tracing when the number of actions can change.
+    This is the default represenation for actions in zea.
+
+    Args:
+        indices (List[int]): List of indices to set to 1.
+        n_possible_actions (int): Total number of possible actions.
+        dtype (str, optional): Data type of the mask. Defaults to _DEFAULT_DTYPE.
+
+    Returns:
+        Tensor: k-hot-encoded vector of shape (n_possible_actions).
+    """
+    mask = ops.zeros(n_possible_actions, dtype=dtype)
+    return ops.scatter_update(
+        mask, ops.expand_dims(indices, axis=1), ops.ones(len(indices), dtype=dtype)
+    )
+
+
+def k_hot_to_indices(selected_lines, n_actions: int, fill_value=-1):
+    """Convert k-hot encoded lines to indices of selected actions.
+
+    Args:
+        selected_lines (Tensor): k-hot encoded lines of shape (batch_size, n_possible_actions).
+        n_actions (int): Number of lines selected.
+        fill_value (int, optional): Value to fill in case there are not enough selected actions.
+            Defaults to -1.
+
+    Returns:
+        Tensor: Indices of selected actions of shape (batch_size, n_actions).
+            If there are fewer than `n_actions` selected, the remaining indices will be
+            filled with `fill_value`.
+    """
+
+    # Find nonzero indices for each frame
+    def get_nonzero(row):
+        return tensor_ops.nonzero(row > 0, size=n_actions, fill_value=fill_value)[0]
+
+    indices = ops.vectorized_map(get_nonzero, selected_lines)
+    return indices
+
+
+def random_uniform_lines(
+    n_actions: int,
+    n_possible_actions: int,
+    n_masks: int,
+    seed: int | keras.random.SeedGenerator = None,
+    dtype=_DEFAULT_DTYPE,
+):
+    """Will generate a mask with random lines.
+
+    Guarantees precisely n_actions.
+
+    Args:
+        n_actions (int): Number of actions to be selected.
+        n_possible_actions (int): Number of possible actions.
+        n_masks (int): Number of masks to generate.
+        seed (int | SeedGenerator | jax.random.key, optional): Seed for random number generation.
+            Defaults to None.
+
+    Returns:
+        Tensor: k-hot-encoded line vectors of shape (n_masks, n_possible_actions).
+                Needs to be converted to image size.
+    """
+    masks = keras.random.uniform([n_masks, n_possible_actions], seed=seed, dtype="float32")
+    masks = hard_straight_through(masks, n_actions)
+    return ops.cast(masks, dtype=dtype)
+
+
+def _assert_equal_spacing(n_actions, n_possible_actions):
+    assert n_possible_actions % n_actions == 0, (
+        "Number of actions must divide evenly into possible actions to use equispaced sampling."
+    )
+
+
+def initial_equispaced_lines(
+    n_actions, n_possible_actions, dtype=_DEFAULT_DTYPE, assert_equal_spacing=True
+):
+    """Generate an initial equispaced k-hot line mask.
+
+    For example, if ``n_actions=2`` and ``n_possible_actions=6``,
+    then ``initial_mask=[1, 0, 0, 1, 0, 0]``.
+
+    Args:
+        n_actions (int): Number of actions to be selected.
+        n_possible_actions (int): Number of possible actions.
+        dtype (str, optional): Data type of the mask. Defaults to _DEFAULT_DTYPE.
+        assert_equal_spacing (bool, optional): If True, asserts that
+            `n_possible_actions` is divisible by `n_actions`, this means that every
+            line will have the exact same spacing. Otherwise, there might be
+            some spacing differences. Defaults to True.
+
+    Returns:
+        Tensor: k-hot-encoded line vector of shape (n_possible_actions).
+            Needs to be converted to image size.
+    """
+    if assert_equal_spacing:
+        _assert_equal_spacing(n_actions, n_possible_actions)
+        selected_indices = ops.arange(0, n_possible_actions, n_possible_actions // n_actions)
+    else:
+        selected_indices = ops.linspace(0, n_possible_actions - 1, n_actions, dtype="int32")
+
+    return indices_to_k_hot(selected_indices, n_possible_actions, dtype=dtype)
+
+
+def next_equispaced_lines(previous_lines, shift=1):
+    """
+    Rolls the previous equispaced mask of shape (..., n_possible_actions) to the right by
+    `shift` which is 1 by default.
+    """
+    return ops.roll(previous_lines, shift=shift, axis=-1)
+
+
+def lines_to_im_size(lines, img_size: tuple):
+    """
+    Convert k-hot-encoded line vectors to image size.
+
+    Args:
+        lines (Tensor): shape is (n_masks, n_possible_actions)
+        img_size (tuple): (height, width)
+
+    Returns:
+        Tensor: Masks of shape (n_masks, img_size, img_size)
+    """
+    height, width = img_size
+
+    remainder = width % ops.shape(lines)[1]
+    assert remainder == 0, (
+        f"Width must be divisible by number of actions. Got remainder: {remainder}."
+    )
+
+    # Repeat till width of image
+    masks = ops.repeat(lines, width // ops.shape(lines)[1], axis=1)
+
+    # Repeat till height of image
+    masks = ops.repeat(masks[:, None], height, axis=1)
+
+    return masks
+
+
+def make_line_mask(
+    line_indices: List[int],
+    image_shape: List[int],
+    line_width: int = 1,
+    dtype=_DEFAULT_DTYPE,
+):
+    """
+    Creates a mask with vertical (i.e. second axis) lines at specified indices.
+
+    Args:
+        line_indices (List[int]): A list of indices where the lines should be drawn.
+        image_shape (List[int]): The shape of the image as [height, width, channels].
+        line_width (int, optional): The width of each line. Defaults to 1.
+        dtype (str, optional): The data type of the mask. Defaults to "float32".
+
+    Returns:
+        mask (Tensor): A tensor of the same shape as `image_shape` with lines drawn
+            at the specified indices.
+    """
+    height, width, channels = image_shape
+
+    # Create k-hot vector for the line indices
+    k_hot = indices_to_k_hot(line_indices, width // line_width, dtype=dtype)
+    # Expand to (1, n_possible_actions) for lines_to_im_size
+    k_hot = ops.expand_dims(k_hot, axis=0)
+    # Use lines_to_im_size to create the mask of shape (1, height, width)
+    mask_2d = lines_to_im_size(k_hot, (height, width))[0]
+
+    # Expand to (height, width, channels)
+    return ops.broadcast_to(mask_2d[..., None], (height, width, channels))