Commit bdade8fb authored by Pierre Paleo's avatar Pierre Paleo
Browse files

Merge branch 'cone_reconstruction' into 'master'

Cone-beam reconstruction

See merge request !180
parents 24032831 8492b06e
Pipeline #78313 passed with stages
in 6 minutes and 46 seconds
__version__ = "2022.1.9"
__version__ = "2022.1.10"
__nabu_modules__ = [
"app",
"cuda",
......
......@@ -79,7 +79,7 @@ class CudaProcessing:
List of array names
"""
for array_name in arrays_names:
old_arr = getattr(self, "_old_" + array_name)
old_arr = getattr(self, "_old_" + array_name, None)
if old_arr is not None:
setattr(self, array_name, old_arr)
......@@ -121,17 +121,37 @@ class CudaProcessing:
Data type. Default is float32.
"""
if isinstance(array_ref, garray.GPUArray):
current_arr = getattr(self, array_name)
current_arr = getattr(self, array_name, None)
setattr(self, "_old_" + array_name, current_arr)
setattr(self, array_name, array_ref)
elif isinstance(array_ref, np.ndarray):
self._allocate_array(array_name, shape, dtype=dtype)
self.allocate_array(array_name, shape, dtype=dtype)
getattr(self, array_name).set(array_ref)
else:
raise ValueError("Expected numpy array or pycuda array")
def get_array(self, array_name):
return getattr(self, array_name, None)
# COMPAT.
_init_arrays_to_none = init_arrays_to_none
_recover_arrays_references = recover_arrays_references
_allocate_array = allocate_array
_set_array = set_array
@staticmethod
def check_array(arr, expected_shape, expected_dtype="f", check_contiguous=True):
"""
Check whether a given array is suitable for being processed (shape, dtype, contiguous)
"""
if arr.shape != expected_shape:
raise ValueError("Expected shape %s but got %s" % (str(expected_shape), str(arr.shape)))
if arr.dtype != np.dtype(expected_dtype):
raise ValueError("Expected data type %s but got %s" % (str(expected_dtype), str(arr.dtype)))
if check_contiguous:
if isinstance(arr, np.ndarray) and not(arr.flags["C_CONTIGUOUS"]):
raise ValueError("Expected C-contiguous array")
if isinstance(arr, garray.GPUArray) and not arr.flags.c_contiguous:
raise ValueError("Expected C-contiguous array")
import numpy as np
try:
import astra
__have_astra__ = True
except ImportError:
__have_astra__ = False
from ..cuda.processing import CudaProcessing
class ConebeamReconstructor:
"""
A reconstructor for cone-beam geometry using the astra toolbox.
"""
def __init__(
self,
sinos_shape,
source_origin_dist,
angles=None,
volume_shape=None,
rot_center=None,
relative_z_position=None,
cuda_options=None,
):
"""
Initialize a cone beam reconstructor. This reconstructor works on slabs of data,
meaning that one partial volume is obtained from one stack of sinograms.
To reconstruct a full volume, the reconstructor must be called on a series of sinograms stacks, with
an updated "relative_z_position" each time.
Parameters
-----------
sinos_shape: tuple
Shape of the sinograms stack, in the form (n_angles, n_y, n_x).
source_origin_dist: float
Distance, in pixel units, between the beam source (cone apex) and the "origin".
The origin is defined as the center of the sample
angles: array, optional
Rotation angles in radians. If provided, its length should be equal to sinos_shape[0].
volume_shape: tuple of int, optional
Shape of the output volume slab, in the form (n_z, n_y, n_x).
If not provided, the output volume slab shape is (sinos_shape[0], sinos_shape[2], sinos_shape[2]).
rot_center: float, optional
Rotation axis position. Default is `(detector_width - 1)/2.0`
relative_z_position: float, optional
Position of the central slice of the slab, with respect to the full stack of slices.
By default it is set to zero, meaning that the current slab is assumed in the middle of the stack
Notes
-----
This reconstructor is using the astra toolbox [1]. Therefore the implementation uses Astra's
reference frame, which is centered on the sample (source and detector move around the sample).
For more information see Fig. 2 of paper [1].
References
-----------
[1] Aarle, Wim & Palenstijn, Willem & Cant, Jeroen & Janssens, Eline & Bleichrodt,
Folkert & Dabravolski, Andrei & De Beenhouwer, Jan & Batenburg, Kees & Sijbers, Jan. (2016).
Fast and flexible X-ray tomography using the ASTRA toolbox.
Optics Express. 24. 25129-25147. 10.1364/OE.24.025129.
"""
self._init_cuda(cuda_options)
self._init_geometry(
sinos_shape,
source_origin_dist,
angles,
volume_shape,
rot_center,
relative_z_position,
)
self._alg_id = None
self._vol_id = None
self._proj_id = None
def _init_cuda(self, cuda_options):
cuda_options = cuda_options or {}
self.cuda = CudaProcessing(**cuda_options)
def _set_sino_shape(self, sinos_shape):
if len(sinos_shape) != 3:
raise ValueError("Expected a 3D shape")
self.sinos_shape = sinos_shape
self.n_sinos, self.n_angles, self.prj_width = sinos_shape
def _init_geometry(
self, sinos_shape, source_origin_dist, angles, volume_shape, rot_center, relative_z_position
):
self._set_sino_shape(sinos_shape)
if angles is None:
self.angles = np.linspace(0, 2 * np.pi, self.n_angles, endpoint=True)
else:
self.angles = angles
if volume_shape is None:
volume_shape = (self.sinos_shape[0], self.sinos_shape[2], self.sinos_shape[2])
self.volume_shape = volume_shape
self.n_z, self.n_y, self.n_x = self.volume_shape
self.source_origin_dist = source_origin_dist
self.vol_geom = astra.create_vol_geom(self.n_y, self.n_x, self.n_z)
self.vol_shape = astra.geom_size(self.vol_geom)
self._cor_shift = 0.0
self.rot_center = rot_center
if rot_center is not None:
self._cor_shift = (self.prj_width - 1) / 2.0 - rot_center
self._create_astra_proj_geometry(relative_z_position)
def _create_astra_proj_geometry(self, relative_z_position):
# This object has to be re-created each time, because once the modifications below are done,
# it is no more a "cone" geometry but a "cone_vec" geometry, and cannot be updated subsequently
# (see astra/functions.py:271)
self.proj_geom = astra.create_proj_geom(
"cone",
# detector spacing in each dimension.
# Normalized to 1, so source_origin and origin_dest have to be put wrt this unit
1.0,
1.0,
self.n_sinos,
self.prj_width,
self.angles,
self.source_origin_dist,
0.0,
)
self.relative_z_position = relative_z_position or 0.0
self.proj_geom = astra.geom_postalignment(self.proj_geom, (self._cor_shift, self.relative_z_position))
# Component 2 of vecs is the z coordinate of the source, component 5 is the z component of the detector position
# We need to re-create the same inclination of the cone beam, thus we need to keep the inclination of the two z positions.
# The detector is centered on the rotation axis, thus moving it up or down, just moves it out of the reconstruction volume.
# We can bring back the detector in the correct volume position, by applying a rigid translation of both the detector and the source.
# The translation is exactly the amount that brought the detector up or down, but in the opposite direction.
vecs = self.proj_geom["Vectors"]
vecs[:, 2] = -vecs[:, 5]
vecs[:, 5] = 0.0
def _set_output(self, volume):
if volume is not None:
self.cuda.check_array(volume, self.vol_shape)
self.cuda.set_array("output", volume, self.vol_shape)
if volume is None:
self.cuda.allocate_array("output", self.vol_shape)
d_volume = self.cuda.get_array("output")
z, y, x = d_volume.shape
self._vol_link = astra.data3d.GPULink(d_volume.ptr, x, y, z, d_volume.strides[-2])
self._vol_id = astra.data3d.link("-vol", self.vol_geom, self._vol_link)
def _set_input(self, sinos):
self.cuda.check_array(sinos, self.sinos_shape)
self.cuda.set_array("sinos", sinos, sinos.shape) # self.cuda.sinos is now a GPU array
# TODO don't create new link/proj_id if ptr is the same ?
# But it seems Astra modifies the input sinogram while doing FDK, so this might be not relevant
d_sinos = self.cuda.get_array("sinos")
self._proj_data_link = astra.data3d.GPULink(
d_sinos.ptr, self.prj_width, self.n_angles, self.n_z, sinos.strides[-2]
)
self._proj_id = astra.data3d.link("-sino", self.proj_geom, self._proj_data_link)
def _update_reconstruction(self):
cfg = astra.astra_dict("FDK_CUDA")
cfg["ReconstructionDataId"] = self._vol_id
cfg["ProjectionDataId"] = self._proj_id
# TODO more options "eg. filter" ?
if self._alg_id is not None:
astra.algorithm.delete(self._alg_id)
self._alg_id = astra.algorithm.create(cfg)
def reconstruct(self, sinos, output=None, relative_z_position=None):
"""
sinos: numpy.ndarray or pycuda.gpuarray
Sinograms, with shape (n_sinograms, n_angles, width)
output: pycuda.gpuarray, optional
Output array. If not provided, a new numpy array is returned
relative_z_position: int, optional
Position of the central slice of the slab, with respect to the full stack of slices.
By default it is set to zero, meaning that the current slab is assumed in the middle of the stack
"""
self._create_astra_proj_geometry(relative_z_position)
self._set_input(sinos)
self._set_output(output)
self._update_reconstruction()
astra.algorithm.run(self._alg_id)
result = self.cuda.get_array("output")
if output is None:
result = result.get()
self.cuda.recover_arrays_references(["sinos", "output"])
return result
def __del__(self):
if self._alg_id is not None:
astra.algorithm.delete(self._alg_id)
if self._vol_id is not None:
astra.data3d.delete(self._vol_id)
if self._proj_id is not None:
astra.data3d.delete(self._proj_id)
import pytest
import numpy as np
from scipy.ndimage import gaussian_filter
try:
import astra
__has_astra__ = True
except ImportError:
__has_astra__ = False
from nabu.cuda.utils import __has_pycuda__, get_cuda_context
from nabu.testutils import __do_large_mem_tests__
if __has_pycuda__:
from nabu.reconstruction.cone import ConebeamReconstructor
from nabu.utils import subdivide_into_overlapping_segment
@pytest.fixture(scope="class")
def bootstrap(request):
cls = request.cls
cls.vol_shape = (128, 126, 126)
cls.n_angles = 180
cls.prj_width = 192 # detector larger than the sample
cls.src_orig_dist = 1000
cls.volume, cls.cone_data = generate_hollow_cube_cone_sinograms(
cls.vol_shape, cls.n_angles, cls.src_orig_dist, prj_width=cls.prj_width
)
if __has_pycuda__:
cls.ctx = get_cuda_context()
@pytest.mark.skipif(not (__has_pycuda__ and __has_astra__), reason="Need pycuda and astra for this test")
@pytest.mark.usefixtures("bootstrap")
class TestCone:
def _create_cone_reconstructor(self, relative_z_position=None):
return ConebeamReconstructor(
self.cone_data.shape,
self.src_orig_dist,
relative_z_position=relative_z_position,
volume_shape=self.volume.shape,
cuda_options={"ctx": self.ctx},
)
def test_simple_cone_reconstruction(self):
C = self._create_cone_reconstructor()
res = C.reconstruct(self.cone_data)
delta = np.abs(res - self.volume)
# Can we do better ? We already had to lowpass-filter the volume!
# First/last slices are OK
assert np.max(delta[:8]) < 1e-5
assert np.max(delta[-8:]) < 1e-5
# Middle region has a relatively low error
assert np.max(delta[40:-40]) < 0.11
# Transition zones between "zero" and "cube" has a large error
assert np.max(delta[10:25]) < 0.3
assert np.max(delta[-25:-10]) < 0.3
# End of transition zones have a smaller error
np.max(delta[25:40]) < 0.15
np.max(delta[-40:-25]) < 0.15
def test_against_explicit_astra_calls(self):
C = self._create_cone_reconstructor()
res = C.reconstruct(self.cone_data)
#
# Check that ConebeamReconstructor is consistent with these calls to astra
#
angles = np.linspace(0, 2 * np.pi, self.n_angles, True)
# "vol_geom" shape layout is (y, x, z). But here this geometry is used for the reconstruction
# (i.e sinogram -> volume)and not for projection (volume -> sinograms).
# So we assume a square slice. Mind that this is a particular case.
vol_geom = astra.create_vol_geom(self.vol_shape[2], self.vol_shape[2], self.vol_shape[0])
proj_geom = astra.create_proj_geom(
"cone", 1.0, 1.0, self.cone_data.shape[0], self.prj_width, angles, self.src_orig_dist, 0.0
)
sino_id = astra.data3d.create("-sino", proj_geom, data=self.cone_data)
rec_id = astra.data3d.create("-vol", vol_geom)
cfg = astra.astra_dict("FDK_CUDA")
cfg["ReconstructionDataId"] = rec_id
cfg["ProjectionDataId"] = sino_id
alg_id = astra.algorithm.create(cfg)
astra.algorithm.run(alg_id)
res_astra = astra.data3d.get(rec_id)
# housekeeping
astra.algorithm.delete(alg_id)
astra.data3d.delete(rec_id)
astra.data3d.delete(sino_id)
assert (
np.max(np.abs(res - res_astra)) < 5e-4
), "ConebeamReconstructor results are inconsistent with plain calls to astra"
def test_full_vs_partial_cone_geometry(self):
"""
In the ideal case, all the data volume (and reconstruction) fits in memory.
In practice this is rarely the case, so we have to reconstruct the volume slabs by slabs.
The slabs should be slightly overlapping to avoid "stitching" artefacts at the edges.
"""
# Astra seems to duplicate the projection data, even if all GPU memory is handled externally
# Let's try with (n_z * n_y * n_x + 2 * n_a * n_z * n_x) * 4 < mem_limit
# 256^3 seems OK with n_a = 200 (180 MB)
n_z = n_y = n_x = 256
n_a = 200
src_orig_dist = 1000
volume, cone_data = generate_hollow_cube_cone_sinograms(
vol_shape=(n_z, n_y, n_x), n_angles=n_a, src_orig_dist=src_orig_dist
)
C_full = ConebeamReconstructor(cone_data.shape, src_orig_dist, cuda_options={"ctx": self.ctx})
proj_id, projs_full_geom = astra.create_sino3d_gpu(volume, C_full.proj_geom, C_full.vol_geom)
astra.data3d.delete(proj_id)
# Do the same slab-by-slab
inner_slab_size = 64
overlap = 16
slab_size = inner_slab_size + overlap * 2
slabs = subdivide_into_overlapping_segment(n_z, slab_size, overlap)
projs_partial_geom = np.zeros_like(projs_full_geom)
for slab in slabs:
z_min, z_inner_min, z_inner_max, z_max = slab
rel_z_pos = (z_min + z_max) / 2 - n_z / 2
subvolume = volume[z_min:z_max, :, :]
C = ConebeamReconstructor(
(z_max - z_min, n_a, n_x), src_orig_dist, cuda_options={"ctx": self.ctx},
relative_z_position=rel_z_pos
)
proj_id, projs = astra.create_sino3d_gpu(subvolume, C.proj_geom, C.vol_geom)
astra.data3d.delete(proj_id)
projs_partial_geom[z_inner_min:z_inner_max] = projs[z_inner_min - z_min:z_inner_max-z_min]
error_profile = [
np.max(np.abs(proj_partial - proj_full))
for proj_partial, proj_full in zip(projs_partial_geom, projs_full_geom)
]
assert np.all(np.isclose(error_profile, 0.0, atol=1/64)), "Mismatch between full-cone and slab geometries"
def generate_hollow_cube_cone_sinograms(vol_shape, n_angles, src_orig_dist, prj_width=None, apply_filter=True):
# Adapted from Astra toolbox python samples
n_z, n_y, n_x = vol_shape
vol_geom = astra.create_vol_geom(n_y, n_x, n_z)
prj_width = prj_width or n_x
prj_height = n_z
angles = np.linspace(0, 2 * np.pi, n_angles, True)
proj_geom = astra.create_proj_geom("cone", 1.0, 1.0, prj_height, prj_width, angles, src_orig_dist, 0.0)
# hollow cube
cube = np.zeros(astra.geom_size(vol_geom), dtype="f")
d = int(min(n_x, n_y)/2 * (1 - np.sqrt(2) / 2))
cube[20:-20, d:-d, d:-d] = 1
cube[40:-40, d+20:-(d+20), d+20:-(d+20)] = 0
# High-frequencies yield cannot be accurately retrieved
if apply_filter:
cube = gaussian_filter(cube, (1.0, 1.0, 1.0))
proj_id, proj_data = astra.create_sino3d_gpu(cube, proj_geom, vol_geom)
astra.data3d.delete(proj_id) # (n_z, n_angles, n_x)
return cube, proj_data
......@@ -26,6 +26,13 @@ if __do_long_tests__:
except:
__do_long_tests__ = False
__do_large_mem_tests__ = os.environ.get("NABU_LARGE_MEM_TESTS", False)
if __do_large_mem_tests__:
try:
__do_large_mem_tests__ = bool(int(__do_large_mem_tests__))
except:
__do_large_mem_tests__ = False
def generate_tests_scenarios(configurations):
"""
......
......@@ -397,6 +397,45 @@ def recursive_copy_dict(dict_):
return res
def subdivide_into_overlapping_segment(N, window_width, half_overlap):
"""
Divide a segment into a number of possibly-overlapping sub-segments.
Parameters
----------
N: int
Total segment length
window_width: int
Length of each segment
half_overlap: int
Half-length of the overlap between each sub-segment.
Returns
-------
segments: list
A list where each item is in the form
(left_margin_start, inner_segment_start, inner_segment_end, right_margin_end)
"""
if half_overlap >= window_width//2:
raise ValueError("overlap must be smaller than window_width")
w_in = window_width - 2*half_overlap
n_segments = N // w_in
inner_start = w_in * np.arange(n_segments)
inner_end = w_in * (np.arange(n_segments) + 1)
margin_left_start = np.maximum(inner_start - half_overlap, 0)
margin_right_end = np.minimum(inner_end + half_overlap, N)
segments = [
(left_start, i_start, i_end, right_end)
for left_start, i_start, i_end, right_end in zip(margin_left_start.tolist(), inner_start.tolist(), inner_end.tolist(), margin_right_end.tolist())
]
if N % w_in:
# additional sub-segment
new_margin_left_start = inner_end[-1] - half_overlap
new_inner_start = inner_end[-1]
new_inner_end = N
segments.append((new_margin_left_start, new_inner_start, new_inner_end, new_inner_end))
return segments
def get_num_threads(n=None):
"""
......@@ -596,4 +635,4 @@ class DictToObj(object):
"""
def __init__(self, dictio):
self.__dict__ = dictio
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment