Working with the zea data format

In this tutorial notebook we will show how to load a zea data file and how to access the data stored in it. There are three common ways to load a zea data file:

1. Loading data from single file with zea.File

2. Loading data from a group of files with zea.Dataset

3. Loading data in batches with dataloading utilities with zea.Dataloader

   

‼️ 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

config_picmus_iq = "hf://zeahub/configs/config_picmus_iq.yaml"

import os

os.environ["KERAS_BACKEND"] = "jax"
os.environ["ZEA_DISABLE_CACHE"] = "1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import keras
import matplotlib.pyplot as plt

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

zea: Using backend 'jax'

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()

Loading a file with zea.File

The zea data format works with HDF5 files. We can open a zea data file using the h5py package and have a look at the contents using the zea.File.summary() function. You can see that every dataset element contains a corresponding description and unit. Note that we now pass a url to a Hugging Face dataset, but you can also use a local file path to a zea data file. Here we will use an example from the PICMUS dataset, converted to zea format and hosted on the Hugging Face Hub.

Tip: You can also use the HDFView tool to view the contents of the zea data file without having to run any code. Or if you use VS Code, you can install the HDF5 extension to view the contents of the file.

You can extract data and acquisition parameters (which are stored together with the data in the zea data file) as follows:

file_path = "hf://zeahub/picmus/database/experiments/contrast_speckle/contrast_speckle_expe_dataset_iq/contrast_speckle_expe_dataset_iq.hdf5"

with zea.File(file_path, mode="r") as file:
    file.summary()

    data = file.load_data("raw_data", indices=0)
    scan = file.scan()
    probe = file.probe()

contrast_speckle_expe_dataset_iq.hdf5/
    ├── description: PICMUS dataset converted to USBMD format
    ├── probe: verasonics_l11_4v
    ├── data/
    │   ├── description: This group contains the data.
    │   └── raw_data/
    │       ├── /data/raw_data (shape=(1, 75, 832, 128, 2))
    │       │   ├── description: The raw_data of shape (n_frames, n_tx, n_el, n_ax, n_ch).
    │       │   ├── unit: unitless
    └── scan/
        ├── description: This group contains the scan parameters.
        ├── azimuth_angles/
        │   ├── /scan/azimuth_angles (shape=(75))
        │   │   ├── description: The azimuthal angles of the transmit beams in radians of shape (n_tx,).
        │   │   ├── unit: rad
        ├── center_frequency/
        │   ├── /scan/center_frequency (shape=())
        │   │   ├── description: The center frequency in Hz.
        │   │   ├── unit: Hz
        ├── focus_distances/
        │   ├── /scan/focus_distances (shape=(75))
        │   │   ├── description: The transmit focus distances in meters of shape (n_tx,). For planewaves this is set to Inf.
        │   │   ├── unit: m
        ├── initial_times/
        │   ├── /scan/initial_times (shape=(75))
        │   │   ├── description: The times when the A/D converter starts sampling in seconds of shape (n_tx,). This is the time between the first element firing and the first recorded sample.
        │   │   ├── unit: s
        ├── n_ax/
        │   ├── /scan/n_ax (shape=())
        │   │   ├── description: The number of axial samples.
        │   │   ├── unit: unitless
        ├── n_el/
        │   ├── /scan/n_el (shape=())
        │   │   ├── description: The number of elements in the probe.
        │   │   ├── unit: unitless
        ├── n_frames/
        │   ├── /scan/n_frames (shape=())
        │   │   ├── description: The number of frames.
        │   │   ├── unit: unitless
        ├── n_tx/
        │   ├── /scan/n_tx (shape=())
        │   │   ├── description: The number of transmits per frame.
        │   │   ├── unit: unitless
        ├── polar_angles/
        │   ├── /scan/polar_angles (shape=(75))
        │   │   ├── description: The polar angles of the transmit beams in radians of shape (n_tx,).
        │   │   ├── unit: rad
        ├── probe_geometry/
        │   ├── /scan/probe_geometry (shape=(128, 3))
        │   │   ├── description: The probe geometry of shape (n_el, 3).
        │   │   ├── unit: m
        ├── sampling_frequency/
        │   ├── /scan/sampling_frequency (shape=())
        │   │   ├── description: The sampling frequency in Hz.
        │   │   ├── unit: Hz
        ├── sound_speed/
        │   ├── /scan/sound_speed (shape=())
        │   │   ├── description: The speed of sound in m/s
        │   │   ├── unit: m/s
        ├── t0_delays/
        │   ├── /scan/t0_delays (shape=(75, 128))
        │   │   ├── description: The t0_delays of shape (n_tx, n_el).
        │   │   ├── unit: s
        └── tx_apodizations/
            ├── /scan/tx_apodizations (shape=(75, 128))
            │   ├── description: The transmit delays for each element defining the wavefront in seconds of shape (n_tx, n_elem). This is the time at which each element fires shifted such that the first element fires at t=0.
            │   ├── unit: unitless

You can also do this exact thing in one go using zea.load_file which will directly return you the data and parameter objects.

data, scan, probe = load_file(file_path, "raw_data", indices=(0, slice(0, 3)))

print("Raw data shape:", data.shape)
print(scan)
print(probe)

