vbeam: a fast and differentiable beamformer
Project description
vbeam: a Fast and Differentiable Beamformer for Optimizing Ultrasound Imaging
vbeam is an open-source beamforming library built on top of machine learning frameworks. This enables fast processing, rapid prototyping, and integrating beamforming into machine learning algorithms.
Run An Example In Google Colab (In Your Browser)
Run one of the examples in your browser via Google Colab:
(Make sure to select an A100 (or V100 if not available) GPU to make the examples run fast!)
- Time-domain REFoCUS with a heterogeneous speed-of-sound map.
- Plane wave imaging (PWI) of a PICMUS dataset.
- Retrospective transmit beamforming (RTB) with optimized scans.
Installation
python3 -m pip install vbeam
You probably also want to install JAX. See JAX installation instructions for how to do so.
IUS 2023 Q&A and Discussion
Visit the GitHub Discussion page for IUS 2023 if you have any questions!
Overview
vbeam is a pixel-based beamformer. This means that it centers around a function called signal_for_point
which returns the delayed, interpolated signal for a single point/pixel. To get a full image, simply repeat signal_for_point
for all pixels, all receiving elements, transmitted waves, frames, etc.
We take advantage of JAX's vmap
transformation to call signal_for_point
in parallel on the GPU. vbeam is a "vmapped beamformer", hence the name.
We refer to the dimensions of our data by name, using a library called spekk. This lets us write beamforming code that looks like this:
beamformer = compose(
signal_for_point,
# Make it run in parallel for all points (pixels), sending-, and receiving- elements
ForAll("points"),
ForAll("senders"),
ForAll("receivers"),
# We sum over all sending-, and receiving- elements
Apply(jnp.sum, [Axis("senders"), Axis("receivers")]),
# Iteratively sum the result of each transmit event to reduce memory usage
Reduce.Sum("transmits"),
# Postprocessing
Apply(setup.scan.unflatten), # Points are flattened by default; reshape them back
Apply(normalized_decibels), # Decibels, and the maximum value is normalized to 0 dB
Apply(scan_convert, setup.scan), # Scan convert from beam-space to cartesian space
# Just-in-time compile the function
Wrap(jax.jit),
# Build the function to let vbeam know about the dimensions of the input data
).build(setup.spec)
But usually, vbeam's out-of-the-box beamformer will be enough:
beamformer = jax.jit(get_das_beamformer(setup))
The above example refers to 4 dimensions: "points"
, "senders"
, "receivers"
, and "transmits"
. In this case, "senders"
refers to the individual contributing elements used during wave transmission, and is part of our time-domain REFoCUS implementation . vbeam allows you to run any arbitrary dimension in parallel as part of your beamformer: perhaps you want to beamform multiple frames in parallel, or with multiple F#s.
If you want to use a different wavefront model, or a different apodization function, you simply pass in the new object when you call the beamformer:
kwargs = setup.data
kwargs["apodization"] = MyCustomApodization()
result = beamformer(**kwargs)
Every beamforming component in vbeam is differentiable, so you can use vbeam to optimize your beamforming algorithm using gradient descent. You can also integrate machine learning into your beamformer, or integrate beamforming as a step in your machine learning algorithm. Using JAX, it is easy to get the gradient of any of the beamformer arguments:
from vbeam.util.jax import grad_for_argname
# Minimize the pixel-values (just for illustration purposes) of the beamformed image by
# "optimizing" the apodization:
my_loss = lambda **kwargs: jnp.sum(beamformer(**kwargs))
loss_grad = grad_for_argname(my_loss, "apodization")
Documentation
You can read the documentation of vbeam's core here.
Citing vbeam
If you use vbeam in your research, please cite the following paper:
Please wait for the in-proceedings article to be published :)
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.