FIX: Conform mask prior to applying

Author: mgxd

nibabies/workflows/anatomical/fit.py

@@ -580,7 +580,6 @@ def init_infant_anat_fit_wf(
                 'A pre-computed T1w brain mask was provided as input and used throughout the '
                 'workflow.'
             )
-        apply_t1w_mask.inputs.in_mask = t1w_mask
         workflow.connect(apply_t1w_mask, 'out_file', t1w_buffer, 't1w_brain')
 
         if not t1w_preproc:
@@ -594,6 +593,7 @@ def init_infant_anat_fit_wf(
             workflow.connect([
                 (t1w_validate, conform_t1w_mask_wf, [('out_file', 'inputnode.ref_file')]),
                 (conform_t1w_mask_wf, t1w_buffer, [('outputnode.out_file', 't1w_mask')]),
+                (conform_t1w_mask_wf, apply_t1w_mask, [('outputnode.out_file', 'in_mask')]),
                 (t1w_validate, t1w_n4_wf, [('out_file', 'inputnode.in_anat')]),
                 (t1w_n4_wf, t1w_buffer, [('outputnode.anat_preproc', 't1w_preproc')]),
                 (t1w_n4_wf, apply_t1w_mask, [('outputnode.anat_preproc', 'in_file')]),
@@ -602,6 +602,7 @@ def init_infant_anat_fit_wf(
             LOGGER.info('ANAT Skipping T1w masking')
             workflow.connect(t1w_validate, 'out_file', apply_t1w_mask, 'in_file')
             t1w_buffer.inputs.t1w_mask = t1w_mask
+            apply_t1w_mask.inputs.in_mask = t1w_mask
 
     # T2w masking logic:
     #
@@ -713,7 +714,6 @@ def init_infant_anat_fit_wf(
                 'A pre-computed T2w brain mask was provided as input and used throughout the '
                 'workflow.'
             )
-        apply_t2w_mask.inputs.in_mask = t2w_mask
         workflow.connect(apply_t2w_mask, 'out_file', t2w_buffer, 't2w_brain')
 
         if not t2w_preproc:
@@ -728,6 +728,7 @@ def init_infant_anat_fit_wf(
             workflow.connect([
                 (t2w_validate, conform_t2w_mask_wf, [('out_file', 'inputnode.ref_file')]),
                 (conform_t2w_mask_wf, t2w_buffer, [('outputnode.out_file', 't2w_mask')]),
+                (conform_t2w_mask_wf, apply_t2w_mask, [('outputnode.out_file', 'in_mask')]),
                 (t2w_validate, t2w_n4_wf, [('out_file', 'inputnode.in_anat')]),
                 (t2w_n4_wf, t2w_buffer, [('outputnode.anat_preproc', 't2w_preproc')]),
                 (t2w_n4_wf, apply_t2w_mask, [('outputnode.anat_preproc', 'in_file')]),
@@ -736,6 +737,7 @@ def init_infant_anat_fit_wf(
             LOGGER.info('ANAT Skipping T2w masking')
             workflow.connect(t2w_validate, 'out_file', apply_t2w_mask, 'in_file')
             t2w_buffer.inputs.t2w_mask = t2w_mask
+            apply_t2w_mask.inputs.in_mask = t2w_mask
 
     # Stage 3: Coregistration
     t1w2t2w_xfm = precomputed.get('t1w2t2w_xfm')
@@ -1746,7 +1748,6 @@ def init_infant_single_anat_fit_wf(
     else:
         LOGGER.info(f'ANAT Found {reference_anat} brain mask')
         desc += 'A pre-computed brain mask was provided as input and used throughout the workflow.'
-        apply_mask.inputs.in_mask = anat_mask
         workflow.connect(apply_mask, 'out_file', anat_buffer, 'anat_brain')
 
         if not anat_preproc:
@@ -1760,6 +1761,7 @@ def init_infant_single_anat_fit_wf(
             workflow.connect([
                 (anat_validate, conform_anat_mask_wf, [('out_file', 'inputnode.ref_file')]),
                 (conform_anat_mask_wf, anat_buffer, [('outputnode.out_file', 'anat_mask')]),
+                (conform_anat_mask_wf, apply_mask, [('outputnode.out_file', 'in_mask')]),
                 (anat_validate, anat_n4_wf, [('out_file', 'inputnode.in_anat')]),
                 (anat_n4_wf, anat_buffer, [('outputnode.anat_preproc', 'anat_preproc')]),
                 (anat_n4_wf, apply_mask, [('outputnode.anat_preproc', 'in_file')]),
@@ -1768,6 +1770,7 @@ def init_infant_single_anat_fit_wf(
             LOGGER.info(f'ANAT Skipping {reference_anat} masking')
             workflow.connect(anat_validate, 'out_file', apply_mask, 'in_file')
             anat_buffer.inputs.anat_mask = anat_mask
+            apply_mask.inputs.in_mask = anat_mask
 
     # Stage 3: Segmentation
     seg_method = 'jlf' if config.execution.segmentation_atlases_dir else 'fast'