Adaptive Beamforming by Deep Learning (ABLE)

Delay-and-sum adds up the time-of-flight corrected channel signals with fixed apodization weights. Adaptive beamformers instead derive the weights per pixel from the data itself, which suppresses off-axis clutter and sharpens the image. The classic example is minimum variance beamforming, and it is expensive: it estimates and inverts a covariance matrix for every single pixel.

ABLE (Adaptive Beamforming by deep LEarning) keeps the adaptive weights but replaces that estimator with a small neural network, which predicts the weights directly from the delayed channel data. Trained to imitate an adaptive beamformer, it approaches its image quality at close to delay-and-sum cost.

In this notebook we beamform a plane wave acquisition three ways: delay-and-sum, minimum variance, and an ABLE model that we train — inside the pipeline — to mimic minimum variance.

Scope of this example: This notebook demonstrates the ABLE approach using a small, publicly available PICMUS example, with training on a single phantom frame. It is not an exact reproduction of the model, training setup, or results reported in the paper.

Reference: Luijten et al., “Adaptive Ultrasound Beamforming Using Deep Learning”, IEEE Transactions on Medical Imaging 39 (12), 2020.

   

‼️ Important: This notebook is optimized for GPU/TPU. Code execution on a CPU may be very slow.

If you are running in Colab, please enable a hardware accelerator via:

Runtime → Change runtime type → Hardware accelerator → GPU/TPU 🚀.

%%capture
%pip install zea

import os

os.environ["KERAS_BACKEND"] = "jax"

import keras
import matplotlib.pyplot as plt
import numpy as np

import zea
from zea import init_device
from zea.display import histogram_match, to_8bit
from zea.models.able import ABLE
from zea.ops import Beamform, EnvelopeDetect, LogCompress, Normalize, Pipeline
from zea.utils import FunctionTimer
from zea.visualize import set_mpl_style

zea: Using backend 'jax'

init_device(verbose=False)
set_mpl_style()

seed = 0
n_steps = 1200
learning_rate = 1e-2
num_patches = 100
n_transmits = None  # None compounds every plane wave; an int takes that many, evenly spaced
# Training differentiates through every transmit at once, so its memory grows with
# transmits x pixels -- at full resolution that is hundreds of GB. The figures are made
# at full resolution; the fit runs on a coarser grid, which costs ABLE nothing because
# it is a stack of 1x1 convolutions and so is not tied to the grid it was fitted on.
# Fewer, larger patches for the fit: the gradient keeps every patch's activations alive
# either way, so more patches only adds overhead -- and on the smaller training grid 100
# patches would leave too few axial rows per patch for minimum variance to average over.
train_num_patches = 20
train_grid_size_z = 192
train_grid_size_x = 128
grid_size_z = None  # None lets zea size the grid at about half a wavelength
grid_size_x = None

Loading data

We use the resolution phantom from the PICMUS dataset, hosted in zea format on the Hugging Face Hub. The series of point targets in that phantom dataset directly relates to the lateral resolution the adaptive beamformer is able to obtain.

path = (
    "hf://zeahub/picmus/database/experiments/resolution_distorsion/"
    "resolution_distorsion_expe_dataset_iq/resolution_distorsion_expe_dataset_iq.hdf5"
)

with zea.File(path) as f:
    data = f.data.raw_data[0][None]  # first frame, with a batch dimension
    parameters = f.load_parameters()

parameters.set_transmits(n_transmits)  # compound the steered plane waves
data = data[:, parameters.selected_transmits]

parameters.xlims = (-0.019, 0.019)
parameters.zlims = (0.0, 0.06)
parameters.f_number = 0


def apply_grid_size(parameters):
    """Shrink the beamforming grid when the notebook test asks for it."""
    if grid_size_z is not None:
        parameters.grid_size_z = grid_size_z
    if grid_size_x is not None:
        parameters.grid_size_x = grid_size_x


apply_grid_size(parameters)

Beamforming with zea.ops.Beamform

Beamform is a pipeline that time-of-flight corrects the channel data and hands it to a beamformer of choice. Swapping the beamformer is a single keyword, so the same B-mode pipeline serves all three variants. Keywords the beamformer itself takes are forwarded to it, which is how minimum variance gets its sub-aperture length below. That length is the main lever on how adaptive it is: a longer sub-aperture gives the beamformer more degrees of freedom to null off-axis energy, while a shorter one is cheaper and better conditioned but lets its weights drift towards the uniform ones of delay-and-sum, until the two images look alike.

