Merge pull request #900 from 36000/density_task

[ENH] Density task

AFQ/api/group.py

@@ -776,6 +776,48 @@ def upload_to_s3(self, s3fs, remote_path):
         if op.exists(self.afqb_path):
             s3fs.put(self.afqb_path, remote_path, recursive=True)
 
+    def export_group_density(self, boolify=True):
+        """
+        Generate a group density map by combining single subject density maps.
+
+        Parameters
+        ----------
+        boolify : bool
+            Whether to turn subject streamline count images into booleans
+            before adding them into the group density map.
+
+        Return
+        ------
+        Path to density nifti file.
+        """
+        densities = self.export("density_maps", collapse=False)
+        ex_density_init = nib.load(densities[
+            self.valid_sub_list[0]][
+                self.valid_ses_list[0]])  # for shape and header
+
+        group_density = np.zeros_like(ex_density_init.get_fdata())
+        self.logger.info("Generating Group Density...")
+        for ii in tqdm(range(len(self.valid_ses_list))):
+            this_sub = self.valid_sub_list[ii]
+            this_ses = self.valid_ses_list[ii]
+            this_density = nib.load(densities[this_sub][this_ses]).get_fdata()
+            if boolify:
+                this_density = this_density.astype(bool)
+
+            group_density = group_density + this_density
+        group_density = group_density / len(self.valid_sub_list)
+        group_density = nib.Nifti1Image(
+            group_density,
+            ex_density_init.affine,
+            header=ex_density_init.header
+        )
+
+        out_fname = op.abspath(op.join(
+            self.afq_path,
+            f"desc-density_subjects-all_space-MNI_dwi.nii.gz"))
+        nib.save(group_density, out_fname)
+        return out_fname
+
     def assemble_AFQ_browser(self, output_path=None, metadata=None,
                              page_title="AFQ Browser", page_subtitle="",
                              page_title_link="", page_subtitle_link=""):

AFQ/tasks/decorators.py

@@ -130,16 +130,18 @@ def wrapper_as_file(*args, **kwargs):
                 tracking_params=tracking_params,
                 segmentation_params=segmentation_params)
             if not op.exists(this_file):
-                img_trk_or_csv, meta = func(*args[:og_arg_count], **kwargs)
+                img_trk_np_or_csv, meta = func(*args[:og_arg_count], **kwargs)
 
                 logger.info(f"Saving {this_file}")
-                if isinstance(img_trk_or_csv, nib.Nifti1Image):
-                    nib.save(img_trk_or_csv, this_file)
-                elif isinstance(img_trk_or_csv, StatefulTractogram):
+                if isinstance(img_trk_np_or_csv, nib.Nifti1Image):
+                    nib.save(img_trk_np_or_csv, this_file)
+                elif isinstance(img_trk_np_or_csv, StatefulTractogram):
                     save_tractogram(
-                        img_trk_or_csv, this_file, bbox_valid_check=False)
+                        img_trk_np_or_csv, this_file, bbox_valid_check=False)
+                elif isinstance(img_trk_np_or_csv, np.ndarray):
+                    np.save(this_file, img_trk_np_or_csv)
                 else:
-                    img_trk_or_csv.to_csv(this_file)
+                    img_trk_np_or_csv.to_csv(this_file)
 
                 if include_seg:
                     meta["dependent"] = "rec"

AFQ/tasks/segmentation.py

@@ -8,7 +8,7 @@
 
 import pimms
 
-from AFQ.tasks.decorators import as_file
+from AFQ.tasks.decorators import as_file, as_img
 from AFQ.tasks.utils import get_fname, with_name, str_to_desc
 import AFQ.segmentation as seg
 from AFQ.utils.path import drop_extension
@@ -17,6 +17,7 @@
 from AFQ.data.s3bids import write_json
 import AFQ.api.bundle_dict as abd
 import AFQ.utils.streamlines as aus
+import AFQ.utils.volume as auv
 
 from dipy.io.streamline import load_tractogram, save_tractogram
 from dipy.io.stateful_tractogram import Space
