Merge pull request #436 from mgxd/fix/mcribs-clipping

ENH: Minimize clipping prior to surface reconstruction with MCRIBS

Author: mgxd

nibabies/interfaces/mcribs.py

@@ -32,7 +32,7 @@ class MCRIBReconAllInputSpec(CommandLineInputSpec):
     t1w_file = File(
         exists=True,
         copyfile=True,
-        desc='T1w to be used for deformable (must be registered to T2w image)',
+        desc='T1w to be used for deformable (must be in register with T2w image)',
     )
     t2w_file = File(
         exists=True,
@@ -47,14 +47,20 @@ class MCRIBReconAllInputSpec(CommandLineInputSpec):
 
     # MCRIBS options
     conform = traits.Bool(
+        False,
+        usedefault=True,
         argstr='--conform',
         desc='Reorients to radiological, axial slice orientation. Resamples to isotropic voxels',
     )
     tissueseg = traits.Bool(
+        False,
+        usedefault=True,
         argstr='--tissueseg',
         desc='Perform tissue type segmentation',
     )
     surfrecon = traits.Bool(
+        False,
+        usedefault=True,
         argstr='--surfrecon',
         desc='Reconstruct surfaces',
     )
@@ -75,6 +81,8 @@ class MCRIBReconAllInputSpec(CommandLineInputSpec):
         desc='Use Deformable fast collision test',
     )
     autorecon_after_surf = traits.Bool(
+        False,
+        usedefault=True,
         argstr='--autoreconaftersurf',
         desc='Do all steps after surface reconstruction',
     )
@@ -99,25 +107,51 @@ class MCRIBReconAll(CommandLine):
     output_spec = MCRIBReconAllOutputSpec
     _no_run = False
 
+    _expected_files = {
+        'surfrecon': {
+            'meshes': (
+                'pial-lh-reordered.vtp',
+                'pial-rh-reordered.vtp',
+                'white-rh.vtp',
+                'white-lh.vtp',
+            )
+        },
+        'autorecon': {
+            'mri': ('T2.mgz', 'aseg.presurf.mgz', 'ribbon.mgz', 'brain.mgz'),
+            'label': ('lh.cortex.label', 'rh.cortex.label'),
+            'stats': ('aseg.stats', 'brainvol.stats', 'lh.aparc.stats', 'rh.curv.stats'),
+            'surf': (
+                'lh.pial', 'rh.pial',
+                'lh.white', 'rh.white',
+                'lh.curv', 'rh.curv',
+                'lh.thickness', 'rh.thickness'),
+        }
+    }  # fmt:skip
+
     @property
     def cmdline(self):
         cmd = super().cmdline
-        # Avoid processing if valid
+        if self.inputs.surfrecon and not self.inputs.t2w_file:
+            raise AttributeError('Missing `t2w_file` input.')
+
+        # If an output directory is provided, check if we can skip the run
         if self.inputs.outdir:
             sid = self.inputs.subject_id
             # Check MIRTK surface recon deformable
             if self.inputs.surfrecon:
                 surfrecon_dir = Path(self.inputs.outdir) / sid / 'SurfReconDeformable' / sid
-                if self._verify_surfrecon_outputs(surfrecon_dir, error=False):
+                if self._verify_outputs('surfrecon', surfrecon_dir):
                     self._no_run = True
             # Check FS directory population
-            elif self.inputs.autorecon_after_surf:
+            if self.inputs.autorecon_after_surf:
                 fs_dir = Path(self.inputs.outdir) / sid / 'freesurfer' / sid
-                if self._verify_autorecon_outputs(fs_dir, error=False):
+                if self._verify_outputs('autorecon', fs_dir):
                     self._no_run = True
+                else:
+                    self._no_run = False
 
             if self._no_run:
-                return 'echo MCRIBSReconAll: nothing to do'
+                return 'echo MCRIBReconAll: nothing to do'
         return cmd
 
     def _setup_directory_structure(self, mcribs_dir: Path) -> None:
@@ -204,8 +238,6 @@ def _run_interface(self, runtime):
         mcribs_dir = self.inputs.outdir or Path(runtime.cwd) / 'mcribs'
         self._mcribs_dir = Path(mcribs_dir)
         if self.inputs.surfrecon:
-            if not self.inputs.t2w_file:
-                raise AttributeError('Missing T2w input')
             self._setup_directory_structure(self._mcribs_dir)
         # overwrite CWD to be in MCRIB subject's directory
         runtime.cwd = str(self._mcribs_dir / self.inputs.subject_id)
