ENH: Add robust co-registration between PET and MRI

docs/usage.rst

@@ -206,6 +206,17 @@ Examples: ::
     $ petprep /data/bids_root /out participant --hmc-fwhm 8 --hmc-start-time 60
     $ petprep /data/bids_root /out participant --hmc-init-frame 10 --hmc-init-frame-fix
 
+Anatomical co-registration
+--------------------------
+*PETPrep* aligns the PET reference volume to the T1-weighted anatomy before
+deriving downstream outputs. By default, FreeSurfer's ``mri_coreg`` performs
+the alignment, with the :option:`--pet2anat-dof` flag controlling the degrees
+of freedom (rigid-body, 6 dof, is the default). When working with low
+signal-to-noise references or challenging anatomy, the
+:option:`--pet2anat-robust` flag enables ``mri_robust_register`` with an NMI
+cost function to improve robustness. This mode is restricted to rigid-body
+alignment and therefore requires ``--pet2anat-dof 6``.
+
 Segmentation
 ----------------
 *PETPrep* can segment the brain into different brain regions and extract time activity curves from these regions.

docs/workflows.rst

@@ -427,42 +427,28 @@ Interpolation uses a Lanczos kernel.
 
 PET to T1w registration
 ~~~~~~~~~~~~~~~~~~~~~~~
-:py:func:`~petprep.workflows.pet.registration.init_bbreg_wf`
+:py:func:`~petprep.workflows.pet.registration.init_pet_reg_wf`
 
 .. workflow::
     :graph2use: hierarchical
     :simple_form: yes
 
