Merge pull request #420 from dPys/development

Development

Author: dPys

pynets/core/interfaces.py

@@ -574,6 +574,7 @@ def _run_interface(self, runtime):
             template_tmp_path,
             self.inputs.simple,
         )
+        time.sleep(0.5)
 
         if self.inputs.mask:
             out_name_mask = fname_presuffix(
@@ -1143,11 +1144,12 @@ class RegisterDWI(SimpleInterface):
 
     def _run_interface(self, runtime):
         import gc
+        import time
         import glob
         import os.path as op
         from pynets.registration import register
-        from pynets.registration import reg_utils as regutils
         from nipype.utils.filemanip import fname_presuffix, copyfile
+        from pynets.registration.reg_utils import check_orient_and_dims
 
         fa_tmp_path = fname_presuffix(
             self.inputs.fa_path, suffix="_tmp", newpath=runtime.cwd
@@ -1205,6 +1207,9 @@ def _run_interface(self, runtime):
                     mask_tmp_path,
                     copy=True,
                     use_hardlink=False)
+                mask_tmp_path = check_orient_and_dims(
+                    mask_tmp_path, runtime.cwd, self.inputs.vox_size
+                )
             else:
                 mask_tmp_path = None
 
@@ -1218,6 +1223,9 @@ def _run_interface(self, runtime):
                 gm_mask,
                 copy=True,
                 use_hardlink=False)
+            gm_mask = check_orient_and_dims(
+                gm_mask, runtime.cwd, self.inputs.vox_size
+            )
         else:
             gm_mask = None
 
@@ -1231,6 +1239,9 @@ def _run_interface(self, runtime):
                 wm_mask,
                 copy=True,
                 use_hardlink=False)
+            wm_mask = check_orient_and_dims(
+                wm_mask, runtime.cwd, self.inputs.vox_size
+            )
         else:
             wm_mask = None
 
@@ -1245,6 +1256,9 @@ def _run_interface(self, runtime):
                 csf_mask,
                 copy=True,
                 use_hardlink=False)
+            csf_mask = check_orient_and_dims(
+                csf_mask, runtime.cwd, self.inputs.vox_size
+            )
         else:
             csf_mask = None
 
@@ -1270,9 +1284,11 @@ def _run_interface(self, runtime):
 
         # Generate T1w brain mask
         reg.gen_mask(mask_tmp_path)
+        time.sleep(0.5)
 
         # Perform anatomical segmentation
         reg.gen_tissue(wm_mask, gm_mask, csf_mask, self.inputs.overwrite)
+        time.sleep(0.5)
 
         # Align t1w to mni template
         # from joblib import Memory
@@ -1284,17 +1300,20 @@ def _run_interface(self, runtime):
         # t1w2mni_align = memory.cache(reg.t1w2mni_align)
         # t1w2mni_align()
         reg.t1w2mni_align()
+        time.sleep(0.5)
 
         if (self.inputs.overwrite is True) or (
                 op.isfile(reg.t1w2dwi) is False):
             # Align t1w to dwi
             reg.t1w2dwi_align()
+            time.sleep(0.5)
 
         if (self.inputs.overwrite is True) or (
             op.isfile(reg.wm_gm_int_in_dwi) is False
         ):
             # Align tissue
             reg.tissue2dwi_align()
+            time.sleep(0.5)
 
         self._results["wm_in_dwi"] = reg.wm_in_dwi
         self._results["gm_in_dwi"] = reg.gm_in_dwi
@@ -1389,6 +1408,7 @@ class RegisterAtlasDWI(SimpleInterface):
 
     def _run_interface(self, runtime):
         import gc
+        import time
         import os
         from pynets.registration import reg_utils as regutils
         from pynets.core.nodemaker import \
@@ -1561,6 +1581,9 @@ def _run_interface(self, runtime):
         base_dir_tmp = f"{runtime.cwd}/atlas_{atlas_name}"
         os.makedirs(base_dir_tmp, exist_ok=True)
 
+        mni2dwi_xfm = f"{base_dir_tmp}{'/'}{atlas_name}" \
+                              f"{'_mni2dwi_xfm.mat'}"
+
         aligned_atlas_t1mni = f"{base_dir_tmp}{'/'}{atlas_name}" \
                               f"{'_t1w_mni.nii.gz'}"
         aligned_atlas_skull = f"{base_dir_tmp}{'/'}{atlas_name}" \
@@ -1597,6 +1620,7 @@ def _run_interface(self, runtime):
             dwi_aligned_atlas,
             dwi_aligned_atlas_wmgm_int,
             B0_mask_tmp_path,
+            mni2dwi_xfm,
             self.inputs.simple,
         )
 
