Using nabu from python to reconstruct a dataset with GPU

This notebook shows how to use the Nabu software for performing a basic reconstruction of a tomography dataset.

The computations are done on a local machine with a GPU and Cuda available.

This tutorial goes a bit further than nabu_basic_reconstruction.ipynb: - GPU implementation of each component is used - We see how to start from a configuration file and devise a simple processing chain accordingly

The same dataset is used (binned scan of a bamboo stick, thanks Ludovic Broche, ESRF ID19).

1 - Load the dataset informations

We must provide nabu with the the configuration file (nabu.conf), describing the path to the dataset and the processing steps. This is the equivalent of the .par file in PyHST2. In this file, no information is given on the detector size, energy, distance, etc: these informations are extracted from the dataset metadata.

import os
from nabu.testutils import utilstest, get_file
from nabu.pipeline.fullfield.processconfig import ProcessConfig

print("Getting dataset (downloading if necessary) ...")
data_path = get_file("bamboo_reduced.nx")
print("... OK")

# Get the configuration file of this dataset
conf_fname = get_file("bamboo_reduced.conf")

# Change directory to the path where the data is located (only useful for this tutorial)
os.chdir(utilstest.data_home)

# Parse this configuration file
conf = ProcessConfig(conf_fname)

Getting dataset (downloading if necessary) ...
... OK
Browsing dataset
Updating dataset information with user configuration
Could not load darks from /tmp/nabu_testdata_pierre/bamboo_reduced_darks.hdf5
Could not load flats from /tmp/nabu_testdata_pierre/bamboo_reduced_flats.hdf5
Saved reduced darks to /tmp/nabu_testdata_pierre/bamboo_reduced_darks.hdf5
Saved reduced flats to /tmp/nabu_testdata_pierre/bamboo_reduced_flats.hdf5
Doing dataset estimations
Estimating center of rotation
CenterOfRotationSlidingWindow.find_shift({'window_width': None, 'roi_yxhw': None, 'median_filt_shape': None, 'padding_mode': None, 'peak_fit_radius': 1, 'high_pass': None, 'low_pass': None, 'return_validity': False, 'cor_options': None})
Estimated center of rotation: 339.486
Doing coupled validation
Cannot do SRCurrent normalization: missing flats and/or projections SRCurrent

Note that ProcessConfig will do quite a few things under the hood: - Parse the configuration file and check parameters correctness - Browse the dataset - Get or compute the reduced flats/darks - Estimate the center of rotation

The resulting object contains all necessary information to process the dataset.

# We can easily get information on the processing steps.
nabu_config = conf.nabu_config
print(nabu_config)
# The same can be done with the dataset structure
dataset_info = conf.dataset_info
# print([getattr(dataset_info, attr) for attr in ["energy", "distance", "n_angles", "radio_dims"]])

{'dataset': {'location': '/tmp/nabu_testdata_pierre/bamboo_reduced.nx', 'hdf5_entry': None, 'nexus_version': 1.0, 'darks_flats_dir': None, 'binning': 1, 'binning_z': 1, 'projections_subsampling': (1, 0), 'exclude_projections': None, 'overwrite_metadata': ''}, 'preproc': {'flatfield': True, 'flat_distortion_correction_enabled': False, 'flat_distortion_params': "tile_size=100; interpolation_kind='linear'; padding_mode='edge'; correction_spike_threshold=None", 'normalize_srcurrent': True, 'ccd_filter_enabled': False, 'ccd_filter_threshold': 0.04, 'detector_distortion_correction': None, 'detector_distortion_correction_options': None, 'double_flatfield_enabled': False, 'dff_sigma': None, 'take_logarithm': True, 'log_min_clip': 1e-06, 'log_max_clip': 10.0, 'sino_normalization': None, 'sino_normalization_file': '', 'processes_file': None, 'sino_rings_correction': 'munch', 'sino_rings_options': None, 'rotate_projections': None, 'rotate_projections_center': None, 'tilt_correction': None, 'autotilt_options': None}, 'phase': {'method': 'paganin', 'delta_beta': 100.0, 'unsharp_coeff': 0.0, 'unsharp_sigma': 0.0, 'unsharp_method': 'gaussian', 'padding_type': 'edge', 'ctf_geometry': 'z1_v=None; z1_h=None; detec_pixel_size=None; magnification=True', 'ctf_advanced_params': 'length_scale=1e-5; lim1=1e-5; lim2=0.2; normalize_by_mean=True'}, 'reconstruction': {'method': 'FBP', 'angles_file': None, 'rotation_axis_position': 'sliding-window', 'cor_options': "side='from_file'", 'cor_slice': None, 'axis_correction_file': None, 'translation_movements_file': None, 'angle_offset': 0.0, 'fbp_filter_type': 'ramlak', 'fbp_filter_cutoff': 1.0, 'source_sample_dist': None, 'sample_detector_dist': None, 'padding_type': 'edge', 'enable_halftomo': 'auto', 'clip_outer_circle': False, 'centered_axis': False, 'start_x': 0, 'end_x': -1, 'start_y': 0, 'end_y': -1, 'start_z': 0, 'end_z': -1, 'iterations': 200, 'optim_algorithm': 'chambolle-pock', 'weight_tv': 0.01, 'preconditioning_filter': True, 'positivity_constraint': True}, 'output': {'location': '/tmp/nabu_testdata_pierre', 'file_prefix': 'bamboo_reduced_rec', 'file_format': 'hdf5', 'overwrite_results': True, 'tiff_single_file': False, 'jpeg2000_compression_ratio': None, 'float_clip_values': None}, 'postproc': {'output_histogram': False, 'histogram_bins': 1000000}, 'resources': {'method': 'local', 'gpus': 1, 'gpu_id': [], 'cpu_workers': 0, 'memory_per_node': (90.0, True), 'threads_per_node': (100.0, True), 'queue': 'gpu', 'walltime': (1, 0, 0)}, 'pipeline': {'save_steps': None, 'resume_from_step': None, 'steps_file': None, 'verbosity': 'info'}, 'about': {}}

