👔 Reorient mri_robust_template outputs

Author: shnizzedy

CPAC/longitudinal/robust_template.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-# Copyright (C) 2024  C-PAC Developers
+# Copyright (C) 2024-2025  C-PAC Developers
 
 # This file is part of C-PAC.
 
@@ -142,11 +142,13 @@ def mri_robust_template(
     nifti_template = pe.Node(MRIConvert(out_type="niigz"), name="NIfTI-template")
     wf.connect(node, "out_file", nifti_template, "in_file")
 
-    nifti_outputs = pe.MapNode(
-        MRIConvert(), name="NIfTI-mapmov", iterfield=["in_file", "out_file"]
-    )
+    nifti_outputs = pe.MapNode(MRIConvert(), name="NIfTI-mapmov", iterfield=["in_file"])
     wf.connect(node, "mapmov", nifti_outputs, "in_file")
-    nifti_outputs.set_input(
+    reorient_outputs = cfg.orientation_node(
+        "reorient_longitudinal_template", pe.MapNode
+    )
+    wf.connect(nifti_outputs, "out_file", reorient_outputs, "in_file")
+    reorient_outputs.set_input(
         "out_file", [f"space-longitudinal{i + 1}.nii.gz" for i in range(num_sessions)]
     )
 

CPAC/utils/configuration/configuration.py

@@ -18,14 +18,14 @@
 
 import os
 import re
-from typing import Any, cast, Literal, Optional
+from typing import Any, cast, Literal, Optional, overload
 from warnings import warn
 
 from click import BadParameter
 import pkg_resources as p
 import yaml
 
-from CPAC.pipeline.nipype_pipeline_engine import Node
+from CPAC.pipeline.nipype_pipeline_engine import MapNode, Node
 from .diff import dct_diff
 
 CONFIG_KEY_TYPE = str | list[str]
@@ -640,15 +640,21 @@ def key_type_error(self, key):
             )
         )
 
-    def orientation_node(self, name: str) -> Node:
+    @overload
+    def orientation_node(self, name: str, node_type: type[MapNode]) -> MapNode: ...
+    @overload
+    def orientation_node(self, name: str, node_type: type[Node]) -> Node: ...
+    def orientation_node(
+        self, name: str, node_type: type[Node | MapNode] = Node
+    ) -> Node | MapNode:
         """Return a node configured to resample an input with AFNI 3dresample."""
         from CPAC.utils.nifti_utils import orientation_node
 
         orientation = cast(
             Literal["RPI", "LPI", "RAI", "LAI", "RAS", "LAS", "RPS", "LPS"],
             self["pipeline_setup", "desired_orientation"],
         )
-        return orientation_node(name=name, orientation=orientation)
+        return orientation_node(name=name, orientation=orientation, node_type=node_type)
 
 
 def check_pname(p_name: str, pipe_config: Configuration) -> str:

CPAC/utils/nifti_utils.py

@@ -17,7 +17,7 @@
 """Utlities for NIfTI images."""
 
 import os
-from typing import Literal
+from typing import Literal, overload
 
 import numpy as np
 import nibabel as nib
@@ -99,17 +99,33 @@ def inverse_nifti_values(image):
     return nib.nifti1.Nifti1Image(out_data, img.affine)
 
 
+@overload
 def orientation_node(
     name: str,
     orientation: Literal["RPI", "LPI", "RAI", "LAI", "RAS", "LAS", "RPS", "LPS"],
-) -> pe.Node:
+    node_type: type[pe.MapNode],
+) -> pe.MapNode: ...
+@overload
+def orientation_node(
+    name: str,
+    orientation: Literal["RPI", "LPI", "RAI", "LAI", "RAS", "LAS", "RPS", "LPS"],
+    node_type: type[pe.Node],
+) -> pe.Node: ...
+def orientation_node(
+    name: str,
+    orientation: Literal["RPI", "LPI", "RAI", "LAI", "RAS", "LAS", "RPS", "LPS"],
+    node_type: type[pe.Node | pe.MapNode] = pe.Node,
+) -> pe.Node | pe.MapNode:
     """Return a node configured to resample an input with AFNI 3dresample."""
-    return pe.Node(
-        interface=afni_utils.Resample(
+    kwargs = {
+        "interface": afni_utils.Resample(
             orientation=orientation,
             outputtype="NIFTI_GZ",
         ),
-        name=name,
-        mem_gb=0,
-        mem_x=(0.0115, "in_file", "t"),
-    )
+        "name": name,
+        "mem_gb": 0,
+        "mem_x": (0.0115, "in_file", "t"),
+    }
+    if node_type == pe.MapNode:
+        kwargs["iterfield"] = ["in_file", "out_file"]
+    return node_type(**kwargs)