@@ -1637,6 +1661,7 @@ def _run_interface(self, runtime):
                 template_tmp_path,
                 self.inputs.simple,
             )
+            time.sleep(0.5)
             os.remove(waymask_tmp_path)
         else:
             waymask_in_dwi = None
@@ -1739,6 +1764,7 @@ class RegisterROIDWI(SimpleInterface):
     def _run_interface(self, runtime):
         import gc
         import os
+        import time
         from pynets.registration import reg_utils as regutils
         from nipype.utils.filemanip import fname_presuffix, copyfile
         import pkg_resources
@@ -1852,6 +1878,7 @@ def _run_interface(self, runtime):
                 template_tmp_path,
                 self.inputs.simple,
             )
+            time.sleep(0.5)
         else:
             roi_in_dwi = None
 
@@ -1910,10 +1937,11 @@ class RegisterFunc(SimpleInterface):
     def _run_interface(self, runtime):
         import gc
         import glob
+        import time
         import os.path as op
         from pynets.registration import register
-        from pynets.registration import reg_utils as regutils
         from nipype.utils.filemanip import fname_presuffix, copyfile
+        from pynets.registration.reg_utils import check_orient_and_dims
 
         anat_mask_existing = [
             i
@@ -1945,6 +1973,9 @@ def _run_interface(self, runtime):
                     mask_tmp_path,
                     copy=True,
                     use_hardlink=False)
+                mask_tmp_path = check_orient_and_dims(
+                    mask_tmp_path, runtime.cwd, self.inputs.vox_size
+                )
             else:
                 mask_tmp_path = None
 
@@ -1959,6 +1990,9 @@ def _run_interface(self, runtime):
                 gm_mask,
                 copy=True,
                 use_hardlink=False)
+            gm_mask = check_orient_and_dims(
+                gm_mask, runtime.cwd, self.inputs.vox_size
+            )
         else:
             gm_mask = None
 
@@ -1973,6 +2007,9 @@ def _run_interface(self, runtime):
                 wm_mask,
                 copy=True,
                 use_hardlink=False)
+            wm_mask = check_orient_and_dims(
+                wm_mask, runtime.cwd, self.inputs.vox_size
+            )
         else:
             wm_mask = None
 
@@ -1995,9 +2032,11 @@ def _run_interface(self, runtime):
 
         # Generate T1w brain mask
         reg.gen_mask(mask_tmp_path)
+        time.sleep(0.5)
 
         # Perform anatomical segmentation
         reg.gen_tissue(wm_mask, gm_mask, self.inputs.overwrite)
+        time.sleep(0.5)
 
         # Align t1w to mni template
         # from joblib import Memory
@@ -2009,6 +2048,7 @@ def _run_interface(self, runtime):
         # t1w2mni_align = memory.cache(reg.t1w2mni_align)
         # t1w2mni_align()
         reg.t1w2mni_align()
+        time.sleep(0.5)
 
         self._results["reg_fmri_complete"] = True
         self._results["basedir_path"] = runtime.cwd
@@ -2055,6 +2095,7 @@ def _run_interface(self, runtime):
         import gc
         import os
         import pkg_resources
+        import time
         from pynets.core.utils import prune_suffices
         from pynets.registration import reg_utils as regutils
         from nipype.utils.filemanip import fname_presuffix, copyfile
@@ -2146,6 +2187,7 @@ def _run_interface(self, runtime):
             t1w2mni_warp_tmp_path,
             self.inputs.simple
         )
+        time.sleep(0.5)
 
         out_dir = f"{self.inputs.dir_path}/t1w_clustered_parcellations/"
         os.makedirs(out_dir, exist_ok=True)
@@ -2216,6 +2258,7 @@ class RegisterAtlasFunc(SimpleInterface):
     def _run_interface(self, runtime):
         import gc
         import os
+        import time
         import glob
         from pynets.registration import reg_utils as regutils
         from pynets.core.nodemaker import \
@@ -2355,6 +2398,7 @@ def _run_interface(self, runtime):
                 aligned_atlas_gm,
                 self.inputs.simple,
             )
+            time.sleep(0.5)
 
             # Correct coords and labels
             [aligned_atlas_gm, coords, labels] = \
@@ -2447,6 +2491,7 @@ class RegisterROIEPI(SimpleInterface):
     def _run_interface(self, runtime):
         import gc
         import os
+        import time
         from pynets.registration import reg_utils as regutils
         from nipype.utils.filemanip import fname_presuffix, copyfile
         import pkg_resources
@@ -2516,6 +2561,7 @@ def _run_interface(self, runtime):
                 template_tmp_path,
                 self.inputs.simple,
             )
+            time.sleep(0.5)
         else:
             roi_in_t1w = None
 
