Keep gcc by indexes (#309)

allow to get rid off external (disconnected) components by tissues

Author: davidmeunier79

macapype/nodes/surface.py

@@ -4,20 +4,24 @@
 from nipype.interfaces.afni.base import AFNICommandBase
 
 
+def getLargestCC(segmentation):
+
+    from skimage.measure import label
+    import numpy as np
+
+    labels = label(segmentation)
+    assert labels.max() != 0  # assume at least 1 CC
+    largestCC = labels == np.argmax(np.bincount(labels.flat)[1:])+1
+    return largestCC, labels
+
+
 def keep_gcc(nii_file):
     import os
     import nibabel as nib
     import numpy as np
     from nipype.utils.filemanip import split_filename as split_f
 
-    def getLargestCC(segmentation):
-
-        from skimage.measure import label
-
-        labels = label(segmentation)
-        assert labels.max() != 0  # assume at least 1 CC
-        largestCC = labels == np.argmax(np.bincount(labels.flat)[1:])+1
-        return largestCC, labels
+    from macapype.nodes.surface import getLargestCC
 
     # nibabel (nifti -> np.array)
     img = nib.load(nii_file)
@@ -50,6 +54,44 @@ def getLargestCC(segmentation):
     return gcc_nii_file
 
 
+def keep_gcc_by_index(nii_file):
+    import os
+    import nibabel as nib
+    import numpy as np
+    from nipype.utils.filemanip import split_filename as split_f
+
+    from macapype.nodes.surface import getLargestCC
+
+    # nibabel (nifti -> np.array)
+    img = nib.load(nii_file)
+    data = img.get_fdata().astype(np.int16)
+
+    path, fname, ext = split_f(nii_file)
+
+    print(np.unique(data))
+
+    new_data = np.zeros(data.shape, data.dtype)
+
+    for index in np.unique(data)[1:]:
+        data_data_bin_index = np.zeros(data.shape, data.dtype)
+        data_data_bin_index[data == index] = 1
+
+        if np.sum(data_data_bin_index):
+            gcc_index_data, _ = getLargestCC(data_data_bin_index)
+            new_data[gcc_index_data] = index
+
+    # nibabel (np.array -> nifti)
+    new_img = nib.Nifti1Image(dataobj=new_data,
+                              header=img.header,
+                              affine=img.affine)
+
+    gcc_nii_file = os.path.abspath(fname + "_gcc" + ext)
+
+    nib.save(new_img, gcc_nii_file)
+
+    return gcc_nii_file
+
+
 def merge_tissues(dseg_file, keep_indexes):
 
     import os

macapype/pipelines/surface.py

@@ -11,7 +11,8 @@
 from macapype.nodes.surface import (Meshify, split_LR_mask,
                                     wrap_nii2mesh,
                                     IsoSurface, merge_tissues,
-                                    keep_gcc)
+                                    keep_gcc,
+                                    keep_gcc_by_index)
 
 from macapype.utils.utils_nodes import parse_key, NodeParams
 
@@ -704,6 +705,16 @@ def create_IsoSurface_brain_pipe(params={},
             fields=["segmented_file"]),
         name='inputnode')
 
+    # keep_gcc_mask
+    keep_gcc_mask = pe.Node(
+        interface=niu.Function(input_names=["nii_file"],
+                               output_names=["gcc_nii_file"],
+                               function=keep_gcc_by_index),
+        name="keep_gcc_mask")
+
+    IsoSurface_brain_pipe.connect(inputnode, 'segmented_file',
+                                  keep_gcc_mask, "nii_file")
+
     # merge_brain_tissues
     merge_brain_tissues = NodeParams(
         interface=niu.Function(input_names=["dseg_file", "keep_indexes"],
@@ -712,7 +723,7 @@ def create_IsoSurface_brain_pipe(params={},
         params=parse_key(params, "merge_brain_tissues"),
         name="merge_brain_tissues")
 
-    IsoSurface_brain_pipe.connect(inputnode, 'segmented_file',
+    IsoSurface_brain_pipe.connect(keep_gcc_mask, "gcc_nii_file",
                                   merge_brain_tissues, 'dseg_file')
 
     # keep_gcc_bin_mask