@@ -217,11 +249,11 @@ def _list_outputs(self):
         if self.inputs.surfrecon:
             # verify surface reconstruction was successful
             surfrecon_dir = self._mcribs_dir / sid / 'SurfReconDeformable' / sid
-            self._verify_surfrecon_outputs(surfrecon_dir, error=True)
+            self._verify_outputs('surfrecon', surfrecon_dir, error=True)
 
         mcribs_fs = self._mcribs_dir / sid / 'freesurfer' / sid
         if self.inputs.autorecon_after_surf:
-            self._verify_autorecon_outputs(mcribs_fs, error=True)
+            self._verify_outputs('autorecon', mcribs_fs, error=True)
 
         outputs['mcribs_dir'] = str(self._mcribs_dir)
         if self.inputs.autorecon_after_surf and self.inputs.subjects_dir:
@@ -241,56 +273,19 @@ def _list_outputs(self):
 
         return outputs
 
-    @staticmethod
-    def _verify_surfrecon_outputs(surfrecon_dir: Path, error: bool) -> bool:
+    def _verify_outputs(self, step: str, root: Path, error: bool = False) -> bool:
         """
-        Sanity check to ensure the surface reconstruction was successful.
+        Method to check to ensure the expected files are present successful.
 
         MCRIBReconAll does not return a failing exit code if a step failed, which leads
         this interface to be marked as completed without error in such cases.
         """
-        # fmt:off
-        surfrecon_files = {
-            'meshes': (
-                'pial-lh-reordered.vtp',
-                'pial-rh-reordered.vtp',
-                'white-rh.vtp',
-                'white-lh.vtp',
-            )
-        }
-        # fmt:on
-        for d, fls in surfrecon_files.items():
-            for fl in fls:
-                if not (surfrecon_dir / d / fl).exists():
-                    if error:
-                        raise FileNotFoundError(f'SurfReconDeformable missing: {fl}')
-                    return False
-        return True
-
-    @staticmethod
-    def _verify_autorecon_outputs(fs_dir: Path, error: bool) -> bool:
-        """
-        Sanity check to ensure the necessary FreeSurfer files have been created.
 
-        MCRIBReconAll does not return a failing exit code if a step failed, which leads
-        this interface to be marked as completed without error in such cases.
-        """
-        # fmt:off
-        fs_files = {
-            'mri': ('T2.mgz', 'aseg.presurf.mgz', 'ribbon.mgz', 'brain.mgz'),
-            'label': ('lh.cortex.label', 'rh.cortex.label'),
-            'stats': ('aseg.stats', 'brainvol.stats', 'lh.aparc.stats', 'rh.curv.stats'),
-            'surf': (
-                'lh.pial', 'rh.pial',
-                'lh.white', 'rh.white',
-                'lh.curv', 'rh.curv',
-                'lh.thickness', 'rh.thickness'),
-        }
-        # fmt:on
-        for d, fls in fs_files.items():
-            for fl in fls:
-                if not (fs_dir / d / fl).exists():
+        expected = self._expected_files[step]
+        for d, files in expected.items():
+            for fl in files:
+                if not (root / d / fl).exists():
                     if error:
-                        raise FileNotFoundError(f'FreeSurfer directory missing: {fl}')
+                        raise FileNotFoundError(f'{step.capitalize()} missing: {fl}')
                     return False
         return True

nibabies/interfaces/tests/test_mcribs.py

