zea: A Toolbox for Cognitive Ultrasound Imaging

Technical Summary: zea: A Toolbox for Cognitive Ultrasound Imaging

Problem Statement
Ultrasound imaging is a ubiquitous, non-invasive, and cost-effective medical modality, yet it faces persistent challenges regarding image quality and interpretation. Traditional processing pipelines rely on naive assumptions during signal processing steps—such as filtering, beamforming, and reconstruction—that fail to account for the complex nature of ultrasound signals. Furthermore, current workflows typically treat acquisition (action) and reconstruction (perception) as disjoint processes. The emerging paradigm of cognitive ultrasound imaging seeks to close this action-perception loop by dynamically adjusting acquisition parameters based on real-time observations. However, realizing this paradigm requires handling high-dimensional data and solving complex inverse problems using powerful deep generative models. Existing software solutions (e.g., Field II, k-Wave, USTB, MUST) are well-suited for simulation and standard reconstruction but lack the flexibility, differentiability, and deep learning integration necessary for cognitive, closed-loop workflows.

Methodology and Architecture
To address these limitations, the authors present zea, a Python package designed as a flexible, modular, and differentiable pipeline for cognitive ultrasound imaging. The toolbox is built upon Keras 3, enabling compatibility with TensorFlow, PyTorch, and JAX backends, which facilitates the seamless integration of custom deep learning models into the reconstruction pipeline.

The architecture of zea is organized into four core components:

1. Data Handling: The package utilizes HDF5 files to store raw data and acquisition parameters together. It provides classes (zea.File, zea.Dataset, zea.Dataloader) optimized for machine learning workflows, including batching, shuffling, caching, and preprocessing. It supports conversion scripts for standard datasets like CAMUS, PICMUS, and EchoNet.
2. Pipeline: A stateless, modular class (zea.Pipeline) allows users to define sequences of operations (zea.Operation). Crucially, the pipeline supports Just-in-Time (JIT) compilation, enabling dynamic parameter adjustment and the real-time integration of deep learning models within the reconstruction loop.
3. Models: zea includes a collection of pre-defined models for image and signal processing, accessible locally or via the Hugging Face Hub. Beyond supervised models, it offers (deep) generative models designed to solve inverse problems within a probabilistic machine learning framework.
4. Agents: This module contains tools to interact with the pipeline and models, allowing for the alteration of pipeline parameters or the selection of acquired data subsets. The agent module is specifically designed to close the action-perception loop, tying together the acquisition and reconstruction phases.

Key Contributions

• A Differentiable Cognitive Pipeline: Unlike traditional toolboxes, zea provides a fully differentiable pipeline that supports dynamic parameter adjustment, a prerequisite for cognitive ultrasound where acquisition strategies must adapt to current observations.
• Integration of Generative Models: The toolbox explicitly supports deep generative models to learn the structured distribution of ultrasound signals, addressing the high-dimensional nature of the data and the challenges of inverse problems.
• Modularity and Interoperability: By leveraging Keras 3, the toolbox offers a high-level interface that allows researchers to integrate custom pipelines into broader machine learning workflows while maintaining compatibility with major deep learning frameworks.
• Open-Source Ecosystem: The package is open-source (Apache License 2.0), available via PyPI and GitHub, and includes 20 example notebooks covering various components. It is supported by a continuous integration pipeline (GitHub Actions) for testing, linting, and documentation.

Results and Validation
The paper does not present a single novel experimental result or a specific benchmark comparison against other tools in this text. Instead, the validation of the toolbox is demonstrated through its adoption and contribution by the research community. The authors note that zea has already been utilized and contributed to in several recent works (e.g., Luijten et al., 2020; Schaft et al., 2025; Stevens et al., 2024; Nolan et al., 2025; Federici et al., 2025). These cited works utilize zea for tasks ranging from adaptive beamforming and active inference for closed-loop transmit beamsteering to dehazing ultrasound using diffusion models and high-volume rate 3D reconstruction.

Significance and Claims
The authors position zea as a complementary tool to existing ultrasound software, specifically targeting the gap in cognitive ultrasound imaging. The paper claims that while the full realization of cognitive ultrasound is an ongoing effort, this toolbox provides the necessary infrastructure to spur further research and development in the field. By offering a modular, extensible library that supports closed-loop workflows and seamless integration with state-of-the-art machine learning models, zea aims to facilitate the redesign of current ultrasound imaging pipelines where parameters are dynamically changed based on past and current observations. The significance lies in enabling the transition from static, disjointed processing to dynamic, closed-loop cognitive systems.