ENH: Refactor PET workflow segmentation process

Author: mnoergaard

petprep/workflows/base.py

@@ -164,6 +164,7 @@ def init_single_subject_wf(subject_id: str):
     )
 
     from petprep.workflows.pet.base import init_pet_wf
+    from petprep.workflows.pet.segmentation import init_segmentation_wf
 
     workflow = Workflow(name=f'sub_{subject_id}_wf')
     workflow.__desc__ = f"""
@@ -526,6 +527,24 @@ def init_single_subject_wf(subject_id: str):
     if config.workflow.anat_only:
         return clean_datasinks(workflow)
 
+    segmentation_wf = init_segmentation_wf(
+        seg=config.workflow.seg,
+        name=f'pet_{config.workflow.seg}_seg_wf',
+    )
+    workflow.connect(
+        [
+            (
+                anat_fit_wf,
+                segmentation_wf,
+                [
+                    ('outputnode.t1w_preproc', 'inputnode.t1w_preproc'),
+                    ('outputnode.subjects_dir', 'inputnode.subjects_dir'),
+                    ('outputnode.subject_id', 'inputnode.subject_id'),
+                ],
+            ),
+        ]
+    )
+
     # Append the PET section to the existing anatomical excerpt
     # That way we do not need to filter down the number of PET datasets
     pet_pre_desc = f"""
@@ -600,6 +619,10 @@ def init_single_subject_wf(subject_id: str):
                     'inputnode.sphere_reg_fsLR',
                 ),
             ]),
+            (segmentation_wf, pet_wf, [
+                ('outputnode.segmentation', 'inputnode.segmentation'),
+                ('outputnode.dseg_tsv', 'inputnode.dseg_tsv'),
+            ]),
         ])  # fmt:skip
 
         if config.workflow.level == 'full':

petprep/workflows/pet/base.py

@@ -115,6 +115,10 @@ def init_pet_wf(
         Registration spheres from fsnative to fsLR space, collated left, then right
     anat_ribbon
         Binary cortical ribbon mask in T1w space
+    segmentation
+        Segmentation file in T1w space
+    dseg_tsv
+        TSV with segmentation statistics
     anat2std_xfm
         Transform from anatomical space to standard space
     std_t1w
@@ -206,6 +210,8 @@ def init_pet_wf(
                 'midthickness_fsLR',
                 'cortex_mask',
                 'anat_ribbon',
+                'segmentation',
+                'dseg_tsv',
                 # Volumetric templates
                 'anat2std_xfm',
                 'std_t1w',
@@ -242,6 +248,8 @@ def init_pet_wf(
             ('subjects_dir', 'inputnode.subjects_dir'),
             ('subject_id', 'inputnode.subject_id'),
             ('fsnative2t1w_xfm', 'inputnode.fsnative2t1w_xfm'),
+            ('segmentation', 'inputnode.segmentation'),
+            ('dseg_tsv', 'inputnode.dseg_tsv'),
         ]),
     ])  # fmt:skip
 

petprep/workflows/pet/fit.py

@@ -51,7 +51,6 @@
 from .ref_tacs import init_pet_ref_tacs_wf
 from .reference_mask import init_pet_refmask_wf
 from .registration import init_pet_reg_wf
