[FIX] Fix GPU Tractography

AFQ/api/group.py

@@ -340,8 +340,8 @@ def load_next_subject():
                 this_bundles_file = self.export(
                     "bundles", collapse=False)[sub][ses]
                 this_mapping = self.export("mapping", collapse=False)[sub][ses]
-                this_img = nib.load(self.export(
-                    "dwi", collapse=False)[sub][ses])
+                this_img = self.export(
+                    "dwi", collapse=False)[sub][ses]
                 seg_sft = aus.SegmentedSFT.fromfile(
                     this_bundles_file,
                     this_img)

AFQ/api/participant.py

@@ -98,11 +98,10 @@ def __init__(self,
         self.output_dir = output_dir
 
         self.kwargs = dict(
-            dwi=dwi_data_file,
+            dwi_path=dwi_data_file,
             bval=bval_file,
             bvec=bvec_file,
             results_dir=output_dir,
-            dwi_affine=nib.load(dwi_data_file).affine,
             bids_info=bids_info,
             base_fname=get_base_fname(output_dir, dwi_data_file),
             **kwargs)

AFQ/definitions/image.py

@@ -157,7 +157,7 @@ def get_name(self):
         return name_from_path(self.fname) if self._from_path else self.suffix
 
     def get_image_getter(self, task_name):
-        def image_getter(dwi, bids_info):
+        def _image_getter_helper(dwi, bids_info):
             # Load data
             image_file, image_data_orig, image_affine = \
                 self.get_path_data_affine(bids_info)
@@ -167,15 +167,20 @@ def image_getter(dwi, bids_info):
                 image_data_orig, image_file)
 
             # Resample to DWI data:
-            dwi_img = nib.load(dwi)
             image_data = _resample_image(
                 image_data,
-                dwi_img.get_fdata(),
+                dwi.get_fdata(),
                 image_affine,
-                dwi_img.affine)
+                dwi.affine)
             return nib.Nifti1Image(
                 image_data.astype(np.float32),
-                dwi_img.affine), meta
+                dwi.affine), meta
+        if task_name == "data":
+            def image_getter(dwi, bids_info):
+                return _image_getter_helper(dwi, bids_info)
+        else:
+            def image_getter(data_imap, bids_info):
+                return _image_getter_helper(data_imap["dwi"], bids_info)
         return image_getter
 
 
@@ -198,11 +203,16 @@ def get_name(self):
         return "entire_volume"
 
     def get_image_getter(self, task_name):
-        def image_getter(dwi):
-            dwi_img = nib.load(dwi)
+        def _image_getter_helper(dwi):
             return nib.Nifti1Image(
-                np.ones(dwi_img.get_fdata()[..., 0].shape, dtype=np.float32),
-                dwi_img.affine), dict(source="Entire Volume")
+                np.ones(dwi.get_fdata()[..., 0].shape, dtype=np.float32),
+                dwi.affine), dict(source="Entire Volume")
+        if task_name == "data":
+            def image_getter(dwi):
+                return _image_getter_helper(dwi)
+        else:
+            def image_getter(data_imap):
+                return _image_getter_helper(data_imap["dwi"])
         return image_getter
 
 
@@ -266,7 +276,7 @@ def get_name(self):
         return "roi"
 
     def get_image_getter(self, task_name):