Raw data shape: (3, 832, 128, 2)
Scan(
    azimuth_angles=array(shape=(75,)),
    center_frequency=5208000.0,
    focus_distances=array(shape=(75,)),
    initial_times=array(shape=(75,)),
    n_ax=832,
    n_el=128,
    n_frames=1,
    n_tx=75,
    polar_angles=array(shape=(75,)),
    probe_geometry=array(shape=(128, 3)),
    sampling_frequency=5208000.0,
    sound_speed=1540.0,
    t0_delays=array(shape=(75, 128)),
    tx_apodizations=array(shape=(75, 128)),
    selected_transmits=[0, 1, 2],
    pixels_per_wavelength=4,
    pfield_kwargs={},
    apply_lens_correction=False,
    grid_type=cartesian,
    bandwidth_percent=200.0,
    attenuation_coef=0.0,
    f_number=1.0
)
<zea.probes.Verasonics_l11_4v object at 0x73782531c590>

Loading data with zea.Dataset

We can also load and manage a group of files (i.e. a dataset) using the zea.Dataset class. Instead of a path to a single file, we can pass a list of file paths or a directory containing multiple zea data files. The zea.Dataset class will automatically load the files and allow you to access the data in a similar way as with zea.File.

dataset_path = "hf://zeahub/picmus/database/experiments"

dataset = zea.Dataset(dataset_path)

print(dataset)

for file in dataset:
    print(file)

dataset.close()

zea: Searching /tmp/zea_cache_83gfgryh/huggingface/datasets/datasets--zeahub--picmus/snapshots/07fe825b53c92b1d423fadb1dfa104ed2a38aa4a/database/experiments for ['.hdf5', '.h5'] files...
zea: Dataset validated. Check /tmp/zea_cache_83gfgryh/huggingface/datasets/datasets--zeahub--picmus/snapshots/07fe825b53c92b1d423fadb1dfa104ed2a38aa4a/database/experiments/validated.flag for details.
Dataset with 4 files
zea HDF5 File: 'contrast_speckle_expe_dataset_rf.hdf5' (mode=r)
zea HDF5 File: 'contrast_speckle_expe_dataset_iq.hdf5' (mode=r)
zea HDF5 File: 'resolution_distorsion_expe_dataset_rf.hdf5' (mode=r)
zea HDF5 File: 'resolution_distorsion_expe_dataset_iq.hdf5' (mode=r)

Loading data with Dataloader

In machine and deep learning workflows, we often want more features like batching, shuffling, and parallel data loading. The zea.Dataloader class provides a convenient way to create a high-performance data loader from a zea dataset. It is built on Grain and does not require TensorFlow. This dataloader is particularly useful for training models. Consistency of shape is preferred, which is not the case for PICMUS. Therefore in this example we will use a small part of the CAMUS dataset.

dataset_path = "hf://zeahub/camus-sample/val"
dataloader = zea.Dataloader(
    dataset_path,
    key="data/image_sc",
    batch_size=4,
    shuffle=True,
    clip_image_range=[-60, 0],
    image_range=[-60, 0],
    normalization_range=[0, 1],
    image_size=(256, 256),
    resize_type="resize",  # or "center_crop or "random_crop"
    seed=4,
)

for batch in dataloader:
    print("Batch shape:", batch.shape)
    break  # Just show the first batch

fig, _ = zea.visualize.plot_image_grid(batch)

zea: Searching /tmp/zea_cache_83gfgryh/huggingface/datasets/datasets--zeahub--camus-sample/snapshots/617cf91a1267b5ffbcfafe9bebf0813c7cee8493/val for ['.hdf5', '.h5'] files...
zea: Dataset validated. Check /tmp/zea_cache_83gfgryh/huggingface/datasets/datasets--zeahub--camus-sample/snapshots/617cf91a1267b5ffbcfafe9bebf0813c7cee8493/val/validated.flag for details.
zea: Caching is globally disabled for _find_h5_file_shapes.
zea: WARNING H5Generator: Not all files have the same shape. This can lead to issues when resizing images later....
zea: H5Generator: Shuffled data.
zea: H5Generator: Shuffled data.
Batch shape: (4, 256, 256, 1)

Processing an example

We will now use one of the zea data files to demonstrate how to process it. A full example can be found in the zea_pipeline_example notebook. Here we will just show a simple example for completeness. We will start by loading a config file, that contains all the required information to initiate a processing pipeline.

config = zea.Config.from_path(config_picmus_iq)

Now we can load the zea data file, extract data and parameters, and then process the data using the pipeline defined by the config file.

with zea.File(config.data.dataset_folder + "/" + config.data.file_path, mode="r") as file:
    # we use config here to overwrite some of the scan parameters
    scan = file.scan(**config.scan)
    data = file.load_data(config.data.dtype)

pipeline = zea.Pipeline.from_config(config)
parameters = pipeline.prepare_parameters(probe=probe, scan=scan)

images = pipeline(data=data, **parameters)["data"]
images = keras.ops.convert_to_numpy(images)

zea: WARNING No transmit origins provided, using zeros
zea: DEBUG [zea.Pipeline] The following input keys are not used by the pipeline: {'xlims', 'n_el', 'zlims', 'center_frequency'}. Make sure this is intended. This warning will only be shown once.

Finally we can plot the result.

image = zea.display.to_8bit(images[0], dynamic_range=(-50, 0))
plt.figure()
# Convert xlims and zlims from meters to millimeters for display
xlims_mm = [v * 1e3 for v in scan.xlims]
zlims_mm = [v * 1e3 for v in scan.zlims]
plt.imshow(image, cmap="gray", extent=[xlims_mm[0], xlims_mm[1], zlims_mm[1], zlims_mm[0]])
plt.xlabel("X (mm)")
plt.ylabel("Z (mm)")

Text(0, 0.5, 'Z (mm)')