zea.models.able¶

Adaptive Beamforming by Deep LEarning (ABLE).

Original implementation of paper:

“Adaptive Ultrasound Beamforming Using Deep Learning”
https://doi.org/10.1109/TMI.2020.3008537
Author: Ben Luijten

Classes

`ABLE`(args, *kwargs)	Adaptive Beamforming by Deep LEarning (ABLE) model.
`ABLEBeamform`(args, *kwargs)	Beamform by summing the ABLE-weighted channel data.

class zea.models.able.ABLE(*args, **kwargs)[source]¶

Bases: BaseModel

Adaptive Beamforming by Deep LEarning (ABLE) model.

Implements a configurable pixel-wise convolutional encoder/decoder that computes per-element adaptive weights from time-of-flight-corrected channel data. Summing the weighted data over elements and transmits forms the image; ABLEBeamform does both in one pipeline operation.

Reference

Luijten, B. et al., “Adaptive Ultrasound Beamforming Using Deep Learning,” IEEE Trans. Med. Imaging 39 (12), 2020. https://doi.org/10.1109/TMI.2020.3008537

Expected input shape (time-of-flight corrected channel data, as produced by TOFCorrection):

(n_tx, n_pix, n_el) — RF data
(n_tx, n_pix, n_el, n_ch) — IQ data

The model maps over the transmit axis (n_tx) internally; each (n_pix, n_el[, n_ch]) slice is processed independently.

Note

Only 1×1 convolutions (kernel_size=1) are currently supported because PatchedGrid processes pixels independently.

Parameters:

latent_dim (int) – Channel size for the hidden layers when latent_layers is not supplied. Default is 32.
kernel_size (int, tuple, or list) –
Kernel size specification.
- int — every convolution uses (k, k).
- tuple (h, w) — every convolution uses (h, w).
- list of ints/tuples — per-layer kernel sizes (length must equal the total number of layers).
Default is 1.
n_latent_layers (int) – Number of hidden layers when latent_layers is not supplied. Must be ≥ 1. Default is 2.
latent_layers (list or None) – Explicit list of channel sizes for the hidden layers. Overrides latent_dim when provided, and must contain exactly n_latent_layers entries. Default is None, which gives every hidden layer latent_dim channels.
axis (int or None) – Reserved for future use. Default is None.
name (str) – Model name forwarded to BaseModel. Default is "able".
**kwargs – Additional keyword arguments forwarded to BaseModel.

antirectifier(x)[source]¶

Apply the anti-rectifier activation function.

This function centers the input, splits it into positive and negative components, and normalizes the result.

Parameters:: x (Tensor) – Input tensor.
Returns:: Transformed tensor with anti-rectifier activation applied.
Return type:: Tensor

apply_model(inputs)[source]¶

Apply the ABLE network to a single transmit slice.

Parameters:: inputs (Tensor) – TOF-corrected data for one transmit event, shape (n_pix, n_el) or (n_pix, n_el, n_ch).
Returns:: Adaptively weighted data with the same shape as inputs.
Return type:: Tensor

build(input_shape)[source]¶

Build the ABLE model based on the input shape.

Parameters:

input_shape (tuple) –

Shape of the input tensor. Supported formats (n_tx acts as the batch axis that call maps over):

(n_tx, n_pix, n_el) — RF data, rank 3
(n_tx, n_pix, n_el, n_ch) — IQ data, rank 4

call(inputs)[source]¶

Apply ABLE to the input data.

Maps apply_model over the first axis (n_tx) using keras.ops.map() so the forward pass is fully traceable by JAX and compatible with gradient computation.

Parameters:: inputs (Tensor) – Shape (n_tx, n_pix, n_el[, n_ch]).
Returns:: Adaptively weighted data with the same shape as inputs.
Return type:: Tensor

get_config()[source]¶

Returns the config of the object.

An object config is a Python dictionary (serializable) containing the information needed to re-instantiate it.

stack_channels(x, axis)[source]¶

Reshape input into 4D for use with Conv2D.

