zea.models¶

Collection of (generative) models for ultrasound imaging.

zea contains a collection of models for various tasks, all located in the zea.models package.

See the following dropdown for a list of available models:

Presets for these models can be found in zea.models.presets. Presets are pre-trained weights for the models, which can be used to initialize the models for inference or further training. Each model class has a presets attribute that lists the available presets for that model. We store the presets on Hugging Face Hub, and they are downloaded automatically when loading a model with a preset.

To use these models, you can import them directly from the zea.models module and load the pretrained weights using the from_preset() method. For example:

>>> from zea.models.unet import UNet

>>> model = UNet.from_preset("unet-echonet-inpainter")

You can list all available presets using the presets attribute:

>>> from zea.models.unet import UNet
>>> presets = list(UNet.presets.keys())
>>> print(f"Available built-in zea presets for UNet: {presets}")
Available built-in zea presets for UNet: ['unet-echonet-inpainter']

Generative models¶

In addition to regular models, zea provides generative models for tasks such as image generation, inpainting, and denoising. The key difference is that generative models have sampling methods implemented. There are two base classes: GenerativeModel for classical models (e.g. a Gaussian mixture model) and DeepGenerativeModel for neural-network-based models — the latter also inherits from BaseModel, adding Keras features like weight saving and preset loading. Both expose the following methods:

fit() for training the model on data
sample() for generating new samples from the learned distribution
posterior_sample() for drawing samples from the posterior given measurements
log_density() for computing the log-probability of data under the model

See the following dropdown for a list of available generative models:

An example of how to use the zea.models.diffusion.DiffusionModel is shown below:

>>> from zea.models.diffusion import DiffusionModel

>>> model = DiffusionModel.from_preset("diffusion-echonet-dynamic")
>>> samples = model.sample(n_samples=4)

Modules

`base`	Base model class for all zea Keras models.
`carotid_segmenter`	Carotid segmentation model.
`deeplabv3`	DeepLabV3+ architecture for multi-class segmentation.
`dense`	Dense models and architectures
`diffusion`	Diffusion models for ultrasound image generation and posterior sampling.
`dit`	Diffusion Transformer (DiT) backend.
`echonet`	Echonet-Dynamic segmentation model for cardiac ultrasound segmentation.
`echonetlvh`	EchoNetLVH model for segmentation of PLAX view cardiac ultrasound.
`flow_matching`	Flow matching generative model for ultrasound image generation and posterior sampling.
`generative`	Generative models for zea.
`gmm`	Gaussian Mixture Model (GMM) implementation
`hvae`	Hierarchical Variational Auto-Encoder for image generation, posterior sampling and inference tasks.
`layers`	Layers used in zea.models
`lpips`	LPIPS model for perceptual similarity.
`lv_segmentation`	nnU-Net segmentation model trained on the augmented CAMUS dataset.
`preset_utils`	Mostly from keras_hub.src.models import preset_utils
`presets`	Model presets for zea.models
`regional_quality`	MobileNetv2 based image quality model for myocardial regions in apical views.
`speckle2self`	Speckle2Self self-supervised speckle reduction model for ultrasound images.
`taesd`	Tiny Autoencoder (TAESD) model.
`unet`	UNet models and architectures.
`utils`	Utilities for models