Merge branch 'master' into 'option_near_tafforensis'

# Conflicts:
#   nabu/__init__.py

nabu/cuda/processing.py

@@ -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")

nabu/reconstruction/cone.py

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

nabu/reconstruction/tests/test_cone.py

+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

nabu/testutils.py

@@ -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):
     """

nabu/utils.py

@@ -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
-    
+