Merge pull request #163 from oesteban/fix/template-ids-cohort

ENH: Fix template keys output in normalization workflow, when cohort present

Author: effigies

smriprep/workflows/norm.py

@@ -133,15 +133,29 @@ def init_anat_norm_wf(
             workflow.__desc__ += (', ', '.')[template == templates[-1][0]]
 
     inputnode = pe.Node(niu.IdentityInterface(fields=[
-        'moving_image', 'moving_mask', 'moving_segmentation', 'moving_tpms',
-        'lesion_mask', 'orig_t1w']),
-        name='inputnode')
-    out_fields = ['standardized', 'anat2std_xfm', 'std2anat_xfm',
-                  'std_mask', 'std_dseg', 'std_tpms', 'template', 'template_spec']
+        'lesion_mask',
+        'moving_image',
+        'moving_mask',
+        'moving_segmentation',
+        'moving_tpms',
+        'orig_t1w',
+        'template',
+    ]), name='inputnode')
+    inputnode.iterables = [('template', templates)]
+
+    out_fields = [
+        'anat2std_xfm',
+        'standardized',
+        'std2anat_xfm',
+        'std_dseg',
+        'std_mask',
+        'std_tpms',
+        'template',
+        'template_spec',
+    ]
     poutputnode = pe.Node(niu.IdentityInterface(fields=out_fields), name='poutputnode')
 
     split_desc = pe.Node(TemplateDesc(), run_without_submitting=True, name='split_desc')
-    split_desc.iterables = [('template', templates)]
 
     tf_select = pe.Node(TemplateFlowSelect(resolution=1 + debug),
                         name='tf_select', run_without_submitting=True)
@@ -169,6 +183,8 @@ def init_anat_norm_wf(
                           iterfield=['input_image'], name='std_tpms')
 
     workflow.connect([
+        (inputnode, split_desc, [('template', 'template')]),
+        (inputnode, poutputnode, [('template', 'template')]),
         (inputnode, trunc_mov, [('moving_image', 'op1')]),
         (inputnode, registration, [
             ('moving_mask', 'moving_mask'),
@@ -198,13 +214,12 @@ def init_anat_norm_wf(
         (std_mask, poutputnode, [('output_image', 'std_mask')]),
         (std_dseg, poutputnode, [('output_image', 'std_dseg')]),
         (std_tpms, poutputnode, [('output_image', 'std_tpms')]),
-        (split_desc, poutputnode, [('name', 'template'),
-                                   ('spec', 'template_spec')]),
+        (split_desc, poutputnode, [('spec', 'template_spec')]),
     ])
 
     # Provide synchronized output
     outputnode = pe.JoinNode(niu.IdentityInterface(fields=out_fields),
-                             name='outputnode', joinsource='split_desc')
+                             name='outputnode', joinsource='inputnode')
     workflow.connect([
         (poutputnode, outputnode, [(f, f) for f in out_fields]),
     ])

smriprep/workflows/outputs.py

@@ -64,7 +64,7 @@ def init_anat_reports_wf(reportlets_dir, freesurfer,
         name='ds_std_t1w_report', run_without_submitting=True)
 
     workflow.connect([
-        (inputnode, tf_select, [('template', 'template'),
+        (inputnode, tf_select, [(('template', _drop_cohort), 'template'),
                                 ('template_spec', 'template_spec')]),
         (inputnode, norm_rpt, [('template', 'before_label')]),
         (inputnode, norm_msk, [('std_t1w', 'after'),
@@ -74,7 +74,7 @@ def init_anat_reports_wf(reportlets_dir, freesurfer,
         (norm_msk, norm_rpt, [('before', 'before'),
                               ('after', 'after')]),
         (inputnode, ds_std_t1w_report, [
-            ('template', 'space'),
+            (('template', _fmt_cohort), 'space'),
             ('source_file', 'source_file')]),
         (norm_rpt, ds_std_t1w_report, [('out_report', 'in_file')]),
     ])
@@ -193,23 +193,23 @@ def init_anat_derivatives_wf(bids_root, freesurfer, num_t1w, output_dir,
         # Template
         (inputnode, ds_t1w_tpl_warp, [
             ('anat2std_xfm', 'in_file'),
-            ('template', 'to')]),
+            (('template', _drop_cohort), 'to')]),
         (inputnode, ds_t1w_tpl_inv_warp, [
             ('std2anat_xfm', 'in_file'),
-            ('template', 'from')]),
+            (('template', _drop_cohort), 'from')]),
         (inputnode, ds_t1w_tpl, [
             ('std_t1w', 'in_file'),
-            ('template', 'space')]),
+            (('template', _fmt_cohort), 'space')]),
         (inputnode, ds_std_mask, [
             ('std_mask', 'in_file'),
-            ('template', 'space'),
+            (('template', _fmt_cohort), 'space'),
             (('template', _rawsources), 'RawSources')]),
-        (inputnode, ds_std_dseg, [('template', 'space')]),
+        (inputnode, ds_std_dseg, [(('template', _fmt_cohort), 'space')]),
         (inputnode, lut_std_dseg, [('std_dseg', 'in_file')]),
         (lut_std_dseg, ds_std_dseg, [('out', 'in_file')]),
         (inputnode, ds_std_tpms, [
             ('std_tpms', 'in_file'),
-            ('template', 'space')]),
+            (('template', _fmt_cohort), 'space')]),
         (t1w_name, ds_t1w_tpl_warp, [('out', 'source_file')]),
         (t1w_name, ds_t1w_tpl_inv_warp, [('out', 'source_file')]),
         (t1w_name, ds_t1w_tpl, [('out', 'source_file')]),
@@ -327,3 +327,11 @@ def _rpt_masks(mask_file, before, after, after_mask=None):
     nb.Nifti1Image(anii.get_fdata() * msk,
                    anii.affine, anii.header).to_filename('after.nii.gz')
     return abspath('before.nii.gz'), abspath('after.nii.gz')
+
+
+def _drop_cohort(in_template):
+    return in_template.split(':')[0]
+
+
+def _fmt_cohort(in_template):
+    return in_template.replace(':', '_')