DYNAMIC_RANGE = (-50.0, 0.0)  # 60 dB leaves the speckle washed out; 50 dB has more bite


def bmode_pipeline(beamformer, jit_options="ops", patches=None, **kwargs):
    """B-mode pipeline: beamform, then envelope detect, normalize and log compress."""
    return Pipeline(
        operations=[
            Beamform(
                beamformer=beamformer,
                num_patches=num_patches if patches is None else patches,
                **kwargs,
            ),
            EnvelopeDetect(),
            Normalize(),
            LogCompress(),
        ],
        with_batch_dim=True,
        jit_options=jit_options,
    )

The pipeline is stateless: all acquisition parameters travel along with the data. We prepare them once and reuse them for every pipeline below.

n_el = data.shape[-2]

das_pipeline = bmode_pipeline("delay_and_sum")
# The longest sub-aperture the covariance still supports. Shortening it pulls minimum
# variance back towards delay-and-sum until the two are hard to tell apart.
mv_pipeline = bmode_pipeline("minimum_variance", subarray_size=n_el // 2)

inputs = das_pipeline.prepare_parameters(parameters)
inputs["dynamic_range"] = DYNAMIC_RANGE

das_bmode = das_pipeline(**inputs, data=data)["data"]
mv_bmode = mv_pipeline(**inputs, data=data)["data"]

Minimum variance tightens the point targets – on this phantom from roughly 0.8 mm wide at -6 dB to a little over 0.5 mm – at the cost of a much heavier computation. That is the image we will ask ABLE to reproduce.

def image_extent(parameters):
    """Grid limits in mm, in the order imshow wants them."""
    return [
        parameters.xlims[0] * 1e3,
        parameters.xlims[1] * 1e3,
        parameters.zlims[1] * 1e3,
        parameters.zlims[0] * 1e3,
    ]


extent_mm = image_extent(parameters)


def plot_bmodes(images, titles, extent=None):
    """Plot a row of B-mode images on a shared dynamic range.

    The panels are kept modest in size on purpose: speckle is noise-like and so compresses
    poorly, and the docs checker caps a single embedded output at 400 kB.
    """
    extent = extent_mm if extent is None else extent
    fig, axes = plt.subplots(
        1, len(images), figsize=(2.55 * len(images), 3.25), constrained_layout=True
    )
    for ax, image, title in zip(np.atleast_1d(axes), images, titles):
        ax.imshow(to_8bit(np.array(image), DYNAMIC_RANGE), cmap="gray", extent=extent)
        ax.set_title(title)
        ax.set_xlabel("Lateral (mm)")
    np.atleast_1d(axes)[0].set_ylabel("Depth (mm)")
    return fig


_ = plot_bmodes([das_bmode[0], mv_bmode[0]], ["Delay-and-sum", "Minimum variance"])

The ABLE model

ABLE looks at the time-of-flight corrected samples of one pixel across the receive elements and predicts an apodization weight for each of them. Because every pixel is treated on its own, the network is a stack of 1×1 convolutions, kept small on purpose: a few thousand parameters is enough.

The model has to be built before it enters the pipeline, so that its weights are created eagerly rather than inside a traced pipeline call.

n_el, n_ch = data.shape[-2], data.shape[-1]

keras.utils.set_random_seed(seed)  # same starting point on every run
able = ABLE(latent_dim=32, n_latent_layers=2)
able.build((1, 1, n_el, n_ch))  # (n_tx, n_pix, n_el, n_ch)

able.summary()

Model: "able"

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape           ┃       Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ conv2d (Conv2D)                 │ (1, 1, 1, 256)         │        65,792 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ conv2d_1 (Conv2D)               │ (1, 1, 1, 32)          │        16,416 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ conv2d_2 (Conv2D)               │ (1, 1, 1, 32)          │         2,080 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ conv2d_3 (Conv2D)               │ (1, 1, 1, 256)         │        16,640 │
└─────────────────────────────────┴────────────────────────┴───────────────┘

 Total params: 100,928 (394.25 KB)

 Trainable params: 100,928 (394.25 KB)

 Non-trainable params: 0 (0.00 B)

ABLE is registered as a beamformer, so it slots into the very same pipeline. The one difference is jit_options=None: compiling the operations separately would cut the gradient path to the model weights, and training would silently stall.

able_pipeline = bmode_pipeline("able", model=able, jit_options=None, patches=train_num_patches)

train_parameters = parameters.copy()
train_parameters.grid_size_z = train_grid_size_z
train_parameters.grid_size_x = train_grid_size_x

train_inputs = able_pipeline.prepare_parameters(train_parameters)
train_inputs["dynamic_range"] = DYNAMIC_RANGE

# The regression target has to live on the same grid as the images ABLE produces, and on
# the same patch layout: minimum variance averages its covariance over axially adjacent
# pixels, and a patch too short in depth truncates that window.
mv_train = bmode_pipeline("minimum_variance", subarray_size=n_el // 2, patches=train_num_patches)(
    **train_inputs, data=data
)["data"]

zea: WARNING width/grid_size_x = 0.0002969 > wavelength/2 = 0.0001478. Consider increasing grid_size_x to 258 or more, or unsetting it to size the grid automatically.
zea: WARNING depth/grid_size_z = 0.0003125 > wavelength/2 = 0.0001478. Consider increasing grid_size_z to 406 or more, or unsetting it to size the grid automatically.

Training

ABLE is trained by regression: the whole pipeline is differentiable, so we can push the log-compressed image through it and compare it against the minimum variance image with a mean squared error. Gradients flow all the way back through the delay-and-sum and the envelope detection into the network weights.

Note that we are training on a single frame here — enough to make the point in a notebook. See the paper for the training regime used in practice.

class ABLEBeamformer(keras.Model):
    """Thin Keras model around the pipeline, exposing the ABLE weights to the optimizer."""

    def __init__(self, model, pipeline, parameters):
        super().__init__()
        self.able = model
        self.pipeline = pipeline
        self.parameters = parameters

    def call(self, data):
        return self.pipeline(**self.parameters, data=data)[self.pipeline.output_key]


trainer = ABLEBeamformer(able, able_pipeline, train_inputs)
trainer.compile(optimizer=keras.optimizers.Adam(learning_rate=learning_rate), loss="mse")

losses = []
progbar = keras.utils.Progbar(n_steps, stateful_metrics=["loss"])

print("Compiling the training step (this takes a while), then training...")

for step in range(1, n_steps + 1):
    loss = float(trainer.train_on_batch(data, mv_train))
    losses.append(loss)
    progbar.update(step, values=[("loss", loss)])

Compiling the training step (this takes a while), then training...
1200/1200 ━━━━━━━━━━━━━━━━━━━━ 508s 398ms/step - loss: 1.5106

The loss should drop by more than an order of magnitude and then flatten out, which is what tells us the run has converged.

fig, ax = plt.subplots(figsize=(6, 3.5), constrained_layout=True)
ax.plot(np.arange(1, len(losses) + 1), losses)
ax.set_yscale("log")
ax.set_xlabel("Training step")
ax.set_ylabel("Loss (MSE)")
_ = ax.set_title("ABLE training loss")

Results

Training is done, so we can hand the model to a compiled pipeline again and time all three beamformers. The first call of each is dropped, since it includes tracing and compilation.

The three images are histogram matched to the delay-and-sum one before display, following Bottenus et al.: adaptive beamformers suppress off-axis energy, which shifts the brightness of the whole image, and without matching that shift reads as a difference in image quality.

trained_pipeline = bmode_pipeline("able", model=able)
able_bmode = trained_pipeline(**inputs, data=data)["data"]

timer = FunctionTimer()
timed = {
    name: timer(lambda p=pipeline: p(**inputs, data=data), name=name)
    for name, pipeline in [
        ("Delay-and-sum", das_pipeline),
        ("Minimum variance", mv_pipeline),
        ("ABLE", trained_pipeline),
    ]
}

# Minimum variance dominates this cell: it inverts a covariance per pixel
for name, run in timed.items():
    print(f"timing {name}...")
    for _ in range(6):  # one warmup + five timed runs
        run()

timer.print(drop_first=True)

timing Delay-and-sum...
timing Minimum variance...
timing ABLE...
Function Timing Statistics
=====================================================================================================
Function              Mean          Median        Std Dev       Min           Max           Count
-----------------------------------------------------------------------------------------------------
Delay-and-sum         0.171504      0.171614      0.000723      0.170585      0.172428      5
Minimum variance      48.839354     48.815990     0.060196      48.766020     48.908843     5
ABLE                  0.879694      0.916365      0.051758      0.822366      0.918061      5

Each beamformer compresses the dynamic range differently, which skews a side-by-side look. Matching all three to the delay-and-sum histogram puts them on a common scale, so what is left to see is the image and not the contrast.

reference = np.array(das_bmode[0])
matched = [
    histogram_match(np.array(image), reference)
    for image in (das_bmode[0], mv_bmode[0], able_bmode[0])
]

titles = [
    f"{name}\n{timer.get_stats(name, drop_first=True)['mean'] * 1e3:.0f} ms" for name in timed
]

_ = plot_bmodes(matched, titles)

Does it transfer?

The network was fitted to a single phantom frame, so the fair question is whether it learned anything beyond that one image. PICMUS also ships in-vivo scans taken with the same probe and the same plane wave sequence, so we can insert the trained model into the very same pipeline. We don’t retrain it.

carotid_path = (
    "hf://zeahub/picmus/in_vivo/carotid_cross/"
    "carotid_cross_expe_dataset_iq/carotid_cross_expe_dataset_iq.hdf5"
)

with zea.File(carotid_path) as f:
    carotid_data = f.data.raw_data[0][None]
    carotid_parameters = f.load_parameters()

carotid_parameters.set_transmits(n_transmits)
carotid_data = carotid_data[:, carotid_parameters.selected_transmits]
carotid_parameters.xlims = (-0.019, 0.019)
carotid_parameters.zlims = (0.0, 0.045)
apply_grid_size(carotid_parameters)
carotid_parameters.f_number = 0

carotid_inputs = das_pipeline.prepare_parameters(carotid_parameters)
carotid_inputs["dynamic_range"] = DYNAMIC_RANGE

carotid_bmodes = [
    pipeline(**carotid_inputs, data=carotid_data)["data"][0]
    for pipeline in (das_pipeline, mv_pipeline, trained_pipeline)
]


def mse(image_a, image_b):
    """Mean squared difference between two log-compressed images."""
    return float(np.mean((np.array(image_a) - np.array(image_b)) ** 2))


for label, (das_i, mv_i, able_i) in [
    ("phantom", (das_bmode[0], mv_bmode[0], able_bmode[0])),
    ("carotid", carotid_bmodes),
]:
    ratio = mse(das_i, mv_i) / mse(able_i, mv_i)
    print(f"{label}: ABLE sits {ratio:4.1f}x closer to MV than DAS does (MSE ratio)")

carotid_reference = np.array(carotid_bmodes[0])
_ = plot_bmodes(
    [histogram_match(np.array(image), carotid_reference) for image in carotid_bmodes],
    ["Delay-and-sum", "Minimum variance", "ABLE"],
    extent=image_extent(carotid_parameters),
)

phantom: ABLE sits  5.9x closer to MV than DAS does (MSE ratio)
carotid: ABLE sits  1.8x closer to MV than DAS does (MSE ratio)

ABLE recovers most of what minimum variance offers over delay-and-sum, at a fraction of its runtime. The covariance inversion is gone, replaced by a handful of 1×1 convolutions. As you can see the difference on in-vivo data is not as pronounced as on the phantom. Although looking at some smaller features (which need the improved lateral resolution) shows that ABLE is still able to recover some of the minimum variance image quality.

A few things to try from here:

• Tune the minimum variance beamformer by varying subarray_size, diagonal_loading, and axial_averaging, and inspect how these settings affect the phantom and in-vivo images. These settings also determine ABLE’s training target: after changing them, regenerate the MV targets and retrain ABLE, using the same MV settings for training and comparison.

• Train against another beamformer by changing a single word: "coherence_factor", "generalized_coherence_factor" or "delay_multiply_and_sum" are all registered too.

• Train on more frames, or on a whole dataset with a zea.Dataloader, which is what makes the model generalize beyond the frame shown here.

• Save the trained weights with able.save_to_preset(...) and reload them later with ABLE.from_preset(...).