2 - Chunk processing

Nabu processes data by chunks of radios (see the documentation for more explanations).

In a first step, we define how to read chunks of radios.

from nabu.io.reader import ChunkReader

What is the largest chunk size we can process ?

The answer is given by inspecting the current GPU memory, and the processing steps.

from nabu.cuda.utils import get_gpu_memory
from nabu.pipeline.fullfield.computations import estimate_max_chunk_size

chunk_size = estimate_max_chunk_size(
    get_gpu_memory(0),
    conf
)
print("Chunk_size = %d" % chunk_size)

Chunk_size = 540

# Load the first 'chunk_size' lines of all the radios
sub_region = (None, None, 0, chunk_size) # start_x, end_x, start_z, end_z
chunk_reader = ChunkReader(
    dataset_info.projections,
    sub_region=sub_region,
    convert_float=True
)

# Load the current chunk
chunk_reader.load_files() # takes some time

print(chunk_reader.files_data.shape)
print(chunk_reader.files_data.dtype)

(1000, 540, 640)
float32

3 - Initialize the GPU

Most of the processing can be done on GPU (or many-core CPU if using OpenCL).

With pycuda.gpuarray (or its OpenCL counterpart pyopencl.array), we manipulate array objects with memory residing on device. This allows to avoid extraneous host <-> device copies.

import pycuda.gpuarray as garray
from nabu.cuda.utils import get_cuda_context
import numpy as np

# Create a Cuda context on device ID 0
# By default, all following GPU processings will be bound on this context
ctx = get_cuda_context(device_id=0)

radios = chunk_reader.files_data
n_angles, n_z, n_x = radios.shape
# transfer the chunk on GPU
d_radios = garray.to_gpu(radios)

4 - Pre-processing

Pre-processing utilities are available in the nabu.preproc module.

Utilities available with the cuda backend are implemented in a module with a _cuda suffix.

4.1 - Flat-field

from nabu.preproc.flatfield_cuda import CudaFlatFieldDataUrls

radios_indices = sorted(conf.dataset_info.projections.keys())
# Configure the `FlatField` processor
cuda_flatfield = CudaFlatFieldDataUrls(
    d_radios.shape,
    dataset_info.flats,
    dataset_info.darks,
    radios_indices=radios_indices,
    sub_region=sub_region,
    convert_float=True,
)

# Perform the normalization on GPU
if nabu_config["preproc"]["flatfield"]:
    print("Doing flat-field")
    cuda_flatfield.normalize_radios(d_radios)

Doing flat-field

4.2 - Phase retrieval

from nabu.preproc.phase_cuda import CudaPaganinPhaseRetrieval

energy = dataset_info.energy
# Phase retrieval is done on each radio individually, with the sub-region specified above
if (nabu_config["phase"]["method"] or "").lower() == "paganin":
    print("Doing phase retrieval")
    cudapaganin = CudaPaganinPhaseRetrieval(
        (n_z, n_x),
        distance=dataset_info.distance,
        energy=energy,
        delta_beta=nabu_config["phase"]["delta_beta"],
        pixel_size=dataset_info.pixel_size * 1e6,
    )
    for i in range(n_angles):
        cudapaganin.apply_filter(d_radios[i], output=d_radios[i])

Doing phase retrieval

/home/pierre/.venv/py311/lib/python3.11/site-packages/skcuda/cublas.py:284: UserWarning: creating CUBLAS context to get version number
  warnings.warn('creating CUBLAS context to get version number')

4.3 - Logarithm

from nabu.preproc.ccd_cuda import CudaLog

if nabu_config["preproc"]["take_logarithm"]:
    print("Taking logarithm")
    cuda_log = CudaLog(d_radios.shape, clip_min=0.01)
    cuda_log.take_logarithm(d_radios)

Taking logarithm

5 - Reconstruction

We use the filtered backprojection with nabu.reconstruction.fbp

from nabu.reconstruction.fbp import Backprojector

rec_options = conf.processing_options["reconstruction"]
B = Backprojector(
    (n_angles, n_x),
    angles=rec_options["angles"],
    rot_center=rec_options["rotation_axis_position"],
    padding_mode="edges"
)
d_recs = garray.zeros((n_z, n_x, n_x), "f")

print("Reconstructing...", end="")
for i in range(n_z):
    B.fbp(radios[:, i, :], output=d_recs[i])
recs = d_recs.get()
print(" ... OK")

Reconstructing... ... OK

6 - Visualize

%pylab nbagg

%pylab is deprecated, use %matplotlib inline and import the required libraries.
Populating the interactive namespace from numpy and matplotlib

figure()
imshow(recs[0], cmap="gray")

<matplotlib.image.AxesImage at 0x7f1d095b1fd0>