Source code for ewokstomo.tasks.online.nabu_pipeline
import numpy as np
from collections import namedtuple
from nabu.pipeline.fullfield.processconfig import ProcessConfig
from nabu.pipeline.fullfield.reconstruction import compute_reconstruction_margin
from nabu.cuda.utils import __has_cupy__
if __has_cupy__:
from nabu.pipeline.fullfield.accumulated_cuda import CudaAccumulatedPipeline
default_pipeline_cls = CudaAccumulatedPipeline
else:
from nabu.pipeline.fullfield.accumulated import AccumulatedPipeline
default_pipeline_cls = AccumulatedPipeline
DataPlaceHolder = namedtuple("DataPlaceHolder", ["shape"])
[docs]
def get_nabu_pipeline(
config_dict, dataset_info, projs_batch_size, pipeline_backend="cuda"
):
"""
Create a reconstruction pipeline object, based on configuration and scan info.
Parameters
----------
config_dict: dict
dict of dict, see 'default_nabu_config' for structure and keys.
dataset_info: nabu.resources.dataset.Dataset
Dataset structure, InMemoryDataset in this case
projs_batch_size: int
Number of projections processed at once ; i.e we trigger the pipeline as soon as we have 'projs_batch_size' in memory
"""
dataset_info.rotation_angles = np.zeros(projs_batch_size, dtype="f")
# Build the ProcessConfig data structure - this does extra checks and pre-calculates things for the pipeline
process_config = ProcessConfig(conf_dict=config_dict, dataset_info=dataset_info)
# Sub-region to be reconstructed.
# The full data volume has shape (n_z, n_y, n_x), but we might want to reconstruct a region of interest
# this is defined by start_{x|y|z} and end_{x|y|z} in the config file
# This holds for slice reconstruction. In this case we also need to compute the processing margin.
(start_z, end_z), margin_ud = get_reconstruction_z_region_and_margin(process_config)
sub_region = (
(0, projs_batch_size),
(start_z, end_z),
(0, dataset_info.radio_dims[0]),
)
margin = None if margin_ud is None else (margin_ud, (0, 0))
# Shape of the data volume that will be ingested by a pipeline instance,
# in the form (n_projections, n_vertic, n_horiz)
chunk_shape = (projs_batch_size,) + tuple(sr[1] - sr[0] for sr in sub_region[1:])
# Add place-holder for some fields, otherwise pipeline will not instantiate
dataset_info.data = DataPlaceHolder(chunk_shape)
dataset_info.current_projections_indices = np.zeros(
projs_batch_size, dtype=np.uint64
)
pipeline_cls = default_pipeline_cls
if pipeline_backend == "cuda" and __has_cupy__:
pipeline_cls = CudaAccumulatedPipeline
pipeline = pipeline_cls(process_config, chunk_shape, margin=margin)
pipeline.target_sub_region = sub_region
return pipeline
[docs]
def get_reconstruction_z_region_and_margin(process_config):
"""
Get parameters relevant for processing pipeline: sub-region to reconstruct, and processing margin.
Parameters
----------
process_config: nabu.pipeline.process_config.ProcessConfig object
Data structure with the processing configuration
Returns
-------
start_end_z: tuple of int
Tuple in the form (start_z, end_z), defining the start and stop of slices indices to be reconstructed
margin_ud: tuple of int
Tuple in the form (margin_up, margin_down): "processing margin" in the vertical direction.
Data will be cropped as: reconstruction[margin_up:-margin_down, ...]
"""
rec_region = process_config.rec_region
start_z = rec_region["start_z"]
end_z = (
rec_region["end_z"] + 1
) # cf difference between nabu cfg interface and Python
n_z, n_x = process_config.dataset_info.radio_dims[::-1]
is_reconstructing_full_volume = True
if start_z > 0 or end_z < n_z - 1:
is_reconstructing_full_volume = False
margin_up_down = None
if not (is_reconstructing_full_volume):
margin_v, margin_h = compute_reconstruction_margin(
process_config,
user_margin=process_config.nabu_config["pipeline"].get(
"processing_margin", None
),
)
new_start_z = max(start_z - margin_v, 0)
new_end_z = min(end_z + margin_v, n_z - 1)
margin_up_down = (start_z - new_start_z), (new_end_z - end_z)
start_z = new_start_z
end_z = new_end_z
return (start_z, end_z), margin_up_down
