zea.models.echonet

Echonet-Dynamic segmentation model for cardiac ultrasound segmentation.

To try this model, simply load one of the available presets:

>>> from zea.models.echonet import EchoNetDynamic

>>> model = EchoNetDynamic.from_preset("echonet-dynamic")

Important

This is a zea implementation of the model. For the original paper and code, see here.

Ouyang, David, et al. “Video-based AI for beat-to-beat assessment of cardiac function.” Nature 580.7802 (2020): 252-256

See also

A tutorial notebook where this model is used: Left ventricle segmentation.

Note

This model is only currently supported with the TensorFlow or JAX Backend. When using TensorFlow as backend, the model will work out of the box. When using JAX as backend, the model is built using TensorFlow and then converted to JAX. This requires both TensorFlow and JAX to be installed, which can be tricky regarding compatible CUDA versions. One option is to run in our Docker container, which has been tested to work with both backends.

Functions

convert_original_weights([output_dir, ...])	Convert the original PyTorch EchoNet-Dynamic weights to a TensorFlow SavedModel.
download_original_weights([weights_folder])	Download the original weights from the EchoNet Github repository.

Classes

EchoNetDynamic(*args, **kwargs)

EchoNet-Dynamic segmentation model for cardiac ultrasound segmentation.

class zea.models.echonet.EchoNetDynamic(*args, **kwargs)

Bases: BaseModel

EchoNet-Dynamic segmentation model for cardiac ultrasound segmentation.

Preprocessing should normalize the input images with mean and standard deviation.

build(input_shape)

Builds the network.

call(inputs)

Segment the input image.

custom_load_weights(preset, **kwargs)

Load the weights for the segmentation model.

maybe_convert_to_jax()

Converts the network to JAX if backend is JAX.

JAX conversion traces the SavedModel using an example input of shape (1, INFERENCE_SIZE, INFERENCE_SIZE, 3). At runtime, call() may pass (B, INFERENCE_SIZE, INFERENCE_SIZE, 3) after resize/tile preprocessing.

zea.models.echonet.convert_original_weights(output_dir=None, weights_folder=None)

Convert the original PyTorch EchoNet-Dynamic weights to a TensorFlow SavedModel.

This is how the echonet-dynamic preset on the Hugging Face Hub was created; it is kept here for reproducibility and is not needed to use the model.

The conversion goes PyTorch -> ONNX -> TensorFlow and therefore needs a few extra packages that are not part of the zea dependencies:

pip install torch torchvision onnx==1.16.1 onnxruntime==1.18.1 onnx2tf \
    onnx-graphsurgeon onnxsim sne4onnx sng4onnx tf-keras

onnxsim is required despite onnx2tf only warning when it is missing, and is unpinned because 0.4.x has no wheel for recent Python versions.

Note

torch and TensorFlow have to coexist in one process here, which not every combination of wheels survives. Prefer the zeahub/all container, and import torch before this module if you hit a crash while building the torch model.

Parameters:

• output_dir (str | Path, optional) – Folder to write the converted model to. Defaults to a timestamped folder under ./temp/zea.

• weights_folder (str | Path, optional) – Folder to download the original weights into. See download_original_weights().

Returns:

Folder containing the converted TensorFlow SavedModel.

Return type:

Path

zea.models.echonet.download_original_weights(weights_folder=None)

Download the original weights from the EchoNet Github repository.

Only used by convert_original_weights() to reproduce the published preset; using the model needs neither.

Parameters:

weights_folder (str | Path, optional) – Folder to download the weights into. Defaults to ./echonet_weights.

Returns:

Path to the downloaded checkpoint.

Return type:

Path