Working with a sequence of images

In this notebook, we will look at how to work with a sequence of images, and possibly change some parameters on the fly. As in all examples, we will start with setting the Keras backend.

   

%%capture
%pip install zea

import os

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

import keras
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML

import zea
from zea import init_device, load_file
from zea.visualize import set_mpl_style

zea: Using backend 'jax'

n_frames = 15
n_tx = 11
n_tx_total = 127

We will work with the GPU if available, and initialize using init_device to pick the best available device. Also, (optionally), we will set the matplotlib style for plotting.

init_device(verbose=False)
set_mpl_style()

We create a small helper function to animate the sequence of images.

def animate_images(images, scan, interval=100, cmap="gray"):
    """Helper function to animate a list of images."""
    fig, ax = plt.subplots(figsize=(5, 4), dpi=80)
    xlims_mm = [v * 1e3 for v in scan.xlims]
    zlims_mm = [v * 1e3 for v in scan.zlims]
    im = ax.imshow(
        np.array(images[0]),
        animated=True,
        cmap=cmap,
        extent=[xlims_mm[0], xlims_mm[1], zlims_mm[1], zlims_mm[0]],
    )
    ax.set_xlabel("X (mm)")
    ax.set_ylabel("Z (mm)")

    def update(frame):
        im.set_array(np.array(images[frame]))
        return [im]

    ani = animation.FuncAnimation(
        fig,
        update,
        frames=len(images),
        blit=True,
        interval=interval,
    )
    plt.close(fig)
    return HTML(ani.to_jshtml(fps=1000 // interval, embed_frames=True, default_mode="reflect"))

Let’s initialize a default B-mode ultrasound image formation pipeline.

pipeline = zea.Pipeline.from_default(pfield=False, with_batch_dim=False)

We will load a sequence of acquired RF data frames (carotid scan) and reconstruct a B-mode image from each frame. We will then animate the sequence of images. But first let’s load the data and parameters.

# this can take a while to download
file_path = "hf://zeahub/zea-carotid-2023/2_cross_bifur_right_0000.hdf5"  # ~25GB
# so let's use the smaller one by default:
file_path = "hf://zeahub/zea-carotid-2023/2_cross_bifur_right_0000_small.hdf5"  # ~2.5GB

frames = list(range(n_frames))  # use first 15 frames for demonstration
data, scan, probe = load_file(file_path, "raw_data", indices=frames)

scan.set_transmits(n_tx)  # reduce number of transmits for faster processing
scan.zlims = (0, 0.04)  # reduce z-limits a bit for better visualizations
scan.xlims = probe.xlims
scan.n_ch = data.shape[-1]  # rf data

config = zea.Config(dynamic_range=(-40, 0))

zea: WARNING The probe geometry in the data file does not match the probe geometry of the probe. The probe geometry has been updated to match the data file.

Reconstructing a sequence of B-mode images

images = []
n_frames = data.shape[0]
progbar = keras.utils.Progbar(n_frames, stateful_metrics=["frame"])

params = pipeline.prepare_parameters(probe, scan, config)

for frame_no in range(n_frames):
    output = pipeline(data=data[frame_no, scan.selected_transmits], **params)
    image = output.pop("data")
    params = output
    images.append(image)
    progbar.update(frame_no + 1)

animate_images(images, scan, interval=100, cmap="gray")

zea: Loading cached result for compute_pfield.
15/15 ━━━━━━━━━━━━━━━━━━━━ 8s 375ms/step

Once Loop Reflect

Change transmits on the fly

We now used 11 transmits throughout for every frame. We can also sweep through the the transmits for each frame to see how it affects the image quality.

images = []
n_frames = len(frames)
progbar = keras.utils.Progbar(n_frames, stateful_metrics=["frame"])

for idx, frame_no in enumerate(frames):
    tx_idx = int(round(idx * (n_tx_total - 1) / (n_frames - 1)))
    scan.set_transmits([tx_idx])
    raw_data_frame = data[frame_no, tx_idx][None, ...]
    with zea.log.set_level("WARNING"):  # to surpress info messages
        params = pipeline.prepare_parameters(probe, scan)
        output = pipeline(data=raw_data_frame, **params)
    images.append(output["data"])
    progbar.update(idx + 1)

animate_images(images, scan, interval=100, cmap="gray")

15/15 ━━━━━━━━━━━━━━━━━━━━ 2s 58ms/step

Once Loop Reflect