-    from petprep.workflows.pet.registration import init_bbreg_wf
-    wf = init_bbreg_wf(
+    from petprep.workflows.pet.registration import init_pet_reg_wf
+    wf = init_pet_reg_wf(
         omp_nthreads=1,
-        use_bbr=True,
-        pet2anat_dof=9,
-        pet2anat_init='t2w',
+        pet2anat_dof=6,
+        mem_gb=3,
+        use_robust_register=False,
     )
 
-``pet2anat_init`` selects the initialization strategy for PET-to-anatomical
-registration. The default ``'auto'`` setting uses available metadata to choose
-between header information and intensity-based approaches.
+The PET reference volume is aligned to the skull-stripped anatomical image
+using FreeSurfer's ``mri_coreg`` with the number of degrees of freedom set via
+the :option:`--pet2anat-dof` flag. The resulting affine is converted to ITK
+format for downstream application, along with its inverse.
 
-The alignment between the reference :abbr:`EPI (echo-planar imaging)` image
-of each run and the reconstructed subject using the gray/white matter boundary
-(FreeSurfer's ``?h.white`` surfaces) is calculated by the ``bbregister`` routine.
-See :func:`petprep.workflows.pet.registration.init_bbreg_wf` for further details.
-
-.. figure:: _static/EPIT1Normalization.svg
-
-    Animation showing :abbr:`EPI (echo-planar imaging)` to T1w registration (FreeSurfer ``bbregister``)
-
-If FreeSurfer processing is disabled, FSL ``flirt`` is run with the
-:abbr:`BBR (boundary-based registration)` cost function, using the
-``fast`` segmentation to establish the gray/white matter boundary.
-See :func:`petprep.workflows.pet.registration.init_fsl_bbr_wf` for further details.
-
-After either :abbr:`BBR (boundary-based registration)` workflow is run, the resulting affine
-transform will be compared to the initial transform found by FLIRT.
-Excessive deviation will result in rejecting the BBR refinement and accepting the
-original, affine registration.
+If co-registration proves challenging, the :option:`--pet2anat-robust` flag
+switches the workflow to FreeSurfer's ``mri_robust_register`` with an NMI cost function and restricted
+to rigid-body (6 dof) transforms. This method is more robust to large initial misalignments.
 
 Resampling PET runs onto standard spaces
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

petprep/cli/parser.py

@@ -357,6 +357,12 @@ def _bids_filter(value, parser):
         help='Degrees of freedom when registering PET to anatomical images. '
         '6 degrees (rotation and translation) are used by default.',
     )
+    g_conf.add_argument(
+        '--pet2anat-robust',
+        action='store_true',
+        help='Use FreeSurfer mri_robust_register with an NMI cost function for'
+        'PET-to-T1w co-registration. This option is limited to 6 dof.',
+    )
     g_conf.add_argument(
         '--force-bbr',
         action=DeprecatedAction,
@@ -746,6 +752,9 @@ def parse_args(args=None, namespace=None):
     parser = _build_parser()
     opts = parser.parse_args(args, namespace)
 
+    if getattr(opts, 'pet2anat_robust', False) and opts.pet2anat_dof != 6:
+        parser.error('--pet2anat-robust requires --pet2anat-dof=6.')
+
     if opts.config_file:
         skip = {} if opts.reports_only else {'execution': ('run_uuid',)}
         config.load(opts.config_file, skip=skip, init=False)

petprep/config.py

@@ -556,6 +556,8 @@ class workflow(_Config):
     """Degrees of freedom of the PET-to-anatomical registration steps."""
     pet2anat_init = 'auto'
     """Initial transform for PET-to-anatomical registration."""
+    pet2anat_robust = False
+    """Use ``mri_robust_register`` for PET-to-anatomical alignment."""
     cifti_output = None
     """Generate HCP Grayordinates, accepts either ``'91k'`` (default) or ``'170k'``."""
     hires = None

petprep/interfaces/reports.py

@@ -227,6 +227,7 @@ def _generate_segment(self):
 class FunctionalSummaryInputSpec(TraitedSpec):
     registration = traits.Enum(
         'mri_coreg',
+        'mri_robust_register',
         'Precomputed',
         mandatory=True,
         desc='PET/anatomical registration method',
@@ -246,8 +247,10 @@ def _generate_segment(self):
         # TODO: Add a note about registration_init below?
         if self.inputs.registration == 'Precomputed':
             reg = 'Precomputed affine transformation'
-        else:
+        elif self.inputs.registration == 'mri_coreg':
             reg = f'FreeSurfer <code>mri_coreg</code> - {dof} dof'
+        else:
+            reg = 'FreeSurfer <code>mri_robust_register</code> (ROBENT cost)'
 
         meta = self.inputs.metadata or {}
         time_zero = meta.get('TimeZero', None)

petprep/interfaces/tests/test_reports.py

@@ -74,11 +74,15 @@ def test_subject_summary_handles_missing_task(tmp_path):
     assert 'Task: <none> (1 run)' in segment
 
 
-def test_functional_summary_with_metadata():
+@pytest.mark.parametrize(
+    'registration',
+    ['mri_coreg', 'mri_robust_register'],
+)
+def test_functional_summary_with_metadata(registration):
     from ..reports import FunctionalSummary
 
     summary = FunctionalSummary(
-        registration='mri_coreg',
+        registration=registration,
         registration_dof=6,
         orientation='RAS',
         metadata={
@@ -92,6 +96,7 @@ def test_functional_summary_with_metadata():
     )
 
     segment = summary._generate_segment()
+    assert registration in segment
     assert 'Radiotracer: [11C]DASB' in segment
     assert 'Injected dose: 100 MBq' in segment
     assert 'Number of frames: 2' in segment

petprep/workflows/pet/fit.py

@@ -226,9 +226,15 @@ def init_pet_fit_wf(
             'Please check your BIDS JSON sidecar.'
         )
 
+    registration_method = 'Precomputed'
+    if not petref2anat_xform:
+        registration_method = (
+            'mri_robust_register' if config.workflow.pet2anat_robust else 'mri_coreg'
+        )
+
     summary = pe.Node(
         FunctionalSummary(
-            registration=('Precomputed' if petref2anat_xform else 'mri_coreg'),
+            registration=registration_method,
             registration_dof=config.workflow.pet2anat_dof,
             orientation=orientation,
             metadata=metadata,
@@ -337,6 +343,7 @@ def init_pet_fit_wf(
             pet2anat_dof=config.workflow.pet2anat_dof,
             omp_nthreads=omp_nthreads,
             mem_gb=mem_gb['resampled'],
+            use_robust_register=config.workflow.pet2anat_robust,
             sloppy=config.execution.sloppy,
         )
 

petprep/workflows/pet/registration.py

@@ -41,6 +41,7 @@ def init_pet_reg_wf(
     pet2anat_dof: AffineDOF,
     mem_gb: float,
     omp_nthreads: int,
+    use_robust_register: bool = False,
     name: str = 'pet_reg_wf',
     sloppy: bool = False,
 ):
@@ -69,6 +70,10 @@ def init_pet_reg_wf(
         Size of PET file in GB
     omp_nthreads : :obj:`int`
         Maximum number of threads an individual process may use
+    use_robust_register : :obj:`bool`
+        Run FreeSurfer ``mri_robust_register`` with an NMI cost function for
+        PET-to-anatomical alignment. Only rigid-body (6 dof) alignment is
+        supported in this mode.
     name : :obj:`str`
         Name of workflow (default: ``pet_reg_wf``)
 
@@ -90,7 +95,7 @@ def init_pet_reg_wf(
         Affine transform from anatomical space to PET space (ITK format)
 
     """
-    from nipype.interfaces.freesurfer import MRICoreg
+    from nipype.interfaces.freesurfer import MRICoreg, RobustRegister
     from niworkflows.engine.workflows import LiterateWorkflow as Workflow
     from niworkflows.interfaces.nibabel import ApplyMask
     from niworkflows.interfaces.nitransforms import ConcatenateXFMs
@@ -107,26 +112,56 @@ def init_pet_reg_wf(
     )
 
     mask_brain = pe.Node(ApplyMask(), name='mask_brain')
-    mri_coreg = pe.Node(
-        MRICoreg(dof=pet2anat_dof, sep=[4], ftol=0.0001, linmintol=0.01),
-        name='mri_coreg',
-        n_procs=omp_nthreads,
-        mem_gb=5,
-    )
+    if use_robust_register:
+        coreg = pe.Node(
+            RobustRegister(
+                auto_sens=False,
+                est_int_scale=False,
+                init_orient=True,
+                args='--cost NMI',
+                max_iterations=10,
+                high_iterations=20,
+                iteration_thresh=0.01,
+            ),
+            name='mri_robust_register',
+            n_procs=omp_nthreads,
+            mem_gb=5,
+        )
+        coreg_target = 'target_file'
+        coreg_output = 'out_reg_file'
+    else:
+        coreg = pe.Node(
+            MRICoreg(dof=pet2anat_dof, sep=[4], ftol=0.0001, linmintol=0.01),
+            name='mri_coreg',
+            n_procs=omp_nthreads,
+            mem_gb=5,
+        )
+        coreg_target = 'reference_file'
+        coreg_output = 'out_lta_file'
     convert_xfm = pe.Node(ConcatenateXFMs(inverse=True), name='convert_xfm')
 
-    workflow.connect([
-        (inputnode, mask_brain, [
-            ('anat_preproc', 'in_file'),
-            ('anat_mask', 'in_mask'),
-        ]),
-        (inputnode, mri_coreg, [('ref_pet_brain', 'source_file')]),
-        (mask_brain, mri_coreg, [('out_file', 'reference_file')]),
-        (mri_coreg, convert_xfm, [('out_lta_file', 'in_xfms')]),
-        (convert_xfm, outputnode, [
-            ('out_xfm', 'itk_pet_to_t1'),
-            ('out_inv', 'itk_t1_to_pet'),
-        ]),
-    ])  # fmt:skip
+    connections = [
+        (
+            inputnode,
+            mask_brain,
+            [
+                ('anat_preproc', 'in_file'),
+                ('anat_mask', 'in_mask'),
+            ],
+        ),
+        (inputnode, coreg, [('ref_pet_brain', 'source_file')]),
+        (mask_brain, coreg, [('out_file', coreg_target)]),
+        (coreg, convert_xfm, [(coreg_output, 'in_xfms')]),
+        (
+            convert_xfm,
+            outputnode,
+            [
+                ('out_xfm', 'itk_pet_to_t1'),
+                ('out_inv', 'itk_t1_to_pet'),
+            ],
+        ),
+    ]
+
+    workflow.connect(connections)  # fmt:skip
 
     return workflow

petprep/workflows/pet/tests/test_fit.py

@@ -287,6 +287,26 @@ def test_pet_fit_stage1_inclusion(bids_root: Path, tmp_path: Path):
     assert not any(name.startswith('pet_hmc_wf') for name in wf2.list_node_names())
 
 
+def test_pet_fit_robust_registration(bids_root: Path, tmp_path: Path):
+    """Robust PET-to-anatomical registration swaps in mri_robust_register."""
+    pet_series = [str(bids_root / 'sub-01' / 'pet' / 'sub-01_task-rest_run-1_pet.nii.gz')]
+    img = nb.Nifti1Image(np.zeros((2, 2, 2, 1)), np.eye(4))
+    for path in pet_series:
+        img.to_filename(path)
+        Path(path).with_suffix('').with_suffix('.json').write_text(
+            '{"FrameTimesStart": [0], "FrameDuration": [1]}'
+        )
+
+    with mock_config(bids_dir=bids_root):
+        config.workflow.pet2anat_robust = True
+        config.workflow.pet2anat_dof = 6
+        wf = init_pet_fit_wf(pet_series=pet_series, precomputed={}, omp_nthreads=1)
+
+    node_names = wf.list_node_names()
+    assert 'pet_reg_wf.mri_robust_register' in node_names
+    assert 'pet_reg_wf.mri_coreg' not in node_names
+
+
 def test_pet_fit_requires_both_derivatives(bids_root: Path, tmp_path: Path):
     """Supplying only one of petref or HMC transforms should raise an error."""
     pet_series = [str(bids_root / 'sub-01' / 'pet' / 'sub-01_task-rest_run-1_pet.nii.gz')]