Support output spaces to arbitrary templates and cohorts (#184)

* seems to work

* remove restriction

qsirecon/interfaces/anatomical.py

@@ -312,11 +312,7 @@ def _run_interface(self, runtime):
 
 
 class _GetTemplateInputSpec(BaseInterfaceInputSpec):
-    template_name = traits.Enum(
-        "MNI152NLin2009cAsym",
-        "MNIInfant",
-        mandatory=True,
-    )
+    template_name = traits.Str(mandatory=True)
 
 
 class _GetTemplateOutputSpec(BaseInterfaceInputSpec):
@@ -331,10 +327,15 @@ class GetTemplate(SimpleInterface):
     def _run_interface(self, runtime):
         from templateflow.api import get as get_template
 
+        template_name = self.inputs.template_name
+        cohort = None
+        if "+" in template_name:
+            template_name, cohort = template_name.split("+")
+
         template_file = str(
             get_template(
-                self.inputs.template_name,
-                cohort=[None, "2"],
+                template_name,
+                cohort=cohort,
                 resolution="1",
                 desc=None,
                 suffix="T1w",
@@ -343,8 +344,8 @@ def _run_interface(self, runtime):
         )
         mask_file = str(
             get_template(
-                self.inputs.template_name,
-                cohort=[None, "2"],
+                template_name,
+                cohort=cohort,
                 resolution="1",
                 desc="brain",
                 suffix="mask",

qsirecon/interfaces/scalar_mapping.py

@@ -23,7 +23,6 @@
 )
 from nipype.utils.filemanip import fname_presuffix
 
-from .. import config
 from .bids import get_bids_params
 
 LOGGER = logging.getLogger("nipype.interface")
@@ -219,6 +218,7 @@ class _TemplateMapperInputSpec(ScalarMapperInputSpec):
     template_reference_image = File(exists=True, mandatory=True)
     to_template_transform = File(exists=True, mandatory=True)
     interpolation = traits.Str("NearestNeighbor", usedefault=True)
+    template_space = traits.Str(mandatory=True)
 
 
 class _TemplateMapperOutputSpec(ScalarMapperOutputSpec):
@@ -263,13 +263,9 @@ def _do_mapping(self, runtime):
             new_metadata["path"] = output_fname
             if "bids" not in new_metadata:
                 raise Exception(f"incomplete metadata spec {new_metadata}")
-            new_metadata["bids"]["space"] = (
-                "MNIInfant" if config.workflow.infant else "MNI152NLin2009cAsym"
-            )
+            new_metadata["bids"]["space"] = self.inputs.template_space
             resampled_image_metadata.append(new_metadata)
 
         self._results["template_space_scalars"] = resampled_images
         self._results["template_space_scalar_info"] = resampled_image_metadata
-        self._results["template_space"] = (
-            "MNIInfant" if config.workflow.infant else "MNI152NLin2009cAsym"
-        )
+        self._results["template_space"] = self.inputs.template_space

qsirecon/utils/bids.py

@@ -184,13 +184,14 @@ def collect_anatomical_data(
     from qsirecon.data import load as load_data
 
     anat_data = {}
-    status = {}
+    status = {"template_output_space": None}
 
     if session_id is None:
         return anat_data, {"has_qsiprep_t1w": False, "has_qsiprep_t1w_transforms": False}
 
     _spec = yaml.safe_load(load_data.readable("io_spec.yaml").read_text())
     queries = _spec["queries"]["anat"]
+
     if config.workflow.infant:
         queries["acpc_to_template_xfm"]["to"] = "MNIInfant"
         queries["template_to_acpc_xfm"]["from"] = "MNIInfant"
@@ -241,10 +242,20 @@ def collect_anatomical_data(
 
         config.loggers.utils.warning("No anat-to-template or template-to-anat transforms found.")
         status["has_qsiprep_t1w_transforms"] = False
-
+    else:
+        # Determine the output space from the transform file
+        status["template_output_space"] = _determine_output_space(anat_data)
     return anat_data, status
 
 
+def _determine_output_space(status):
+    """Determine what output space the transform maps to/from"""
+    if not status["template_to_acpc_xfm"]:
+        return None
+
+    return get_entity(status["template_to_acpc_xfm"], "from")
+
+
 def write_derivative_description(
     bids_dir,
     deriv_dir,

qsirecon/workflows/base.py

@@ -148,6 +148,7 @@ def init_single_subject_recon_wf(subject_id):
             needs_t1w_transform=bool(config.execution.atlases),
             bids_filters=config.execution.bids_filters or {},
         )
+
         config.loggers.workflow.info(
             f"Anatomical data available for {anat_input_file.path}:\n"
             f"{yaml.dump(anat_data, default_flow_style=False, indent=4)}"
@@ -250,7 +251,6 @@ def init_single_subject_recon_wf(subject_id):
             name=f"{wf_name}_dwi_specific_anat_wf",
             **highres_anat_statuses[anat_input.path],
         )
-
         inputs_dict = {
             "dwi_file": dwi_file,
             "dwi_metadata": config.execution.layout.get_metadata(dwi_file),

qsirecon/workflows/recon/anatomical.py

@@ -318,6 +318,7 @@ def init_dwi_recon_anatomical_workflow(
     has_freesurfer,
     extras_to_make,
     name,
+    template_output_space,
     prefer_dwi_mask=False,
 ):
     """Ensure that anatomical data is available for the reconstruction workflows.
@@ -398,12 +399,12 @@ def _get_status():
             "has_freesurfer_5tt_hsvs": has_freesurfer_5tt_hsvs,
             "has_qsiprep_t1w": has_qsiprep_t1w,
             "has_qsiprep_t1w_transforms": has_qsiprep_t1w_transforms,
+            "template_output_space": template_output_space,
         }
 
-    # XXX: This is a temporary solution until QSIRecon supports flexible output spaces.
     get_template = pe.Node(
         GetTemplate(
-            template_name="MNI152NLin2009cAsym" if not config.workflow.infant else "MNIInfant",
+            template_name=template_output_space,
         ),
         name="get_template",
     )

qsirecon/workflows/recon/scalar_mapping.py

@@ -165,7 +165,10 @@ def init_scalar_to_template_wf(
         name="outputnode",
     )
     workflow = Workflow(name=name)
-    template_mapper = pe.Node(TemplateMapper(**params), name="template_mapper")
+    template_mapper = pe.Node(
+        TemplateMapper(template_space=inputs_dict["template_output_space"], **params),
+        name="template_mapper",
+    )
 
     scalar_output_wf = init_scalar_output_wf()
     workflow.connect([