Merge pull request #113 from nipreps/refactor_segmentation

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"""
@@ -523,6 +524,24 @@ def init_single_subject_wf(subject_id: str):
                 ]),
             ])  # fmt:skip
 
+    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'),
+                ],
+            ),
+        ]
+    )
+
     if config.workflow.anat_only:
         return clean_datasinks(workflow)
 
@@ -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
 
@@ -391,9 +399,7 @@ def init_pet_wf(
                 ('t1w_tpms', 'inputnode.t1w_tpms'),
                 ('subjects_dir', 'inputnode.subjects_dir'),
                 ('subject_id', 'inputnode.subject_id'),
-            ]),
-            (pet_fit_wf, pet_pvc_wf, [
-                ('outputnode.segmentation', 'inputnode.segmentation'),
+                ('segmentation', 'inputnode.segmentation'),
             ]),
             (petref_t1w, pet_pvc_wf, [('output_image', 'inputnode.petref')]),
             (pet_pvc_wf, psf_meta, [
@@ -711,9 +717,9 @@ def init_pet_wf(
 
     workflow.connect([
         (pet_t1w_src, pet_tacs_wf, [(pet_t1w_field, 'inputnode.pet_anat')]),
-        (pet_fit_wf, pet_tacs_wf, [
-            ('outputnode.segmentation', 'inputnode.segmentation'),
-            ('outputnode.dseg_tsv', 'inputnode.dseg_tsv'),
+        (inputnode, pet_tacs_wf, [
+            ('segmentation', 'inputnode.segmentation'),
+            ('dseg_tsv', 'inputnode.dseg_tsv'),
         ]),
         (pet_tacs_wf, ds_pet_tacs, [('outputnode.timeseries', 'in_file')]),
     ])  # 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
     -------
@@ -169,6 +172,8 @@ def init_pet_fit_wf(
                 'subjects_dir',
                 'subject_id',
                 'fsnative2t1w_xfm',
+                'segmentation',
+                'dseg_tsv',
             ],
         ),
         name='inputnode',
@@ -182,8 +187,6 @@ def init_pet_fit_wf(
                 'pet_mask',
                 'motion_xfm',
                 'petref2anat_xfm',
-                'segmentation',
-                'dseg_tsv',
                 'refmask',
             ],
         ),
@@ -397,41 +400,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 +459,10 @@ def init_pet_fit_wf(
         workflow.connect(
             [
                 (
-                    segmentation_wf,
+                    inputnode,
                     refmask_wf,
                     [
-                        ('outputnode.segmentation', 'inputnode.seg_file'),
+                        ('segmentation', 'inputnode.seg_file'),
                     ],
                 ),
                 (
@@ -572,7 +544,7 @@ def init_pet_fit_wf(
             ]
         )
     else:
-        config.loggers.workflow.info('PET Stage 5: Reference mask generation skipped')
+        config.loggers.workflow.info('PET Stage 4: Reference mask generation skipped')
 
     return workflow
 

petprep/workflows/pet/tests/test_base.py

@@ -205,8 +205,8 @@ def test_pvc_receives_segmentation(bids_root: Path):
 
         wf = init_pet_wf(pet_series=pet_series, precomputed={})
 
-    edge = wf._graph.get_edge_data(wf.get_node('pet_fit_wf'), wf.get_node('pet_pvc_wf'))
-    assert ('outputnode.segmentation', 'inputnode.segmentation') in edge['connect']
+    edge = wf._graph.get_edge_data(wf.get_node('inputnode'), wf.get_node('pet_pvc_wf'))
+    assert ('segmentation', 'inputnode.segmentation') in edge['connect']
 
 
 def test_pet_tacs_wf_connections(bids_root: Path):
@@ -226,9 +226,9 @@ def test_pet_tacs_wf_connections(bids_root: Path):
     edge_anat = wf._graph.get_edge_data(wf.get_node('pet_anat_wf'), wf.get_node('pet_tacs_wf'))
     assert ('outputnode.pet_file', 'inputnode.pet_anat') in edge_anat['connect']
 
-    edge_fit = wf._graph.get_edge_data(wf.get_node('pet_fit_wf'), wf.get_node('pet_tacs_wf'))
-    assert ('outputnode.segmentation', 'inputnode.segmentation') in edge_fit['connect']
-    assert ('outputnode.dseg_tsv', 'inputnode.dseg_tsv') in edge_fit['connect']
+    edge_input = wf._graph.get_edge_data(wf.get_node('inputnode'), wf.get_node('pet_tacs_wf'))
+    assert ('segmentation', 'inputnode.segmentation') in edge_input['connect']
+    assert ('dseg_tsv', 'inputnode.dseg_tsv') in edge_input['connect']
 
     edge_ds = wf._graph.get_edge_data(wf.get_node('pet_tacs_wf'), wf.get_node('ds_pet_tacs'))
     assert ('outputnode.timeseries', 'in_file') in edge_ds['connect']

petprep/workflows/tests/test_base.py

@@ -13,7 +13,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 +96,53 @@ 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,