@@ -260,6 +261,27 @@ def export_bundle_lengths(data_imap,
     return counts_df, dict(sources=bundles_files)
 
 
+@pimms.calc("density_maps")
+@as_file('_desc-density_dwi.nii.gz', include_track=True, include_seg=True)
+@as_img
+def export_density_maps(clean_bundles, dwi, data_imap):
+    """
+    full path to 4d nifti file containing streamline counts per voxel
+    per bundle, where the 4th dimension encodes the bundle
+    """
+    bundle_dict = data_imap["bundle_dict"]
+    seg_sft = aus.SegmentedSFT.fromfile(
+        clean_bundles)
+    entire_density_map = np.zeros((*nib.load(dwi).shape[:3], len(bundle_dict)))
+    for ii, bundle_name in enumerate(bundle_dict.keys()):
+        bundle_sl = seg_sft.get_bundle(bundle_name)
+        bundle_density = auv.density_map(bundle_sl).get_fdata()
+        entire_density_map[..., ii] = bundle_density
+
+    return entire_density_map, dict(
+        source=clean_bundles, bundles=list(bundle_dict.keys()))
+
+
 @pimms.calc("profiles")
 @as_file('_desc-profiles_dwi.csv', include_track=True, include_seg=True)
 def tract_profiles(clean_bundles, data_imap,
@@ -394,6 +416,7 @@ def get_segmentation_plan(kwargs):
         export_sl_counts,
         export_bundle_lengths,
         export_bundles,
+        export_density_maps,
         clean_bundles,
         segment,
         tract_profiles])

AFQ/utils/volume.py

@@ -89,7 +89,7 @@ def patch_up_roi(roi, bundle_name="ROI", make_convex=True):
     return hole_filled
 
 
-def density_map(tractogram, n_sls=None, to_vox=False, normalize=False):
+def density_map(tractogram, n_sls=None, normalize=False):
     """
     Create a streamline density map.
     based on:
@@ -104,10 +104,6 @@ def density_map(tractogram, n_sls=None, to_vox=False, normalize=False):
         n_sls to randomly select to make the density map.
         If None, all streamlines are used.
         Default: None
-    to_vox : bool, optional
-        Whether to put the stateful tractogram in VOX space before making
-        the density map.
-        Default: False
     normalize : bool, optional
         Whether to normalize maximum values to 1.
         Default: False
@@ -116,8 +112,7 @@ def density_map(tractogram, n_sls=None, to_vox=False, normalize=False):
     -------
     Nifti1Image containing the density map.
     """
-    if to_vox:
-        tractogram.to_vox()
+    tractogram.to_vox()
 
     sls = tractogram.streamlines
     if n_sls is not None:

examples/plot_afq_callosal.py

@@ -6,6 +6,8 @@
 http://hdl.handle.net/1773/34926
 """
 import os.path as op
+import matplotlib.pyplot as plt
+import nibabel as nib
 
 import plotly
 
@@ -74,6 +76,22 @@
 myafq.export_all()
 
 
+##########################################################################
+# Create Group Density Maps:
+# -------------------------
+#
+# pyAFQ can make density maps of streamline counts per subject/session
+# by calling `myafq.export("density_map")`. When using `GroupAFQ`, you can also
+# combine these into one file by calling `myafq.export_group_density()`.
+group_density = myafq.export_group_density()
+group_density = nib.load(group_density).get_fdata()
+fig, ax = plt.subplots(1)
+ax.matshow(
+    group_density[:, :, group_density.shape[-1] // 2, 0],
+    cmap='viridis')
+ax.axis("off")
+
+
 ##########################################################################
 # Visualizing bundles and tract profiles:
 # ---------------------------------------

examples/plot_callosal_tract_profile.py

@@ -331,7 +331,7 @@
     save_tractogram(tractogram, op.join(working_dir, 'dti_streamlines.trk'),
                     bbox_valid_check=False)
 
-    tractogram_img = density_map(tractogram, n_sls=1000, to_vox=True)
+    tractogram_img = density_map(tractogram, n_sls=1000)
     nib.save(tractogram_img, op.join(working_dir,
                                      'afq_dti_density_map.nii.gz'))
 else:
@@ -402,7 +402,7 @@
     save_tractogram(tractogram, op.join(working_dir, f'afq_{bundle}_seg.trk'),
                     bbox_valid_check=False)
 
-    tractogram_img = density_map(tractogram, n_sls=1000, to_vox=True)
+    tractogram_img = density_map(tractogram, n_sls=1000)
     nib.save(tractogram_img, op.join(working_dir,
                                      f'afq_{bundle}_seg_density_map.nii.gz'))
     show_anatomical_slices(tractogram_img.get_fdata(),
@@ -448,7 +448,7 @@
     save_tractogram(tractogram, op.join(working_dir, f'afq_{bundle}.trk'),
                     bbox_valid_check=False)
 
-    tractogram_img = density_map(tractogram, n_sls=1000, to_vox=True)
+    tractogram_img = density_map(tractogram, n_sls=1000)
     nib.save(tractogram_img, op.join(working_dir,
                                      f'afq_{bundle}_density_map.nii.gz'))
     show_anatomical_slices(tractogram_img.get_fdata(),