nabu 2024.1.9__py3-none-any.whl → 2024.2.0__py3-none-any.whl

nabu/reconstruction/cone.py

@@ -1,5 +1,10 @@
+from math import sqrt
 import numpy as np
-from ..utils import calc_padding_lengths1D, nextpow2
+
+from ..cuda.kernel import CudaKernel
+from ..cuda.processing import CudaProcessing
+from ..reconstruction.filtering_cuda import CudaSinoFilter
+from ..utils import get_cuda_srcfile
 
 try:
     import astra
@@ -7,8 +12,6 @@ try:
     __have_astra__ = True
 except ImportError:
     __have_astra__ = False
-from ..cuda.processing import CudaProcessing
-from ..processing.padding_cuda import CudaPadding
 
 
 class ConebeamReconstructor:
@@ -16,6 +19,16 @@ class ConebeamReconstructor:
     A reconstructor for cone-beam geometry using the astra toolbox.
     """
 
+    default_extra_options = {
+        "axis_correction": None,
+        "clip_outer_circle": False,
+        "scale_factor": None,
+        "filter_cutoff": 1.0,
+        "outer_circle_value": 0.0,
+        # "use_astra_fdk": True,
+        "use_astra_fdk": False,
+    }
+
     def __init__(
         self,
         sinos_shape,
@@ -25,12 +38,12 @@ class ConebeamReconstructor:
         volume_shape=None,
         rot_center=None,
         relative_z_position=None,
-        axis_correction=None,
         pixel_size=None,
-        scale_factor=None,
         padding_mode="zeros",
+        filter_name=None,
         slice_roi=None,
         cuda_options=None,
+        extra_options=None,
     ):
         """
         Initialize a cone beam reconstructor. This reconstructor works on slabs of data,
@@ -78,6 +91,10 @@ class ConebeamReconstructor:
             This parameter must be in the form (start_x, end_x, start_y, end_y) with no negative values.
             Note that the current implementation just crops the final reconstructed volume,
             i.e there is no speed or memory benefit.
+        use_astra_fdk: bool
+            Whether to use the native Astra Toolbox FDK implementation.
+            If set to False, the cone-beam pre-weighting and projections padding/filtering is done by nabu.
+            Note that this parameter is automatically set to False if padding_mode != "zeros".
 
         Notes
         ------
