zea.models.inversionnet¶

InversionNet: full-waveform inversion of ultrasound computed tomography data.

InversionNet is an encoder-decoder network that maps raw multi-source waveform data straight to a speed-of-sound (SOS) map, without an iterative solver. It was introduced for seismic FWI (Wu & Lin, 2020), and is the reference baseline of the OpenPros limited-view prostate USCT benchmark.

Usage¶

from zea.models.inversionnet import InversionNet

model = InversionNet.from_preset("inversionnet-openpros")
sos = model(waveforms)  # (B, 1000, 161, 40) -> (B, 401, 161, 1) in [-1, 1]

Input preprocessing¶

The preset only works on input scaled the way it was trained, and mis-scaled input degrades the reconstruction silently rather than raising. Reproduce the OpenPros preprocessing exactly:

  1. Sign-preserving log compression, \(t(x) = \mathrm{sign}(x)\log(1 + |kx|)\), with k = 1e5.

  2. Min-max normalize [t(data_min), t(data_max)] to [-1, 1], with the OpenPros dataset constants data_min = -0.25 and data_max = 0.45.

The output is in [-1, 1] and maps linearly onto the OpenPros label range of 1300-3600 m/s — undo it with Normalize(input_range=(-1, 1), output_range=(1300, 3600)).

Note

k is effectively part of the weights, not a free parameter. The OpenPros job scripts pass --k 1e9, but k = 1e5 is what reproduces the released checkpoint on the released data: on the OpenPros sample it gives a mean absolute error of 14 m/s against the ground-truth map, where k = 1e9 gives 244 m/s. The optimum is sharp — an order of magnitude either way costs roughly 3x in error — so do not tune it.

Architecture notes¶

  • Encoder: a stride-2 stack that collapses the (time, receiver) plane to 1 x 1 at 512 channels. The first levels stride over time only, because the waveform axis (1000 samples) is much longer than the receiver axis (161).

  • Decoder: transposed convolutions back up to 448 x 192, cropped to the 401 x 161 image grid, followed by a tanh output block.

  • The output is in [-1, 1]; map it to physical units with Normalize (OpenPros uses 1300-3600 m/s).

References

Y. Wu and Y. Lin. InversionNet: An Efficient and Accurate Data-Driven Full Waveform Inversion. IEEE Transactions on Computational Imaging, 6:419-433, 2020. DOI: 10.1109/TCI.2019.2956866 (arXiv:1811.07875)

H. Wang, Y. Wu, Y. Feng, P. Jin, L. Zhang, S. Feng, J. Wiskin, B. Turkbey, P. A. Pinto, B. J. Wood, S. Luo, Y. Chen, E. Boctor and Y. Lin. OpenPros: A Large-Scale Dataset for Limited View Prostate Ultrasound Computed Tomography. 2025. arXiv:2505.12261

Important

This is a zea implementation of the model. The inversionnet-openpros weights are the pretrained baseline released with the OpenPros benchmark (dataset, code, CC-BY-4.0). Please cite both papers above when you use them.

Note

Because the encoder ends in a fixed 8 x 6 convolution and the decoder in a fixed crop, the input size is part of the architecture. The OpenPros preset expects exactly OPENPROS_INPUT_SHAPE.

Module Attributes

OPENPROS_INPUT_SHAPE

Input shape (time, receivers, sources) the OpenPros preset was trained on.

Classes

InversionNet(*args, **kwargs)

Encoder-decoder network mapping USCT waveforms to a speed-of-sound map.
class zea.models.inversionnet.InversionNet(*args, **kwargs)[source]¶

Bases: BaseModel

Encoder-decoder network mapping USCT waveforms to a speed-of-sound map.

Parameters:

  • waveform_shape (tuple) – Input shape (time, receivers, channels), without the batch axis. The architecture is tied to it — see the module note — so it is part of the config. Defaults to OPENPROS_INPUT_SHAPE.

  • enc_ch (tuple) – Base-2 exponents of the encoder channel widths. The first and last entry are the input and bottleneck convolution; the entries in between each become a stride-2 block plus a stride-1 block.

  • enc_side (tuple) – Per intermediate encoder level, whether to stride over the time axis only (1) instead of both axes (0). Must have len(enc_ch) - 2 entries.

  • bottle_conv (tuple) – Kernel of the final, "valid" encoder convolution. It collapses the feature map to 1 x 1, so it must equal that map’s spatial size.

  • bottle_deconv (tuple) – Kernel of the first decoder transposed convolution, which expands 1 x 1 back to bottle_deconv.

  • dec_ch (tuple) – Base-2 exponents of the decoder channel widths. Every entry past the first doubles the spatial size.

  • crop (tuple) – (top, bottom, left, right) pixels to remove from the decoder output to reach the image grid.

Example

import numpy as np
from zea.models.inversionnet import InversionNet, OPENPROS_INPUT_SHAPE

model = InversionNet.from_preset("inversionnet-openpros")
waveforms = np.zeros((1, *OPENPROS_INPUT_SHAPE), dtype="float32")
sos = model(waveforms)  # (1, 401, 161, 1), in [-1, 1]
call(inputs, training=None)[source]¶

Reconstruct a speed-of-sound map from waveform data.

Parameters:

  • inputs (array-like) – Waveforms of shape (B, time, receivers, channels), scaled exactly as during training — see the preprocessing section in the module documentation, which the preset depends on.

  • training (bool, optional) – Forwarded to the batch normalization layers, which use the running statistics unless this is True.

Returns:

Speed-of-sound maps of shape (B, height, width, 1), in [-1, 1].

Return type:

Tensor

custom_load_weights(preset, backend='keras', **kwargs)[source]¶

Load weights from a preset (Hugging Face handle or local directory).

Parameters:

  • preset (str) – Preset identifier passed from from_preset().

  • backend (str) – "keras" loads model.weights.h5; "torch" converts the original model.pth checkpoint and needs PyTorch.

classmethod from_pth(pth_path, **kwargs)[source]¶

Create an InversionNet from an original PyTorch checkpoint.

Parameters:

  • pth_path (str) – Path to the .pth state dict.

  • **kwargs – Passed to the constructor.

Returns:

Model with the converted weights.

Return type:

InversionNet

get_config()[source]¶

Serialize the architecture arguments.

zea.models.inversionnet.OPENPROS_INPUT_SHAPE = (1000, 161, 40)¶

Input shape (time, receivers, sources) the OpenPros preset was trained on.