PatchedGrid passes per-pixel data with the spatial dimensions collapsed:

Rank 2: (pixels, elements)
    -> (pixels, 1, 1, elements)

Rank 3: (pixels, elements, n_ch)
    -> (pixels, 1, 1, elements * n_ch)

    Elements and ``n_ch`` are merged into a single channel axis, with
    ``n_ch`` varying fastest so that :meth:`unstack_channels` can restore
    the original layout.

unstack_channels(x, meta)[source]¶

Inverse of stack_channels. Restores the original shape from Conv2D output.

Expects x with shape (pixels, 1, 1, C) and restores: - rank 2 input: (pixels, 1, 1, elements) -> (pixels, elements) - rank 3, ch==1: (pixels, 1, 1, elements) -> (pixels, elements, 1) - rank 3, ch>1: (pixels, 1, 1, n_ch*elements) -> (pixels, elements, n_ch)

class zea.models.able.ABLEBeamform(*args, **kwargs)[source]¶

Bases: Operation

Beamform by summing the ABLE-weighted channel data.

Drop-in replacement for DelayAndSum: where delay-and-sum sums the time-of-flight corrected channels with uniform weights, this operation first weights them with the content-adaptive apodization predicted by ABLE.

Being registered as a beamformer, it slots straight into Beamform:

from zea.models.able import ABLE
from zea.ops import Beamform

model = ABLE()
model.build((n_tx, n_pix, n_el, n_ch))  # eagerly, see note below
beamform = Beamform(beamformer="able", model=model, num_patches=10)

Since the weights are trainable, the ABLE model has to be built before the pipeline is first called: building it inside a traced pipeline would create the convolution weights as tracers that escape their scope.

Warning

Train with jit_options=None on the enclosing pipeline. The default ("ops") compiles the surrounding operations separately, which cuts the gradient path to the model weights: the loss then simply does not move. Compiling is fine once the model is trained.

Parameters:

model (ABLE or None) – The ABLE model predicting the apodization weights. A default ABLE is created when omitted.
**kwargs – Forwarded to Operation.
input_data_type (DataTypes or None) – Expected data type of the input tensor. Used for pipeline data-type validation; pass None to skip.
output_data_type (DataTypes or None) – Data type produced by this operation.
key (str or None) – Dict key the operation reads from (and writes to by default). Defaults to "data".
output_key (str or None) – Dict key the operation writes its result to. Defaults to key. Set to a different value to preserve the original input under key while producing a new key for downstream operations.
cache_inputs (bool) – When True, values stored via set_input_cache() are merged into every call. False means the cache is empty by default. Selective per-key caching is not supported; use set_input_cache() directly to control which keys are stored.
cache_outputs (bool) – Memoize outputs keyed by a hash of the merged inputs.
jit_compile (bool) – Wrap call() with jit() for faster execution. Disable for easier interactive debugging.
with_batch_dim (bool) – Whether inputs carry a leading batch dimension. Affects default axis selection in filter-type operations.
jit_kwargs (dict or None) – Extra keyword arguments forwarded to the JIT compiler.
jittable (bool) – Mark the operation as JIT-compilable. Set to False for operations that use Python control flow incompatible with tracing.
additional_output_keys (list of str or None) – Extra dict keys this operation may produce beyond output_key. Used for pipeline key-availability validation. Defaults to the class-level ADD_OUTPUT_KEYS list.

call(**kwargs)[source]¶

Apply the ABLE weights to TOF-corrected data and sum.

Parameters:

data (ops.Tensor) – The TOF corrected input of shape (n_tx, prod(grid.shape), n_el, n_ch) with optional batch dimension.

Returns:

Dictionary containing beamformed_data: of shape (prod(grid.shape), n_ch) with optional batch dimension.

Return type:

dict

get_dict(compact=True)[source]¶

Serialize the operation, leaving out the ABLE model.

The model carries trained weights, which do not belong in a pipeline config; save and load it separately (see BaseModel). A pipeline restored from this config therefore gets a freshly initialized ABLE.