@@ -106,10 +123,10 @@ class ConebeamReconstructor:
         Fast and flexible X-ray tomography using the ASTRA toolbox.
         Optics Express. 24. 25129-25147. 10.1364/OE.24.025129.
         """
+        self._configure_extra_options(extra_options)
         self._init_cuda(cuda_options)
-        self.scale_factor = scale_factor
         self._set_sino_shape(sinos_shape)
-        self._init_padding(padding_mode)
+        self._orig_prog_geom = None
         self._init_geometry(
             source_origin_dist,
             origin_detector_dist,
@@ -118,13 +135,17 @@ class ConebeamReconstructor:
             volume_shape,
             rot_center,
             relative_z_position,
-            axis_correction,
             slice_roi,
         )
+        self._init_fdk(padding_mode, filter_name)
         self._alg_id = None
         self._vol_id = None
         self._proj_id = None
 
+    def _configure_extra_options(self, extra_options):
+        self.extra_options = self.default_extra_options.copy()
+        self.extra_options.update(extra_options or {})
+
     def _init_cuda(self, cuda_options):
         cuda_options = cuda_options or {}
         self.cuda = CudaProcessing(**cuda_options)
@@ -135,19 +156,25 @@ class ConebeamReconstructor:
         self.sinos_shape = sinos_shape
         self.n_sinos, self.n_angles, self.prj_width = sinos_shape
 
-    def _init_padding(self, padding_mode):
-        self._pad_data = False
+    def _init_fdk(self, padding_mode, filter_name):
         self.padding_mode = padding_mode
-        if padding_mode == "zeros":
+        self._use_astra_fdk = bool(self.extra_options.get("use_astra_fdk", True))
+        self._use_astra_fdk &= padding_mode in ["zeros", "constant", None, "none"]
+        if self._use_astra_fdk:
             return
-        self._pad_data = True
-        n_x = self.prj_width
-        x_pad_lens = calc_padding_lengths1D(n_x, nextpow2(n_x * 2))
-        self.padder = CudaPadding(
-            (self.n_angles, n_x), ((0, 0),) + (x_pad_lens,), mode=padding_mode, cuda_options={"ctx": self.cuda.ctx}
+        self.sino_filter = CudaSinoFilter(
+            self.sinos_shape[1:],
+            filter_name=filter_name,
+            padding_mode=self.padding_mode,
+            # TODO (?) configure FFT backend
+            extra_options={"cutoff": self.extra_options.get("filter_cutoff", 1.0)},
+            cuda_options={"ctx": self.cuda.ctx},
         )
-        self._sinos_padded_shape = (self.n_sinos, self.n_angles, self.padder.padded_shape[-1])
-        self.prj_width = self.padder.padded_shape[-1]  # will impact translations
+        # In astra, FDK pre-weighting does the "n_a/(pi/2) multiplication"
+        # TODO not sure where this "magnification **2" factor comes from ?
+        mult_factor = self.n_angles / 3.141592 * 2 / (self.magnification**2)
+        self.sino_filter.set_filter(self.sino_filter.filter_f * mult_factor, normalize=False)
+        #
 
     def _set_pixel_size(self, pixel_size):
         if pixel_size is None:
@@ -173,12 +200,6 @@ class ConebeamReconstructor:
             self._vol_geom_n_x = self.n_x - start_x * 2
             self._vol_geom_n_y = self.n_y - start_y * 2
         else:
-            # self._crop_data = True
-            # self._output_cropped_shape = (
-            #     self.n_z,
-            #     np.arange(self.n_y)[start_x:end_x].size,
-            #     np.arange(self.n_x)[start_y:end_y].size,
-            # )
             raise NotImplementedError(
                 "Cone-beam geometry supports only slice_roi centered around origin (got slice_roi=%s with n_x=%d, n_y=%d)"
                 % (str(slice_roi), self.n_x, self.n_y)
@@ -193,7 +214,6 @@ class ConebeamReconstructor:
         volume_shape,
         rot_center,
         relative_z_position,
-        axis_correction,
         slice_roi,
     ):
         if angles is None:
@@ -215,7 +235,7 @@ class ConebeamReconstructor:
         if rot_center is not None:
             self._cor_shift = (self.sinos_shape[-1] - 1) / 2.0 - rot_center
         self._set_pixel_size(pixel_size)
-        self._axis_corrections = axis_correction
+        self._axis_corrections = self.extra_options.get("axis_correction", None)
         self._create_astra_proj_geometry(relative_z_position)
 
     def _create_astra_proj_geometry(self, relative_z_position):
@@ -234,6 +254,8 @@ class ConebeamReconstructor:
         )
         self.relative_z_position = relative_z_position or 0.0
         # This will turn the geometry of type "cone" into a geometry of type "cone_vec"
+        if self._orig_prog_geom is None:
+            self._orig_prog_geom = self.proj_geom
         self.proj_geom = astra.geom_postalignment(self.proj_geom, (self._cor_shift, 0))
         # (src, detector_center, u, v) = (srcX, srcY, srcZ, dX, dY, dZ, uX, uY, uZ, vX, vY, vZ)
         vecs = self.proj_geom["Vectors"]
@@ -274,23 +296,29 @@ class ConebeamReconstructor:
         # But it seems Astra modifies the input sinogram while doing FDK, so this might be not relevant
         d_sinos = self.cuda.get_array("sinos")
 
-        if self._pad_data:
-            sinos_padded = self.cuda.allocate_array("sinos_padded", self._sinos_padded_shape, dtype="f")
-            for i in range(self.n_sinos):
-                self.padder.pad(self.cuda.sinos[i], output=sinos_padded[i])
-            d_sinos = sinos_padded
-
         # self._proj_data_link = astra.data3d.GPULink(d_sinos.ptr, self.prj_width, self.n_angles, self.n_z, sinos.strides[-2])
         self._proj_data_link = astra.data3d.GPULink(
             d_sinos.ptr, self.prj_width, self.n_angles, self.n_sinos, d_sinos.strides[-2]
         )
         self._proj_id = astra.data3d.link("-sino", self.proj_geom, self._proj_data_link)
 
+    def _preprocess_data(self):
+        if self._use_astra_fdk:
+            return
+        d_sinos = self.cuda.sinos
+        fdk_preweighting(
+            d_sinos, self._orig_prog_geom, relative_z_position=self.relative_z_position, cor_shift=self._cor_shift
+        )
+        for i in range(d_sinos.shape[0]):
+            self.sino_filter.filter_sino(d_sinos[i], output=d_sinos[i])
+
     def _update_reconstruction(self):
-        cfg = astra.astra_dict("FDK_CUDA")
+        if self._use_astra_fdk:
+            cfg = astra.astra_dict("FDK_CUDA")
+        else:
+            cfg = astra.astra_dict("BP3D_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)
@@ -308,17 +336,20 @@ class ConebeamReconstructor:
         self._create_astra_proj_geometry(relative_z_position)
         self._set_input(sinos)
         self._set_output(output)
+        self._preprocess_data()
         self._update_reconstruction()
         astra.algorithm.run(self._alg_id)
+        #
+        # NB: Could also be done with
+        # from astra.experimental import direct_BP3D
+        # projector_id = astra.create_projector("cuda3d", self.proj_geom, self.vol_geom, options=None)
+        # direct_BP3D(projector_id, self._vol_link, self._proj_data_link)
+        #
         result = self.cuda.get_array("output")
         if output is None:
             result = result.get()
-        if self.scale_factor is not None:
-            result *= np.float32(self.scale_factor)  # in-place for pycuda
-        # if self._crop_data:
-        #     self.cuda.allocate_array("output_cropped", self._output_cropped_shape, dtype=np.float32)
-        #     for i in range(self.n_z):
-        #         output
+        if self.extra_options.get("scale_factor", None) is not None:
+            result *= np.float32(self.extra_options["scale_factor"])  # in-place for pycuda
         self.cuda.recover_arrays_references(["sinos", "output"])
         return result
 
@@ -345,3 +376,36 @@ def roi_is_centered(shape, slice_):
     Return True if "slice_" define a selection that is centered on the middle of the array.
     """
     return all([selection_is_centered(shp, s.start, s.stop) for shp, s in zip(shape, slice_)])
