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

Author: oesteban

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

@@ -63,8 +63,11 @@ def init_anat_reports_wf(reportlets_dir, freesurfer,
         DerivativesDataSink(base_directory=reportlets_dir, suffix='T1w'),
         name='ds_std_t1w_report', run_without_submitting=True)
 
+    def _drop_cohort(in_template):
+        return in_template.split(':')[0]
+
     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 +77,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', _drop_cohort), 'space'),
             ('source_file', 'source_file')]),
         (norm_rpt, ds_std_t1w_report, [('out_report', 'in_file')]),
     ])