FIX: update pet mask workflow

Author: mnoergaard

fmriprep/workflows/pet/confounds.py

@@ -36,7 +36,6 @@
 from ...config import DEFAULT_MEMORY_MIN_GB
 from ...interfaces import DerivativesDataSink
 from ...interfaces.confounds import (
-    FilterDropped,
     PETSummary,
     FramewiseDisplacement,
     FSLMotionParams,
@@ -148,7 +147,7 @@ def init_pet_confs_wf(
     from niworkflows.interfaces.fixes import FixHeaderApplyTransforms as ApplyTransforms
     from niworkflows.interfaces.images import SignalExtraction
     from niworkflows.interfaces.morphology import BinaryDilation, BinarySubtraction
-    from niworkflows.interfaces.nibabel import ApplyMask, Binarize
+    from niworkflows.interfaces.nibabel import Binarize
     from niworkflows.interfaces.utility import AddTSVHeader, DictMerge
 
     from ...interfaces.confounds import aCompCorMasks
@@ -395,8 +394,8 @@ def _select_cols(table):
         (acompcor_tfm, acompcor_bin, [('output_image', 'in_file')]),
         (acompcor_bin, merge_rois, [
             (('out_mask', _last), 'in3'),
-            (('out_mask', lambda l: l[0]), 'in1'),
-            (('out_mask', lambda l: l[1]), 'in2'),
+            (('out_mask', lambda masks: masks[0]), 'in1'),
+            (('out_mask', lambda masks: masks[1]), 'in2'),
         ]),
         (merge_rois, signals, [('out', 'label_files')]),
 

fmriprep/workflows/pet/fit.py

@@ -198,11 +198,7 @@ def init_pet_fit_wf(
 
     summary = pe.Node(
         FunctionalSummary(
-            registration=(
-                'Precomputed'
-                if petref2anat_xform
-                else 'mri_coreg'
-            ),
+            registration=('Precomputed' if petref2anat_xform else 'mri_coreg'),
             registration_dof=config.workflow.pet2anat_dof,
             orientation=orientation,
         ),
@@ -220,7 +216,6 @@ def init_pet_fit_wf(
         (petref_buffer, outputnode, [
             ('petref', 'petref'),
         ]),
-        (merge_mask, outputnode, [('out', 'pet_mask')]),
         (hmc_buffer, outputnode, [
             ('hmc_xforms', 'motion_xfm'),
         ]),
@@ -265,11 +260,17 @@ def init_pet_fit_wf(
         # Ensure all stage-1 workflows were created successfully before
         # attempting to connect them. Nipype's ``connect`` call will fail
         # with a ``NoneType`` error if any node is undefined.
-        stage1_nodes = [petref_wf, petref_buffer, ds_petref_wf,
-                        func_fit_reports_wf, petref_source_buffer]
+        stage1_nodes = [
+            petref_wf,
+            petref_buffer,
+            ds_petref_wf,
+            func_fit_reports_wf,
+            petref_source_buffer,
+        ]
         if any(node is None for node in stage1_nodes):
-            raise RuntimeError('PET reference stage could not be built - '
-                               'check inputs and configuration.')
+            raise RuntimeError(
+                'PET reference stage could not be built - check inputs and configuration.'
+            )
 
         workflow.connect([
             (petref_wf, petref_buffer, [
@@ -373,14 +374,16 @@ def init_pet_fit_wf(
     else:
         t1w_mask_tfm.inputs.transforms = petref2anat_xform
 
-    workflow.connect([
-        (inputnode, t1w_mask_tfm, [('t1w_mask', 'input_image')]),
-        (petref_buffer, t1w_mask_tfm, [('petref', 'reference_image')]),
-        (petref_buffer, petref_mask, [('petref', 'in_file')]),
-        (petref_mask, merge_mask, [('out_mask', 'mask1')]),
-        (t1w_mask_tfm, merge_mask, [('output_image', 'mask2')]),
-        (merge_mask, petref_buffer, [('out', 'pet_mask')]),
-    ])
+    workflow.connect(
+        [
+            (inputnode, t1w_mask_tfm, [('t1w_mask', 'input_image')]),
+            (petref_buffer, t1w_mask_tfm, [('petref', 'reference_image')]),
+            (petref_buffer, petref_mask, [('petref', 'in_file')]),
+            (petref_mask, merge_mask, [('out_mask', 'mask1')]),
+            (t1w_mask_tfm, merge_mask, [('output_image', 'mask2')]),
+            (merge_mask, outputnode, [('out', 'pet_mask')]),
+        ]
+    )
 
     ds_petmask_wf = init_ds_petmask_wf(
         output_dir=config.execution.petprep_dir,
@@ -482,18 +485,13 @@ def init_pet_native_wf(
     )
     outputnode.inputs.metadata = metadata
 
-    petbuffer = pe.Node(
-        niu.IdentityInterface(fields=['pet_file']), name='petbuffer'
-    )
+    petbuffer = pe.Node(niu.IdentityInterface(fields=['pet_file']), name='petbuffer')
 
     # PET source: track original PET file(s)
     # The Select interface requires an index to choose from ``inlist``. Since
     # ``pet_file`` is a single path, explicitly set the index to ``0`` to avoid
     # missing mandatory input errors when the node runs.
-    pet_source = pe.Node(
-        niu.Select(inlist=[pet_file], index=0),
-        name='pet_source'
-    )
+    pet_source = pe.Node(niu.Select(inlist=[pet_file], index=0), name='pet_source')
     validate_pet = pe.Node(ValidateImage(), name='validate_pet')
     workflow.connect([
         (pet_source, validate_pet, [('out', 'in_file')]),

fmriprep/workflows/pet/tests/test_fit.py

@@ -134,3 +134,23 @@ def test_pet_native_precomputes(
 
     flatgraph = wf._create_flat_graph()
     generate_expanded_graph(flatgraph)
+
+
+def test_pet_fit_mask_connections(bids_root: Path, tmp_path: Path):
+    """Ensure the PET mask is generated and connected correctly."""
+    pet_file = 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))
+    img.to_filename(pet_file)
+
+    with mock_config(bids_dir=bids_root):
+        wf = init_pet_fit_wf(pet_file=pet_file, precomputed={}, omp_nthreads=1)
+
+    assert 'merge_mask' in wf.list_node_names()
+    assert 'ds_petmask_wf.ds_petmask' in wf.list_node_names()
+
+    merge_mask = wf.get_node('merge_mask')
+    edge = wf._graph.get_edge_data(merge_mask, wf.get_node('outputnode'))
+    assert ('out', 'pet_mask') in edge['connect']
+
+    ds_edge = wf._graph.get_edge_data(merge_mask, wf.get_node('ds_petmask_wf'))
+    assert ('out', 'inputnode.petmask') in ds_edge['connect']

fmriprep/workflows/pet/tests/test_reference.py

@@ -6,22 +6,22 @@
 
 def test_reference_frame_select(tmp_path):
     img = nb.Nifti1Image(np.zeros((5, 5, 5, 4)), np.eye(4))
-    pet_file = tmp_path / "pet.nii.gz"
+    pet_file = tmp_path / 'pet.nii.gz'
     img.to_filename(pet_file)
 
     wf = init_raw_petref_wf(pet_file=str(pet_file), reference_frame=2)
     node_names = [n.name for n in wf._get_all_nodes()]
-    assert "extract_frame" in node_names
-    assert "gen_avg" not in node_names
-    node = wf.get_node("extract_frame")
+    assert 'extract_frame' in node_names
+    assert 'gen_avg' not in node_names
+    node = wf.get_node('extract_frame')
     assert node.interface.inputs.t_min == 2
 
 
 def test_reference_frame_average(tmp_path):
     img = nb.Nifti1Image(np.zeros((5, 5, 5, 4)), np.eye(4))
-    pet_file = tmp_path / "pet.nii.gz"
+    pet_file = tmp_path / 'pet.nii.gz'
     img.to_filename(pet_file)
 
-    wf = init_raw_petref_wf(pet_file=str(pet_file), reference_frame="average")
+    wf = init_raw_petref_wf(pet_file=str(pet_file), reference_frame='average')
     node_names = [n.name for n in wf._get_all_nodes()]
-    assert "gen_avg" in node_names
+    assert 'gen_avg' in node_names