+
+
+def fdk_preweighting(d_sinos, proj_geom, relative_z_position=0.0, cor_shift=0.0):
+
+    preweight_kernel = CudaKernel(
+        "devFDK_preweight",
+        filename=get_cuda_srcfile("cone.cu"),
+        signature="Piiifffffiii",
+    )
+
+    # n_angles, n_z, n_x = d_sinos.shape
+    n_z, n_angles, n_x = d_sinos.shape
+    det_origin = sqrt(proj_geom["DistanceOriginDetector"] ** 2 + cor_shift**2)
+
+    block = (32, 16, 1)
+    grid = (((n_x + 32 - 1) // 32) * ((n_z + 32 - 1) // 32), (n_angles + 16 - 1) // 16, 1)
+
+    preweight_kernel(
+        d_sinos,
+        np.uint32(n_x),  # unsigned int projPitch,
+        np.uint32(0),  # unsigned int startAngle,
+        np.uint32(n_angles),  # unsigned int endAngle,
+        np.float32(proj_geom["DistanceOriginSource"]),  # float fSrcOrigin,
+        np.float32(det_origin),  # float fDetOrigin,
+        np.float32(relative_z_position),  # float fZShift,
+        np.float32(proj_geom["DetectorSpacingX"]),  # float fDetUSize,
+        np.float32(proj_geom["DetectorSpacingY"]),  # float fDetVSize,
+        np.int32(n_angles),  # dims.iProjAngles;
+        np.int32(n_x),  # dims.iProjU; // number of detectors in the U direction
+        np.int32(n_z),  # dims.iProjV // number of detectors in the V direction
+        block=block,
+        grid=grid,
+    )

nabu/reconstruction/fbp.py

@@ -73,12 +73,15 @@ class CudaBackprojector(BackprojectorBase):
             self._kernel_options["kernel_name"],
             filename=self._kernel_options["file_name"],
             options=self._kernel_options["sourcemodule_options"],
+            silent_compilation_warnings=True,  # textures and Cuda 11
         )
         if self.halftomo and self.rot_center < self.dwidth:
             self.sino_mult = CudaSinoMult(self.sino_shape, self.rot_center, ctx=self._processing.ctx)
         self._prepare_textures()  # has to be done after compilation for Cuda (to bind texture to built kernel)
 
     def _transfer_to_texture(self, sino, do_checks=True):
+        if do_checks and not (sino.flags.c_contiguous):
+            raise ValueError("Expected C-Contiguous array")
         if self._use_textures:
             copy_array(self._d_sino_cua, sino, check=do_checks)
         else:

nabu/reconstruction/fbp_base.py

@@ -21,6 +21,7 @@ class BackprojectorBase:
         "clip_outer_circle": False,
         "scale_factor": None,
         "filter_cutoff": 1.0,
+        "outer_circle_value": 0.0,
     }
 
     kernel_filename = None
@@ -335,13 +336,17 @@ class BackprojectorBase:
         self._output_is_ndarray = isinstance(output, np.ndarray)
         if output is None or self._output_is_ndarray:
             self._processing.allocate_array("_d_slice", self.slice_shape, dtype=np.float32)
-            return self._processing._d_slice  # pylint: disable=E1101
-        if check:
+            output = self._processing._d_slice  # pylint: disable=E1101
+        elif check:
             assert output.dtype == np.float32
             assert output.shape == self.slice_shape, "Expected output shape %s but got %s" % (
                 self.slice_shape,
                 output.shape,
             )
