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)

Adding a new model¶

New models are welcome! Please follow the Contributing guide for the general workflow (forking, branches, pull requests, etc.). The steps below walk you through what is specific to adding a model.

Create a new file in zea/models/ for your model, e.g. zea/models/mymodel.py.
Add a model class that inherits from zea.models.base.BaseModel. For generative models, use GenerativeModel or DeepGenerativeModel as the base class. Implement the call method.
Upload the pretrained weights to our Hugging Face. The expected files are a config.json and a model.weights.h5. See the Keras documentation for how to save these. You can drag and drop the files directly on the Hugging Face website.

Tip

Alternate saving methods are also possible. See zea.models.echonet.EchoNet for an example — in that case you need to implement a custom_load_weights method in your model class.
Add a preset for your model in zea.models.presets. Presets let you register multiple sets of weights for the same model architecture.
In your model file, import the presets module and call register_presets with your model class to activate the presets.
Import your model in zea/models/__init__.py to make it part of the package.

Adding non-Keras (custom) models¶

The recommended approach for any model is to implement it as a native Keras 3 model. This gives you backend-agnostic execution (JAX, TensorFlow, PyTorch) and the full preset/weight-loading infrastructure for free.

For models originally trained in PyTorch, the typical workflow is:

Vendor the architecture — keep the PyTorch network code at hand while you port it. If you want it in the module long-term (e.g. for ONNX export), put it in a _build_torch_classes() helper that imports torch lazily. PyTorch is then required for weight conversion and ONNX export, but never for inference — which is what keeps it out of the runtime dependencies.
Implement the Keras architecture — write keras.layers.Layer subclasses that replicate each block. Key API differences to handle:
- Padding for strided Conv2D: Keras padding='same' is asymmetric for stride > 1; use ZeroPadding2D(p) + Conv2D(padding='valid') to match PyTorch’s symmetric padding=p.
- ConvTranspose: Keras Conv2DTranspose(padding='valid') gives the full output, from which PyTorch’s padding=p crops p on every side — so follow it with Cropping2D(p). A PyTorch output_padding adds to the end only, so padding=1, output_padding=1 becomes x[:, 1:, 1:, :].
- Normalization: BatchNormalization needs epsilon=1e-5 to match PyTorch’s BatchNorm2d default, and takes its weights as [weight, bias, running_mean, running_var]. For InstanceNorm2d(affine=False) use GroupNormalization(groups=C, scale=False, center=False, epsilon=1e-5).
- Weight axes: Conv2D — (2,3,1,0); Conv2DTranspose — (2,3,1,0) (same permutation, different semantics).
- Input format: Keras defaults to channels-last (NHWC); transpose NCHW → NHWC in call() and back before returning.
- Layer names: do not name custom layer classes with a leading underscore. Keras derives the layer name from the class name, and the TensorFlow backend rejects names starting with _.
Write a weight-loading helper — a function that maps PyTorch state-dict keys to the Keras layer tree and calls layer.set_weights([...]). The weights only exist once the model is built; building it symbolically with self(keras.Input(batch_shape=...)) avoids pushing real data through it.
Add from_pth(path) classmethod — wraps the weight loader for convenient local testing.
Verify against the reference implementation — before uploading the converted weights, run both models on the same input and compare. A correct port agrees to ~1e-5 in float32. If it does not, compare block by block: the layer where the difference appears is the one that was ported wrong.

Warning

Run that comparison on CPU. On a supported NVIDIA GPU all three backends — JAX, TensorFlow and PyTorch — use TF32 for convolutions by default, which costs ~3 decimal digits and makes a correct port look broken (differences of ~1e-2 instead of ~1e-5). This applies to any pair of them, so a TensorFlow-versus-PyTorch comparison is no safer than a JAX-versus-PyTorch one.
Optionally add an ONNX fallback — for environments that have onnxruntime but not torch, you can keep a from_onnx(path) classmethod and an _onnx_sess attribute; override call() to dispatch to the ONNX path when the session is set.
Follow steps 3-6 from the standard guide above for HF upload, presets, and registration.

Two worked examples of this pattern: zea.models.inversionnet (native Keras plus a PyTorch state-dict converter) and zea.models.speckle2self (the same, plus a vendored PyTorch architecture and an optional ONNX fallback).