@@ -0,0 +1,65 @@
+import shutil
+from pathlib import Path
+
+import pytest
+
+from nibabies.interfaces.mcribs import MCRIBReconAll
+
+SUBJECT_ID = 'X'
+
+
+@pytest.fixture
+def mcribs_directory(tmp_path):
+    def make_tree(path, tree):
+        for d, fls in tree.items():
+            (path / d).mkdir(exist_ok=True)
+            for f in fls:
+                (path / d / f).touch()
+
+    root = tmp_path / 'mcribs'
+    surfrecon = root / SUBJECT_ID / 'SurfReconDeformable' / SUBJECT_ID
+    surfrecon.mkdir(parents=True, exist_ok=True)
+    make_tree(surfrecon, MCRIBReconAll._expected_files['surfrecon'])
+    autorecon = root / SUBJECT_ID / 'freesurfer' / SUBJECT_ID
+    autorecon.mkdir(parents=True, exist_ok=True)
+    make_tree(autorecon, MCRIBReconAll._expected_files['autorecon'])
+
+    yield root
+
+    shutil.rmtree(root)
+
+
+def test_MCRIBReconAll(mcribs_directory):
+    t2w = Path('T2w.nii.gz')
+    t2w.touch()
+
+    surfrecon = MCRIBReconAll(
+        subject_id=SUBJECT_ID,
+        surfrecon=True,
+        surfrecon_method='Deformable',
+        join_thresh=1.0,
+        fast_collision=True,
+    )
+
+    # Requires T2w input
+    with pytest.raises(AttributeError):
+        surfrecon.cmdline  # noqa
+
+    surfrecon.inputs.t2w_file = t2w
+    # Since no existing directory is found, will run fresh
+    assert 'MCRIBReconAll --deformablefastcollision --deformablejointhresh' in surfrecon.cmdline
+
+    # But should not need to run again
+    surfrecon.inputs.outdir = mcribs_directory
+    assert surfrecon.cmdline == 'echo MCRIBReconAll: nothing to do'
+
+    t2w.unlink()
+
+    autorecon = MCRIBReconAll(
+        subject_id=SUBJECT_ID,
+        autorecon_after_surf=True,
+    )
+    # No need for T2w here
+    assert autorecon.cmdline == 'MCRIBReconAll --autoreconaftersurf X'
+    autorecon.inputs.outdir = mcribs_directory
+    assert autorecon.cmdline == 'echo MCRIBReconAll: nothing to do'

nibabies/workflows/anatomical/fit.py

@@ -1030,7 +1030,6 @@ def init_infant_anat_fit_wf(
         surface_recon_wf = init_mcribs_surface_recon_wf(
             omp_nthreads=omp_nthreads,
             use_aseg=bool(anat_aseg),
-            use_mask=True,
             precomputed=precomputed,
             mcribs_dir=str(config.execution.mcribs_dir),
         )
@@ -1040,10 +1039,8 @@ def init_infant_anat_fit_wf(
                 ('subject_id', 'inputnode.subject_id'),
                 ('subjects_dir', 'inputnode.subjects_dir'),
             ]),
-            (t2w_buffer, surface_recon_wf, [
-                ('t2w_preproc', 'inputnode.t2w'),
-                ('t2w_mask', 'inputnode.in_mask'),
-            ]),
+            (t2w_validate, surface_recon_wf, [('out_file', 'inputnode.t2w')]),
+            (t2w_buffer, surface_recon_wf, [('t2w_mask', 'inputnode.in_mask'),]),
             (aseg_buffer, surface_recon_wf, [
                 ('anat_aseg', 'inputnode.in_aseg'),
             ]),
@@ -1950,7 +1947,6 @@ def init_infant_single_anat_fit_wf(
         surface_recon_wf = init_mcribs_surface_recon_wf(
             omp_nthreads=omp_nthreads,
             use_aseg=bool(anat_aseg),
-            use_mask=True,
             precomputed=precomputed,
             mcribs_dir=str(config.execution.mcribs_dir),
         )
@@ -1960,10 +1956,8 @@ def init_infant_single_anat_fit_wf(
                 ('subject_id', 'inputnode.subject_id'),
                 ('subjects_dir', 'inputnode.subjects_dir'),
             ]),
-            (anat_buffer, surface_recon_wf, [
-                ('anat_preproc', 'inputnode.t2w'),
-                ('anat_mask', 'inputnode.in_mask'),
-            ]),
+            (anat_validate, surface_recon_wf, [('out_file', 'inputnode.t2w')]),
+            (anat_buffer, surface_recon_wf, [('anat_mask', 'inputnode.in_mask')]),
             (aseg_buffer, surface_recon_wf, [
                 ('anat_aseg', 'inputnode.in_aseg'),
             ]),

nibabies/workflows/anatomical/surfaces.py

@@ -6,6 +6,7 @@
 from nipype.interfaces import freesurfer as fs
 from nipype.interfaces import io as nio
 from nipype.interfaces import utility as niu
