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:

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:

Available 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.

  1. Create a new file in zea/models/ for your model, e.g. zea/models/mymodel.py.

  2. 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.

  3. 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.

  4. Add a preset for your model in zea.models.presets. Presets let you register multiple sets of weights for the same model architecture.

  5. In your model file, import the presets module and call register_presets with your model class to activate the presets.

  6. 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:

  1. Vendor the architecture — copy the PyTorch network code into your module (e.g. inside a _build_torch_classes() helper that imports torch lazily so that PyTorch is only required for weight conversion, not inference).

  2. Implement the Keras architecture — write keras.layers.Layer subclasses that replicate each block. Key API differences to handle:

    • Padding for stride-2 Conv2D: Keras padding='same' is asymmetric for stride > 1; use ZeroPadding2D(1) + Conv2D(padding='valid') to match PyTorch’s symmetric padding=1.

    • ConvTranspose: Keras Conv2DTranspose(padding='valid') gives the full output; crop x[:, 1:, 1:, :] (NHWC) to reproduce PyTorch’s ConvTranspose2d(padding=1, output_padding=1) alignment.

    • InstanceNorm: use GroupNormalization(groups=C, scale=False, center=False, epsilon=1e-5) for InstanceNorm2d(affine=False).

    • 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.

  3. Write a weight-loading helper — a function that maps PyTorch state-dict keys to the Keras layer tree and calls layer.set_weights([...]).

  4. Add from_pth(path) classmethod — wraps the weight loader for convenient local testing.

  5. 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.

  6. Follow steps 3-6 from the standard guide above for HF upload, presets, and registration.

See zea.models.speckle2self for a complete worked example of this pattern (native Keras + PyTorch weight loading + optional ONNX fallback).