+        if self.extra_options.get("clip_outer_circle", False):
+            out_circle_val = self.extra_options.get("outer_circle_value", 0)
+            if out_circle_val != 0:
+                output.fill(out_circle_val)
         return output
 
     def _set_kernel_slice_arg(self, d_slice):

nabu/reconstruction/filtering.py

@@ -162,9 +162,10 @@ class SinoFilter:
             in self.d_sino_padded.
         """
         self._check_array(sino)
-        sino_padded = np.pad(
-            sino, ((0, 0), (0, self.dwidth_padded - self.dwidth)), mode=self.padding_mode
-        )  # pad with a FFT-friendly layout
+        # sino_padded = np.pad(
+        #     sino, ((0, 0), (0, self.dwidth_padded - self.dwidth)), mode=self.padding_mode
+        # )  # pad with a FFT-friendly layout
+        sino_padded = np.pad(sino, ((0, 0), (self.pad_left, self.pad_right)), mode=self.padding_mode)
         sino_padded_f = rfft(sino_padded, axis=1, workers=get_num_threads(self.extra_options["fft_threads"]))
         sino_padded_f *= self.filter_f
         sino_filtered = irfft(sino_padded_f, axis=1, workers=get_num_threads(self.extra_options["fft_threads"]))
@@ -173,9 +174,11 @@ class SinoFilter:
         else:
             res = output
         if self.ndim == 2:
-            res[:] = sino_filtered[:, : self.dwidth]  # pylint: disable=E1126 # ?!
+            # res[:] = sino_filtered[:, : self.dwidth]  # pylint: disable=E1126 # ?!
+            res[:] = sino_filtered[:, self.pad_left : -self.pad_right]  # pylint: disable=E1126 # ?!
         else:
-            res[:] = sino_filtered[:, :, : self.dwidth]  # pylint: disable=E1126 # ?!
+            # res[:] = sino_filtered[:, :, : self.dwidth]  # pylint: disable=E1126 # ?!
+            res[:] = sino_filtered[:, :, self.pad_left : -self.pad_right]  # pylint: disable=E1126 # ?!
         return res
 
     __call__ = filter_sino
@@ -209,12 +212,22 @@ def filter_sinogram(
         frequency cutoff for filter
     """
     n_angles, width = sinogram.shape
-    sinogram_padded = np.pad(sinogram, ((0, 0), (0, padded_width - width)), mode=padding_mode, **padding_kwargs)
+
+    # Initially, padding was done this way
+    # sinogram_padded = np.pad(sinogram, ((0, 0), (0, padded_width - width)), mode=padding_mode, **padding_kwargs)
+
+    #
+    pad_left = (padded_width - width) // 2
+    pad_right = padded_width - width - pad_left
+    sinogram_padded = np.pad(sinogram, ((0, 0), (pad_left, pad_right)), mode=padding_mode, **padding_kwargs)
+    #
+
     fourier_filter = compute_fourier_filter(padded_width, filter_name, cutoff=filter_cutoff)
     if normalize:
         fourier_filter *= np.pi / n_angles
     fourier_filter = fourier_filter[: padded_width // 2 + 1]  # R2C
     sino_f = rfft(sinogram_padded, axis=1)
     sino_f *= fourier_filter
-    sino_filtered = irfft(sino_f, axis=1)[:, :width]  # pylint: disable=E1126 # ?!
+    # sino_filtered = irfft(sino_f, axis=1)[:, :width]  # pylint: disable=E1126 # ?!
+    sino_filtered = irfft(sino_f, axis=1)[:, pad_left:-pad_right]  # pylint: disable=E1126 # ?!
     return sino_filtered

nabu/reconstruction/filtering_cuda.py

@@ -7,7 +7,7 @@ from .filtering import SinoFilter
 
 
 class CudaSinoFilter(SinoFilter):
-    default_extra_options = {**SinoFilter.default_extra_options, **{"fft_backend": "skcuda"}}
+    default_extra_options = {**SinoFilter.default_extra_options, **{"fft_backend": "vkfft"}}
 
     def __init__(
         self,
@@ -76,6 +76,12 @@ class CudaSinoFilter(SinoFilter):
         self._check_array(sino)
         if not (isinstance(sino, self.cuda.array_class)):
             sino = self.cuda.set_array("sino", sino)
+        elif not (sino.flags.c_contiguous):
+            # Transfer the device array into another, c-contiguous, device array
+            # We can throw an error as well in this case, but often we so something like fbp(radios[:, i, :])
+            sino_tmp = self.cuda.allocate_array("sino_contig", sino.shape)
+            sino_tmp.set(sino)
+            sino = sino_tmp
 
         # Padding
         self.padding_kernel(sino, output=self.d_sino_padded)