+from nipype.interfaces.ants import N4BiasFieldCorrection
 from nipype.pipeline import engine as pe
 from niworkflows.engine.workflows import LiterateWorkflow
 from niworkflows.interfaces.freesurfer import (
@@ -14,6 +15,7 @@
 from niworkflows.interfaces.freesurfer import (
     PatchedRobustRegister as RobustRegister,
 )
+from niworkflows.interfaces.morphology import BinaryDilation
 from niworkflows.interfaces.patches import FreeSurferSource
 from smriprep.interfaces.freesurfer import MakeMidthickness
 from smriprep.interfaces.workbench import SurfaceResample
@@ -42,7 +44,6 @@ def init_mcribs_surface_recon_wf(
     *,
     omp_nthreads: int,
     use_aseg: bool,
-    use_mask: bool,
     precomputed: dict,
     mcribs_dir: str | None = None,
     name: str = 'mcribs_surface_recon_wf',
@@ -119,7 +120,27 @@ def init_mcribs_surface_recon_wf(
     fs_to_mcribs = pe.Node(MapLabels(mappings=fs2mcribs), name='fs_to_mcribs')
 
     t2w_las = pe.Node(ReorientImage(target_orientation='LAS'), name='t2w_las')
-    seg_las = t2w_las.clone(name='seg_las')
+    seg_las = pe.Node(ReorientImage(target_orientation='LAS'), name='seg_las')
+
+    # dilated mask and use in recon-neonatal-cortex
+    mask_dil = pe.Node(BinaryDilation(radius=3), name='mask_dil')
+    mask_las = pe.Node(ReorientImage(target_orientation='LAS'), name='mask_las')
+
+    # N4BiasCorrection occurs in MCRIBTissueSegMCRIBS (which is skipped)
+    # Run it (with mask to rescale intensities) prior injection
+    n4_mcribs = pe.Node(
+        N4BiasFieldCorrection(
+            dimension=3,
+            bspline_fitting_distance=200,
+            save_bias=True,
+            copy_header=True,
+            n_iterations=[50] * 5,
+            convergence_threshold=1e-7,
+            rescale_intensities=True,
+            shrink_factor=4,
+        ),
+        name='n4_mcribs',
+    )
 
     mcribs_recon = pe.Node(
         MCRIBReconAll(
@@ -128,43 +149,36 @@ def init_mcribs_surface_recon_wf(
             join_thresh=1.0,
             fast_collision=True,
             nthreads=omp_nthreads,
+            outdir=mcribs_dir,
         ),
         name='mcribs_recon',
         mem_gb=5,
     )
-    if mcribs_dir:
-        mcribs_recon.inputs.outdir = mcribs_dir
-        mcribs_recon.config = {'execution': {'remove_unnecessary_outputs': False}}
-
-    if use_mask:
-        # If available, dilated mask and use in recon-neonatal-cortex
-        from niworkflows.interfaces.morphology import BinaryDilation
-
-        mask_dil = pe.Node(BinaryDilation(radius=3), name='mask_dil')
-        mask_las = t2w_las.clone(name='mask_las')
-        workflow.connect([
-            (inputnode, mask_dil, [('in_mask', 'in_mask')]),
-            (mask_dil, mask_las, [('out_mask', 'in_file')]),
-            (mask_las, mcribs_recon, [('out_file', 'mask_file')]),
-        ])  # fmt:skip
+    mcribs_recon.config = {'execution': {'remove_unnecessary_outputs': False}}
 
     mcribs_postrecon = pe.Node(
         MCRIBReconAll(autorecon_after_surf=True, nthreads=omp_nthreads),
         name='mcribs_postrecon',
         mem_gb=5,
     )
+    mcribs_postrecon.config = {'execution': {'remove_unnecessary_outputs': False}}
 
     fssource = pe.Node(FreeSurferSource(), name='fssource', run_without_submitting=True)
     midthickness_wf = init_make_midthickness_wf(omp_nthreads=omp_nthreads)
 
     workflow.connect([
         (inputnode, t2w_las, [('t2w', 'in_file')]),
         (inputnode, fs_to_mcribs, [('in_aseg', 'in_file')]),
+        (inputnode, mask_dil, [('in_mask', 'in_mask')]),
+        (mask_dil, mask_las, [('out_mask', 'in_file')]),
+        (mask_las, mcribs_recon, [('out_file', 'mask_file')]),
         (fs_to_mcribs, seg_las, [('out_file', 'in_file')]),
         (inputnode, mcribs_recon, [
             ('subjects_dir', 'subjects_dir'),
             ('subject_id', 'subject_id')]),
-        (t2w_las, mcribs_recon, [('out_file', 't2w_file')]),
+        (t2w_las, n4_mcribs, [('out_file', 'input_image')]),
+        (mask_las, n4_mcribs, [('out_file', 'mask_image')]),
+        (n4_mcribs, mcribs_recon, [('output_image', 't2w_file')]),
         (seg_las, mcribs_recon, [('out_file', 'segmentation_file')]),
         (inputnode, mcribs_postrecon, [
             ('subjects_dir', 'subjects_dir'),