-        def _image_getter_helper(dwi_affine, mapping,
+        def _image_getter_helper(mapping,
                                  data_imap, segmentation_params):
             image_data = None
             bundle_dict = data_imap["bundle_dict"]
@@ -280,7 +290,7 @@ def _image_getter_helper(dwi_affine, mapping,
                 bundle_entry = bundle_dict.transform_rois(
                     bundle_name,
                     mapping,
-                    dwi_affine)
+                    data_imap["dwi_affine"])
                 rois = []
                 if self.use_endpoints:
                     rois.extend(
@@ -316,25 +326,25 @@ def _image_getter_helper(dwi_affine, mapping,
                     f"an ROI Image: {bundle_dict._dict}"))
             return nib.Nifti1Image(
                 image_data.astype(np.float32),
-                dwi_affine), dict(source="ROIs")
+                data_imap["dwi_affine"]), dict(source="ROIs")
 
         if task_name == "data":
             raise ValueError((
                 "RoiImage cannot be used in this context, as they"
                 "require later derivatives to be calculated"))
         elif task_name == "mapping":
             def image_getter(
-                    dwi_affine, mapping,
+                    mapping,
                     data_imap, segmentation_params):
                 return _image_getter_helper(
-                    dwi_affine, mapping,
+                    mapping,
                     data_imap, segmentation_params)
         else:
             def image_getter(
-                    dwi_affine, mapping_imap,
+                    mapping_imap,
                     data_imap, segmentation_params):
                 return _image_getter_helper(
-                    dwi_affine, mapping_imap["mapping"],
+                    mapping_imap["mapping"],
                     data_imap, segmentation_params)
         return image_getter
 

AFQ/segmentation.py

@@ -321,9 +321,9 @@ def _read_tg(self, tg=None):
 
         return tg
 
-    def segment(self, bundle_dict, tg, mapping, fdata=None, fbval=None,
-                fbvec=None, reg_prealign=None,
-                reg_template=None, img_affine=None, reset_tg_space=False):
+    def segment(self, bundle_dict, tg, mapping, img,
+                reg_prealign=None,
+                reg_template=None, reset_tg_space=False):
         """
         Segment streamlines into bundles based on either waypoint ROIs
         [Yeatman2012]_ or RecoBundles [Garyfallidis2017]_.
@@ -341,17 +341,14 @@ def segment(self, bundle_dict, tg, mapping, fdata=None, fbval=None,
             Bundles to segment
         mapping : DiffeomorphicMap, or equivalent interface
             A mapping between DWI space and a template.
-        fdata, fbval, fbvec : str
-            Full path to data, bvals, bvecs
+        img : Nifti1Image
+            Image to use as reference.
         reg_prealign : array, optional.
             The linear transformation to be applied to align input images to
             the reference space before warping under the deformation field.
             Default: None.
         reg_template : str or nib.Nifti1Image, optional.
             Template to use for registration. Default: MNI T2.
-        img_affine : array, optional.
-            The spatial transformation from the measurement to the scanner
-            space.
         reset_tg_space : bool, optional
             Whether to reset the space of the input tractogram after
             segmentation is complete. Default: False.
@@ -371,19 +368,8 @@ def segment(self, bundle_dict, tg, mapping, fdata=None, fbval=None,
         bundles using local and global streamline-based registration and
         clustering, Neuroimage, 2017.
         """
-        if img_affine is not None:
-            if (mapping is None
-                or fdata is not None
-                or fbval is not None
-                    or fbvec is not None):
-
-                self.logger.error(
-                    "Provide either the full path to data, bvals, bvecs,"
-                    + "or provide the affine of the image and the mapping")
-
-        self.logger.info("Preparing Segmentation Parameters")
-        self.img_affine = img_affine
-        self.prepare_img(fdata, fbval, fbvec)
+        self.img_affine = img.affine
+        self.img = img
         self.logger.info("Preprocessing Streamlines")
         tg = self._read_tg(tg)
 
