ENH: Add workflow to normalize CSF prior to spatial normalization

Author: mgxd

nibabies/workflows/anatomical/preproc.py

@@ -68,27 +68,50 @@ def init_anat_preproc_wf(
     return wf
 
 
-def init_anat_csf_norm_wf(name='anat_csf_norm_wf') -> LiterateWorkflow:
+def init_csf_norm_wf(name: str = 'csf_norm_wf') -> LiterateWorkflow:
     """Replace low intensity voxels within the CSF mask with the median value."""
 
     workflow = LiterateWorkflow(name=name)
-    inputnode = niu.IdentityInterface(fields=['anat_preproc', 'anat_dseg'], name='inputnode')
+    workflow.__desc__ = (
+        'The CSF mask was used to normalize the anatomical template by the median of voxels '
+        'within the mask.'
+    )
+    inputnode = niu.IdentityInterface(fields=['anat_preproc', 'anat_tpms'], name='inputnode')
     outputnode = niu.IdentityInterface(fields=['anat_preproc'], name='outputnode')
 
-    applymask = pe.Node(ApplyMask(), name='applymask')
+    # select CSF from BIDS-ordered list (GM, WM, CSF)
+    select_csf = pe.Node(niu.Select(index=2), name='select_csf')
+    norm_csf = pe.Node(niu.Function(function=_normalize_roi), name='norm_csf')
 
-    norm2median = pe.Node(niu.Function(function=_normalize_csf), name='norm2median')
-    # 1. mask brain with CSF mask
-    # fslmaths input.nii.gz -mas aseg_label-CSF_mask.nii.gz input_CSF.nii.gz
-    # 2. get median intensity of nonzero voxels in mask
-    # fslstats input_CSF.nii.gz -P 50
-    # 3. normalize CSF-masked T2w to the median
-    # fslmaths input_CSF.nii.gz -bin -mul <median intensity from (> median_CSF.nii.gz
-    # 4. make the modified T2w, setting voxel intensity to be the max between the original T2w's,
-    # and the normalized mask from (3)'s:
-    # fslmaths input.nii.gz -max median_CSF.nii.gz input_floorCSF.nii.gz
+    workflow.connect([
+        (inputnode, select_csf, [('anat_tpms', 'in_list')]),
+        (select_csf, norm_csf, [('out', 'mask_file')]),
+        (inputnode, norm_csf, [('anat_preproc', 'in_file')]),
+        (norm_csf, outputnode, [('out', 'anat_preproc')]),
+    ])  # fmt:skip
 
     return workflow
 
 
-def _normalize_csf(in_file): ...
+def _normalize_roi(in_file, mask_file, threshold=0.2, out_file=None):
+    """Normalize low intensity voxels that fall within a given mask."""
+    import nibabel as nb
+    import numpy as np
+
+    img = nb.load(in_file)
+    img_data = img.get_fdata()
+    mask_img = nb.load(mask_file)
+    # binary mask
+    bin_mask = mask_img.get_fdata() > threshold
+    mask_data = bin_mask * img_data
+
+    median = np.median(mask_data[mask_data > 0])
+    normed_data = np.maximum(img_data, bin_mask * median)
+
+    oimg = img.__class__(normed_data, img.affine, img.header)
+    if not out_file:
+        from nipype.utils.filemanip import fname_presuffix
+
+        out_file = fname_presuffix(in_file, suffix='normed')
+    oimg.to_filename(out_file)
+    return out_file