-from .segmentation import init_segmentation_wf
 
 
 def init_pet_fit_wf(
@@ -103,6 +102,10 @@ def init_pet_fit_wf(
         FreeSurfer subject ID
     fsnative2t1w_xfm
         LTA-style affine matrix translating from FreeSurfer-conformed subject space to T1w
+    segmentation
+        Segmentation file in T1w space
+    dseg_tsv
+        TSV with segmentation statistics
 
     Outputs
     -------
@@ -116,6 +119,10 @@ def init_pet_fit_wf(
     petref2anat_xfm
         Affine transform mapping from PET reference space to the anatomical
         space.
+    segmentation
+        Segmentation file in T1w space
+    dseg_tsv
+        TSV with segmentation statistics
 
     See Also
     --------
@@ -169,6 +176,8 @@ def init_pet_fit_wf(
                 'subjects_dir',
                 'subject_id',
                 'fsnative2t1w_xfm',
+                'segmentation',
+                'dseg_tsv',
             ],
         ),
         name='inputnode',
@@ -192,6 +201,11 @@ def init_pet_fit_wf(
 
     # If all derivatives exist, inputnode could go unconnected, so add explicitly
     workflow.add_nodes([inputnode])
+    workflow.connect(
+        [
+            (inputnode, outputnode, [('segmentation', 'segmentation'), ('dseg_tsv', 'dseg_tsv')]),
+        ]
+    )
 
     petref_buffer = pe.Node(
         niu.IdentityInterface(fields=['petref', 'pet_file']),
@@ -397,41 +411,10 @@ def init_pet_fit_wf(
     ds_petmask_wf.inputs.inputnode.source_files = [pet_file]
     workflow.connect([(merge_mask, ds_petmask_wf, [('out', 'inputnode.petmask')])])
 
-    # Stage 4: Segmentation
-    config.loggers.workflow.info(
-        'PET Stage 4: Adding segmentation workflow using the segmentation: %s', config.workflow.seg
-    )
-    segmentation_wf = init_segmentation_wf(
-        seg=config.workflow.seg,
-        name=f'pet_{config.workflow.seg}_seg_wf',
-    )
-
-    workflow.connect(
-        [
-            (
-                inputnode,
-                segmentation_wf,
-                [
-                    ('t1w_preproc', 'inputnode.t1w_preproc'),
-                    ('subject_id', 'inputnode.subject_id'),
-                    ('subjects_dir', 'inputnode.subjects_dir'),
-                ],
-            ),
-            (
-                segmentation_wf,
-                outputnode,
-                [
-                    ('outputnode.segmentation', 'segmentation'),
-                    ('outputnode.dseg_tsv', 'dseg_tsv'),
-                ],
-            ),
-        ]
-    )
-
-    # Stage 5: Reference mask generation
+    # Stage 4: Reference mask generation
     if config.workflow.ref_mask_name:
         config.loggers.workflow.info(
-            'PET Stage 5: Generating %s reference mask',
+            'PET Stage 4: Generating %s reference mask',
             config.workflow.ref_mask_name,
         )
 
@@ -487,10 +470,10 @@ def init_pet_fit_wf(
         workflow.connect(
             [
                 (
-                    segmentation_wf,
+                    inputnode,
                     refmask_wf,
                     [
-                        ('outputnode.segmentation', 'inputnode.seg_file'),
+                        ('segmentation', 'inputnode.seg_file'),
                     ],
                 ),
                 (

petprep/workflows/tests/test_base.py

@@ -4,6 +4,8 @@
 from pathlib import Path
 from unittest.mock import patch
 
+from pathlib import Path
+
 import nibabel as nb
 import numpy as np
 import pytest
@@ -13,7 +15,7 @@
 from niworkflows.utils.testing import generate_bids_skeleton
 
 from ... import config
-from ..base import init_petprep_wf
+from ..base import init_petprep_wf, init_single_subject_wf
 from ..tests import mock_config
 
 BASE_LAYOUT = {
@@ -96,6 +98,55 @@ def bids_root(tmp_path_factory):
     return bids_dir
 
 
+@pytest.fixture(scope='module')
+def multisession_bids_root(tmp_path_factory):
+    base = tmp_path_factory.mktemp('multisession')
+    bids_dir = base / 'bids'
+    bids_dir.mkdir(parents=True, exist_ok=True)
+    img = nb.Nifti1Image(np.zeros((10, 10, 10, 10)), np.eye(4))
+    (bids_dir / 'dataset_description.json').write_text(
+        '{"Name": "Test", "BIDSVersion": "1.8.0"}'
+    )
+    for ses in ['01', '02']:
+        anat_dir = bids_dir / 'sub-01' / f'ses-{ses}' / 'anat'
+        pet_dir = bids_dir / 'sub-01' / f'ses-{ses}' / 'pet'
+        anat_dir.mkdir(parents=True, exist_ok=True)
+        pet_dir.mkdir(parents=True, exist_ok=True)
+        img.to_filename(anat_dir / f'sub-01_ses-{ses}_T1w.nii.gz')
+        pet_path = pet_dir / f'sub-01_ses-{ses}_task-rest_run-1_pet.nii.gz'
+        img.to_filename(pet_path)
+        (pet_path.with_suffix('').with_suffix('.json')).write_text(
+            '{"FrameTimesStart": [0], "FrameDuration": [1]}'
+        )
+    return bids_dir
+
+
+def test_segmentation_shared_across_runs(multisession_bids_root):
+    with mock_config(bids_dir=multisession_bids_root):
+        wf = init_single_subject_wf('01')
+    flatgraph = wf._create_flat_graph()
+    generate_expanded_graph(flatgraph)
+
+    seg_wf_name = f'pet_{config.workflow.seg}_seg_wf'
+    seg_nodes = [n for n in wf.list_node_names() if n.startswith(seg_wf_name)]
+    assert seg_nodes
+
+    pet_wf_names = [
+        n
+        for n in {name.split('.')[0] for name in wf.list_node_names() if name.startswith('pet_')}
+        if n != seg_wf_name
+    ]
+    assert len(pet_wf_names) == 2
+
+    seg_node = wf.get_node(seg_wf_name)
+    for name in pet_wf_names:
+        pet_node = wf.get_node(name)
+        edge = wf._graph.get_edge_data(seg_node, pet_node)
+        assert ('outputnode.segmentation', 'inputnode.segmentation') in edge['connect']
+        assert ('outputnode.dseg_tsv', 'inputnode.dseg_tsv') in edge['connect']
+        assert all('_seg_wf' not in n for n in pet_node.list_node_names())
+
+
 def _make_params(
     pet2anat_init: str = 'auto',
     medial_surface_nan: bool = False,