@@ -416,24 +402,6 @@ def segment(self, bundle_dict, tg, mapping, fdata=None, fbval=None,
 
         return fiber_groups
 
-    def prepare_img(self, fdata, fbval, fbvec):
-        """
-        Prepare image data from DWI data.
-        Parameters
-        ----------
-        fdata, fbval, fbvec : str
-            Full path to data, bvals, bvecs
-        """
-        if self.img_affine is None:
-            self.img, _, _, _ = \
-                ut.prepare_data(fdata, fbval, fbvec,
-                                b0_threshold=self.b0_threshold)
-            self.img_affine = self.img.affine
-
-        self.fdata = fdata
-        self.fbval = fbval
-        self.fbvec = fbvec
-
     def cross_streamlines(self, fgarray, template=None):
         """
         Classify the streamlines by whether they cross the midline.
@@ -945,9 +913,7 @@ def segment_reco(self, tg=None):
                 self.presegment_bundle_dict,
                 self.tg,
                 self.mapping,
-                self.fdata,
-                self.fbval,
-                self.fbvec,
+                self.img,
                 reg_template=self.reg_template,
                 reg_prealign=self.reg_prealign)
             roiseg_fg = roiseg.fiber_groups

AFQ/tasks/data.py

@@ -1,5 +1,7 @@
 import nibabel as nib
+import nibabel.orientations as nio
 import numpy as np
+import logging
 
 from dipy.io.gradients import read_bvals_bvecs
 import dipy.core.gradients as dpg
@@ -33,16 +35,20 @@
     extract_odf, anisotropic_index, anisotropic_power)
 
 
+logger = logging.getLogger('AFQ')
+
+
 DIPY_GH = "https://github.com/dipy/dipy/blob/master/dipy/"
 
 
-@pimms.calc("data", "gtab", "img")
-def get_data_gtab(dwi, bval, bvec, min_bval=None,
+@pimms.calc("data", "gtab", "dwi", "dwi_affine")
+def get_data_gtab(dwi_path, bval, bvec, min_bval=None,
                   max_bval=None, filter_b=True, b0_threshold=50):
     """
     DWI data as an ndarray for selected b values,
     A DIPY GradientTable with all the gradient information,
-    and unaltered DWI data in a Nifti1Image.
+    and DWI data in a Nifti1Image,
+    and the affine transformation of the DWI data.
 
     Parameters
     ----------
@@ -61,9 +67,10 @@ def get_data_gtab(dwi, bval, bvec, min_bval=None,
         The value of b under which
         it is considered to be b0. Default: 50.
     """
-    img = nib.load(dwi)
-    data = img.get_fdata()
+    img = nib.load(dwi_path)
     bvals, bvecs = read_bvals_bvecs(bval, bvec)
+
+    data = img.get_fdata()
     if filter_b and (min_bval is not None):
         valid_b = np.logical_or(
             (bvals >= min_bval), (bvals <= b0_threshold))
@@ -79,25 +86,27 @@ def get_data_gtab(dwi, bval, bvec, min_bval=None,
     gtab = dpg.gradient_table(
         bvals, bvecs,
         b0_threshold=b0_threshold)
-    return data, gtab, img
+    img = nib.Nifti1Image(data, img.affine)
+    return data, gtab, img, img.affine
 
 
 @pimms.calc("b0")
 @as_file('_desc-b0_dwi.nii.gz')
-def b0(dwi, data, gtab, img):
+def b0(dwi_path, gtab):
     """
     full path to a nifti file containing the mean b0
     """
-    mean_b0 = np.mean(data[..., gtab.b0s_mask], -1)
-    mean_b0_img = nib.Nifti1Image(mean_b0, img.affine)
+    data = nib.load(dwi_path)
+    mean_b0 = np.mean(data.get_fdata()[..., gtab.b0s_mask], -1)
+    mean_b0_img = nib.Nifti1Image(mean_b0, data.affine)
     meta = dict(b0_threshold=gtab.b0_threshold,
-                source=dwi)
+                source=dwi_path)
     return mean_b0_img, meta
 
 
 @pimms.calc("masked_b0")
 @as_file('_desc-maskedb0_dwi.nii.gz')
-def b0_mask(base_fname, b0, brain_mask):
+def b0_mask(b0, brain_mask):
     """
     full path to a nifti file containing the
     mean b0 after applying the brain mask
@@ -970,7 +979,7 @@ def get_data_plan(kwargs):
         if kwargs["bids_info"] is not None:
             bm_def.find_path(
                 kwargs["bids_info"]["bids_layout"],
-                kwargs["dwi"],
+                kwargs["dwi_path"],
                 kwargs["bids_info"]["subject"],
                 kwargs["bids_info"]["session"])
         data_tasks["brain_mask_res"] = pimms.calc("brain_mask")(