ENH: Precomputed derivatives (#173)

* WIP: Anatomical derivatives

* RF: Ensure derivatives are not selected in existing queries

* ENH: Allow the use of precalculated derivatives

* STY: black

* FIX: Syntax errors, remove anatomical derivatives

* ENH: Port `--derivatives` to wrapper script

* ENH: Fast-track if precomputed mask is available

* ENH: Fast-track if precomputed aseg is available

* STY: black

* DOC: Add section on `--derivatives` usage

* Apply suggestions from code review

Co-authored-by: Chris Markiewicz <[email protected]>

* FIX: Remove session restriction for now, add info log

* FIX: Remove duplicate connection

* ENH: Add flag to control derivatives filtering

* ENH: Add filtering options to wrapper arguments

* STY: black

this is the final straw

* FIX: Remove excess print

* Update nibabies/utils/bids.py

Co-authored-by: Chris Markiewicz <[email protected]>

Co-authored-by: Chris Markiewicz <[email protected]>

Author: mgxd

docs/faqs.md

@@ -1,6 +1,29 @@
 # Tips and FAQs
 
-## Providing your own segmentations
+## Leveraging precomputed results
+
+Whether to allow for manual intervention for tough cases, or simply to reduce processing time, *NiBabies* can allow the use of certain pre-computed files during processing.
+Initial support is limited to the following files:
+- Anatomical mask in T1w space
+- Antomical segmentation (aseg) in T1w space
+
+To use pre-computed results, one or more [BIDS Derivatives](https://bids-specification.readthedocs.io/en/stable/05-derivatives/01-introduction.html#bids-derivatives) directories must be passed in to *NiBabies* using the `--derivatives` flag.
+Derivative directories must include a [`dataset_description.json` and the required fields](https://bids-specification.readthedocs.io/en/stable/03-modality-agnostic-files.html#derived-dataset-and-pipeline-description).
+Additionally, files must include the `space-orig` key-value pair in the name.
+
+A sample layout of a derivatives directory can be found below:
+
+```
+my_precomputed/
+├── dataset_description.json
+└── sub-01
+    └── anat
+        ├── sub-01_space-orig_desc-aseg_dseg.nii.gz
+        ├── sub-01_space-orig_desc-brain_mask.json
+        └── sub-01_space-orig_desc-brain_mask.nii.gz
+```
+
+## Multi-atlas segmentation with joint label fusion
 
 By default, *NiBabies* will run [FSL FAST](https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/FAST) for tissue segmentation, and Infant FreeSurfer for segmentation labels.
 However, you can instead use ANTs Joint Label Fusion to generate both, granted you provide multiple atlases with anatomicals / segmentations via the `--segmentation-atlases-dir` flag.

nibabies/cli/parser.py

@@ -625,7 +625,21 @@ def _slice_time_ref(value, parser):
         "--fd-radius",
         type=float,
         default=45,
-        help="Head radius in mm for framewise displacement calculation",
+        help="Head radius in mm for framewise displacement calculation.",
+    )
+    g_baby.add_argument(
+        "-d",
+        "--derivatives",
+        type=PathExists,
+        nargs="+",
+        help="One or more directory containing pre-computed derivatives.",
+    )
+    g_baby.add_argument(
+        "--deriv-filter-file",
+        dest="derivatives_filters",
+        type=_bids_filter,
+        metavar="FILE",
+        help="A JSON file for customizing the derivatives queries.",
     )
     return parser
 
@@ -758,7 +772,7 @@ def parse_args(args=None, namespace=None):
 
     # Force initialization of the BIDSLayout
     config.execution.init()
-    all_subjects = config.execution.layout.get_subjects()
+    all_subjects = config.execution.layout.get_subjects(scope="raw")
     if config.execution.participant_label is None:
         config.execution.participant_label = all_subjects
 

nibabies/config.py

@@ -356,6 +356,10 @@ class execution(_Config):
     """Run in sloppy mode (meaning, suboptimal parameters that minimize run-time)."""
     debug = []
     """Debug mode(s)."""
+    derivatives = None
+    """One or more paths where pre-computed derivatives are found."""
+    derivatives_filters = None
+    """A dictionary of BIDS selection filters"""
     echo_idx = None
     """Select a particular echo for multi-echo EPI datasets."""
     fs_license_file = _fs_license
@@ -451,18 +455,25 @@ def init(cls):
                 database_path=_db_path,
                 reset_database=cls.bids_database_dir is None,
                 indexer=_indexer,
+                derivatives=cls.derivatives or False,
             )
             cls.bids_database_dir = _db_path
         cls.layout = cls._layout
-        if cls.bids_filters:
+
+        def _unserialize_bids_queries(queries):
             from bids.layout import Query
 
-            # unserialize pybids Query enum values
-            for acq, filters in cls.bids_filters.items():
-                cls.bids_filters[acq] = {
+            for acq, filters in queries.items():
+                queries[acq] = {
                     k: getattr(Query, v[7:-4]) if not isinstance(v, Query) and "Query" in v else v
                     for k, v in filters.items()
                 }
+            return queries
+
+        if cls.bids_filters:
+            cls.bids_filters = _unserialize_bids_queries(cls.bids_filters)
+        if cls.derivatives_filters:
+            cls.derivatives_filters = _unserialize_bids_queries(cls.derivatives_filters)
 
         if "all" in cls.debug:
             cls.debug = list(DEBUG_MODES)

nibabies/utils/bids.py

@@ -158,7 +158,7 @@ def group_bolds_ref(*, layout, subject):
 
     for ses, suffix in sorted(
         product(
-            layout.get_sessions(subject=subject) or (None,),
+            layout.get_sessions(subject=subject, scope="raw") or (None,),
             {
                 "bold",
             },
@@ -304,3 +304,46 @@ def validate_input_dir(exec_env, bids_dir, participant_label):
             subprocess.check_call(["bids-validator", str(bids_dir), "-c", temp.name])
         except FileNotFoundError:
             print("bids-validator does not appear to be installed", file=sys.stderr)
+
+
+def collect_precomputed_derivatives(layout, subject_id, derivatives_filters=None):
+    """
+    Query and collect precomputed derivatives.
+
+    This function is used to determine which workflow steps can be skipped,
+    based on the files found.
+    """
+
+    deriv_queries = {
+        'anat_mask': {
+            'datatype': 'anat',
+            'desc': 'brain',
+            'space': 'orig',
+            'suffix': 'mask',
+        },
+        'anat_aseg': {
+            'datatype': 'anat',
+            'desc': 'aseg',
+            'space': 'orig',
+            'suffix': 'dseg',
+        },
+    }
+    if derivatives_filters is not None:
+        deriv_queries.update(derivatives_filters)
+
+    derivatives = {}
+    for deriv, query in deriv_queries.items():
+        res = layout.get(
+            scope='derivatives',
+            subject=subject_id,
+            extension=['.nii', '.nii.gz'],
+            **query,
+        )
+        if not res:
+            continue
+        if len(res) > 1:  # Some queries may want multiple results
+            raise Exception(
+                f"When searching for <{deriv}>, found multiple results: {[f.path for f in res]}"
+            )
+        derivatives[deriv] = res[0]
+    return derivatives