@@ -2688,7 +2734,6 @@ def _run_interface(self, runtime):
             B0_mask_tmp_path,
             copy=True,
             use_hardlink=False)
-        B0_mask_data = nib.load(B0_mask_tmp_path, mmap=True).get_fdata()
 
         # Fit diffusion model
         # Save reconstruction to .npy
@@ -2744,6 +2789,7 @@ def _run_interface(self, runtime):
                 copy=True,
                 use_hardlink=False,
             )
+            del model
         else:
             print(
                 f"Found existing reconstruction with "
@@ -2754,7 +2800,6 @@ def _run_interface(self, runtime):
                 copy=True,
                 use_hardlink=False,
             )
-            model = np.load(recon_path, mmap_mode='r+').astype('float32')
 
         dwi_img.uncache()
         del dwi_data
@@ -2780,12 +2825,6 @@ def _run_interface(self, runtime):
             copy=True,
             use_hardlink=False)
 
-        atlas_img = nib.load(labels_im_file_tmp_path, mmap=True)
-        atlas_data = np.array(atlas_img.dataobj).astype("uint16")
-        atlas_data_wm_gm_int = np.asarray(
-            nib.load(labels_im_file_tmp_path_wm_gm_int, mmap=True).dataobj
-        ).astype("uint16")
-
         t1w2dwi_tmp_path = fname_presuffix(
             self.inputs.t1w2dwi, suffix="_tmp",
             newpath=runtime.cwd
@@ -2836,22 +2875,8 @@ def _run_interface(self, runtime):
                 waymask_tmp_path,
                 copy=True,
                 use_hardlink=False)
-            waymask_data = np.asarray(nib.load(waymask_tmp_path,
-                                               mmap=True).dataobj
-                                      ).astype("bool")
         else:
-            waymask_data = None
-
-        # Build mask vector from atlas for later roi filtering
-        parcels = []
-        i = 0
-        intensities = [i for i in np.unique(atlas_data) if i != 0]
-        for roi_val in intensities:
-            parcels.append(atlas_data == roi_val)
-            i += 1
-
-        del atlas_data
-        gc.collect()
+            waymask_tmp_path = None
 
         # Iteratively build a list of streamlines for each ROI while tracking
         print(
@@ -2889,9 +2914,9 @@ def _run_interface(self, runtime):
         # Commence Ensemble Tractography
         streamlines = track_ensemble(
             self.inputs.target_samples,
-            atlas_data_wm_gm_int,
-            parcels,
-            model,
+            labels_im_file_tmp_path_wm_gm_int,
+            labels_im_file_tmp_path,
+            recon_path,
             get_sphere(sphere),
             self.inputs.directget,
             self.inputs.curv_thr_list,
@@ -2900,14 +2925,14 @@ def _run_interface(self, runtime):
             self.inputs.maxcrossing,
             int(roi_neighborhood_tol),
             self.inputs.min_length,
-            waymask_data,
-            B0_mask_data,
+            waymask_tmp_path,
+            B0_mask_tmp_path,
             t1w2dwi_tmp_path, gm_in_dwi_tmp_path,
             vent_csf_in_dwi_tmp_path, wm_in_dwi_tmp_path,
-            self.inputs.tiss_class, B0_mask_tmp_path
+            self.inputs.tiss_class,
+            runtime.cwd
         )
 
-        del model, parcels, atlas_data_wm_gm_int
         gc.collect()
 
         # Save streamlines to trk
@@ -2951,15 +2976,16 @@ def _run_interface(self, runtime):
         if use_life is True:
             print('Using LiFE to evaluate streamline plausibility...')
             from pynets.dmri.dmri_utils import evaluate_streamline_plausibility
-            dwi_img = nib.load(dwi_file_tmp_path, mmap=True)
+            dwi_img = nib.load(dwi_file_tmp_path)
             dwi_data = dwi_img.get_fdata().astype('float32')
             orig_count = len(streamlines)
 
             if self.inputs.waymask:
-                mask_data = waymask_data
+                mask_data = nib.load(waymask_tmp_path
+                                     ).get_fdata().astype('bool').astype('int')
             else:
                 mask_data = nib.load(wm_in_dwi_tmp_path
-                                     ).get_fdata().astype('float32')
+                                     ).get_fdata().astype('bool').astype('int')
             try:
                 streamlines = evaluate_streamline_plausibility(
                     dwi_data, gtab, mask_data, streamlines,
@@ -2971,9 +2997,7 @@ def _run_interface(self, runtime):
             if len(streamlines) < 0.5*orig_count:
                 raise ValueError('LiFE revealed no plausible streamlines in '
                                  'the tractogram!')
-            del dwi_data
-
-        del B0_mask_data, waymask_data
+            del dwi_data, mask_data
 
         stf = StatefulTractogram(
             streamlines,