ENH: Conform precomputed mask/aseg if available (only for single runs)

Author: mgxd

nibabies/workflows/anatomical/template.py

@@ -12,6 +12,8 @@ def init_anat_template_wf(
     longitudinal: bool = False,
     bspline_fitting_distance: int = 200,
     sloppy: bool = False,
+    precomputed_mask: bool = False,
+    precomputed_aseg: bool = False,
     name: str = "anat_template_wf",
 ) -> LiterateWorkflow:
     """
@@ -78,9 +80,20 @@ def init_anat_template_wf(
 `mri_robust_template` [FreeSurfer {fs_ver}, @fs_template].
 """
 
-    inputnode = pe.Node(niu.IdentityInterface(fields=["anat_files"]), name="inputnode")
+    inputnode = pe.Node(
+        niu.IdentityInterface(fields=["anat_files", "anat_mask", "anat_aseg"]), name="inputnode"
+    )
     outputnode = pe.Node(
-        niu.IdentityInterface(fields=["out_file", "valid_list", "realign_xfms", "out_report"]),
+        niu.IdentityInterface(
+            fields=[
+                "out_file",
+                "valid_list",
+                "realign_xfms",
+                "out_report",
+                "anat_mask",
+                "anat_aseg",
+            ],
+        ),
         name="outputnode",
     )
 
@@ -101,6 +114,28 @@ def init_anat_template_wf(
     # fmt:on
 
     if num_files == 1:
+        if precomputed_mask:
+            anat_mask_conform = pe.Node(Conform(), name='anat_mask_conform')
+            # fmt:off
+            wf.connect([
+                (inputnode, anat_mask_conform, [('aseg', 'in_file')]),
+                (anat_ref_dimensions, anat_mask_conform, [
+                    ('target_zooms', 'target_zooms'),
+                    ('target_shape', 'target_shape')]),
+                (anat_mask_conform, outputnode, [('out_file', 'anat_mask')]),
+            ])
+            # fmt:on
+        if precomputed_aseg:
+            anat_aseg_conform = pe.Node(Conform(), name='anat_aseg_conform')
+            # fmt:off
+            wf.connect([
+                (inputnode, anat_aseg_conform, [('aseg', 'in_file')]),
+                (anat_ref_dimensions, anat_aseg_conform, [
+                    ('target_zooms', 'target_zooms'),
+                    ('target_shape', 'target_shape')]),
+                (anat_aseg_conform, outputnode, [('out_file', 'anat_aseg')])
+            ])
+            # fmt:on
         get1st = pe.Node(niu.Select(index=[0]), name="get1st")
         outputnode.inputs.anat_realign_xfm = [
             get_file("smriprep", "data/itkIdentityTransform.txt")