Merge pull request #50 from nipreps/add_gm_threshold

WIP: Add gray matter threshold option to reference mask workflow

Author: mnoergaard

petprep/data/reference_mask/config.json

@@ -5,7 +5,16 @@
     "refmask_indices": [47, 8], 
     "erode_by_voxels": 1,
     "exclude_indices": [7, 46, 172, 1007, 2007, 1009, 2009, 1011, 2011, 1013, 2013, 16], 
-    "dilate_by_voxels": 3
+    "dilate_by_voxels": 3,
+    "gm_prob_threshold": 0.8
+    },
+    "thalamus": 
+    {
+    "refmask_indices": [10, 49], 
+    "erode_by_voxels": 1,
+    "exclude_indices": [14, 24], 
+    "dilate_by_voxels": 3,
+    "gm_prob_threshold": 0.2
     }
   },
   "wm": {
@@ -16,7 +25,8 @@
       "exclude_indices": [10, 49, 12, 51, 13, 52, 11, 50, 4, 43, 31, 63],
       "dilate_by_voxels": 3, 
       "smooth_fwhm_mm": 10,
-      "target_volume_ml": 40
+      "target_volume_ml": 40,
+      "gm_prob_threshold": 0
     }
   }
 }

petprep/interfaces/reference_mask.py

@@ -14,6 +14,7 @@
 
 class ExtractRefRegionInputSpec(BaseInterfaceInputSpec):
     seg_file = File(exists=True, mandatory=True, desc='Segmentation NIfTI file')
+    gm_probseg = File(exists=True, desc='Gray matter probability map for thresholding')
     config_file = File(exists=True, mandatory=True, desc='Path to the config.json file')
     segmentation_type = traits.Str(mandatory=True, desc="Type of segmentation (e.g. 'gtm', 'wm')")
     region_name = traits.Str(
@@ -32,6 +33,9 @@ class ExtractRefRegion(SimpleInterface):
 
     def _run_interface(self, runtime):
         seg_img = nib.load(self.inputs.seg_file)
+        gm_prob_img = None
+        if isdefined(self.inputs.gm_probseg):
+            gm_prob_img = nib.load(self.inputs.gm_probseg)
 
         if isdefined(self.inputs.override_indices):
             cfg = {'refmask_indices': list(self.inputs.override_indices)}
@@ -50,7 +54,11 @@ def _run_interface(self, runtime):
 
         from petprep.utils.reference_mask import generate_reference_region
 
-        refmask_img = generate_reference_region(seg_img=seg_img, config=cfg)
+        refmask_img = generate_reference_region(
+            seg_img=seg_img,
+            config=cfg,
+            gm_probseg_img=gm_prob_img,
+        )
 
         out_file = os.path.abspath('refmask.nii.gz')
         nib.save(refmask_img, out_file)

petprep/interfaces/tests/test_reference_mask.py

@@ -116,3 +116,34 @@ def test_extract_refregion_override_missing_config(tmp_path):
     out = nb.load(res.outputs.refmask_file).get_fdata()
     assert out[2, 2, 2] == 1
     assert out.sum() == 1
+
+
+def test_extract_refregion_gm_threshold(tmp_path):
+    data = np.zeros((5, 5, 5), dtype=np.uint8)
+    data[1, 1, 1] = 1
+    data[2, 2, 2] = 1
+    seg = tmp_path / 'seg.nii.gz'
+    nb.Nifti1Image(data, np.eye(4)).to_filename(seg)
+    gm = np.zeros((5, 5, 5), dtype=np.float32)
+    gm[1, 1, 1] = 0.4
+    gm[2, 2, 2] = 0.8
+    gm_file = tmp_path / 'gm.nii.gz'
+    nb.Nifti1Image(gm, np.eye(4)).to_filename(gm_file)
+
+    cfg = _create_config(tmp_path, [1], {'gm_prob_threshold': 0.5})
+
+    node = pe.Node(
+        ExtractRefRegion(
+            seg_file=str(seg),
+            gm_probseg=str(gm_file),
+            config_file=str(cfg),
+            segmentation_type='testseg',
+            region_name='region',
+        ),
+        name='er5',
+        base_dir=str(tmp_path),
+    )
+    res = node.run()
+    out = nb.load(res.outputs.refmask_file).get_fdata()
+    assert out.sum() == 1
+    assert out[2, 2, 2] == 1

petprep/utils/reference_mask.py

@@ -4,7 +4,11 @@
 from skimage.morphology import ball, binary_dilation, binary_erosion
 
 
-def generate_reference_region(seg_img: nib.Nifti1Image, config: dict) -> nib.Nifti1Image:
+def generate_reference_region(
+    seg_img: nib.Nifti1Image,
+    config: dict,
+    gm_probseg_img: nib.Nifti1Image | None = None,
+) -> nib.Nifti1Image:
     """Generate a reference region using a flexible config.
 
     Config keys:
@@ -14,6 +18,8 @@ def generate_reference_region(seg_img: nib.Nifti1Image, config: dict) -> nib.Nif
         - dilate_by_voxels (int, optional): Dilation radius for excluded regions.
         - smooth_fwhm_mm (float, optional): FWHM for smoothing the target region.
         - target_volume_ml (float, optional): Keep only the top N voxels by smoothed value.
+        - gm_prob_threshold (float, optional): Threshold the final mask using a
+          gray matter probability map. Requires ``gm_probseg_img``.
 
     Returns:
         nib.Nifti1Image: Final reference mask.
@@ -52,4 +58,12 @@ def generate_reference_region(seg_img: nib.Nifti1Image, config: dict) -> nib.Nif
             threshold = values[-target_voxels]
         mask = ((smoothed >= threshold) & (mask > 0)).astype(np.uint8)
 
+    # Step 5: Optional gray matter probability thresholding
+    if gm_probseg_img is not None and 'gm_prob_threshold' in config:
+        gm_prob = gm_probseg_img.get_fdata()
+        if gm_prob.shape != mask.shape:
+            raise ValueError('gm_probseg_img does not match segmentation shape')
+        mask = (mask > 0) & (gm_prob >= config['gm_prob_threshold'])
+        mask = mask.astype(np.uint8)
+
     return nib.Nifti1Image(mask, affine, header)

petprep/utils/tests/test_reference_mask.py

@@ -62,3 +62,22 @@ def test_generate_reference_region_large_target_volume():
     out_img = generate_reference_region(img, config)
     mask = out_img.get_fdata()
     assert mask.sum() == seg.sum()
+
+
+def test_generate_reference_region_gm_threshold():
+    seg = np.zeros((5, 5, 5), dtype=np.uint8)
+    seg[2, 2, 2] = 1
+    seg[3, 3, 3] = 1
+    seg_img = nb.Nifti1Image(seg, np.eye(4))
+
+    gm = np.zeros((5, 5, 5), dtype=np.float32)
+    gm[2, 2, 2] = 0.4
+    gm[3, 3, 3] = 0.8
+    gm_img = nb.Nifti1Image(gm, np.eye(4))
+
+    config = {'refmask_indices': [1], 'gm_prob_threshold': 0.5}
+
+    out = generate_reference_region(seg_img, config, gm_probseg_img=gm_img)
+    data = out.get_fdata()
+    assert data.sum() == 1
+    assert data[3, 3, 3] == 1

petprep/workflows/pet/base.py

@@ -243,6 +243,7 @@ def init_pet_wf(
             ('t1w_preproc', 'inputnode.t1w_preproc'),
             ('t1w_mask', 'inputnode.t1w_mask'),
             ('t1w_dseg', 'inputnode.t1w_dseg'),
+            ('t1w_tpms', 'inputnode.t1w_tpms'),
             ('subjects_dir', 'inputnode.subjects_dir'),
             ('subject_id', 'inputnode.subject_id'),
             ('fsnative2t1w_xfm', 'inputnode.fsnative2t1w_xfm'),

petprep/workflows/pet/fit.py

@@ -164,6 +164,7 @@ def init_pet_fit_wf(
                 't1w_preproc',
                 't1w_mask',
                 't1w_dseg',
+                't1w_tpms',
                 'subjects_dir',
                 'subject_id',
                 'fsnative2t1w_xfm',
@@ -448,6 +449,8 @@ def init_pet_fit_wf(
             name='refmask_report_wf',
         )
 
+        gm_select = pe.Node(niu.Select(index=0), name='select_gm_probseg')
+
         pet_ref_tacs_wf = init_pet_ref_tacs_wf(name='pet_ref_tacs_wf')
         pet_ref_tacs_wf.inputs.inputnode.metadata = str(
             Path(pet_file).with_suffix('').with_suffix('.json')
@@ -470,6 +473,8 @@ def init_pet_fit_wf(
         )
         ds_ref_tacs.inputs.source_file = pet_file
 
+        workflow.connect([(inputnode, gm_select, [('t1w_tpms', 'inlist')])])
+
         workflow.connect(
             [
                 (
@@ -479,6 +484,11 @@ def init_pet_fit_wf(
                         ('outputnode.segmentation', 'inputnode.seg_file'),
                     ],
                 ),
+                (
+                    gm_select,
+                    refmask_wf,
+                    [('out', 'inputnode.gm_probseg')],
+                ),
                 (
                     refmask_wf,
                     outputnode,

petprep/workflows/pet/reference_mask.py

@@ -42,7 +42,10 @@ def init_pet_refmask_wf(
 
     workflow = pe.Workflow(name=name)
 
-    inputnode = pe.Node(IdentityInterface(fields=['seg_file']), name='inputnode')
+    inputnode = pe.Node(
+        IdentityInterface(fields=['seg_file', 'gm_probseg']),
+        name='inputnode',
+    )
     outputnode = pe.Node(IdentityInterface(fields=['refmask_file']), name='outputnode')
 
     extract_mask = pe.Node(ExtractRefRegion(), name='extract_refregion')
@@ -56,7 +59,14 @@ def init_pet_refmask_wf(
 
     workflow.connect(
         [
-            (inputnode, extract_mask, [('seg_file', 'seg_file')]),
+            (
+                inputnode,
+                extract_mask,
+                [
+                    ('seg_file', 'seg_file'),
+                    ('gm_probseg', 'gm_probseg'),
+                ],
+            ),
             (extract_mask, outputnode, [('refmask_file', 'refmask_file')]),
         ]
     )

petprep/workflows/pet/tests/test_fit.py

@@ -235,6 +235,10 @@ def test_refmask_report_connections(bids_root: Path, tmp_path: Path):
         'inputnode.source_files',
     ) in petref_edge['connect']
 
+    gm_node = wf.get_node('select_gm_probseg')
+    edge_prob = wf._graph.get_edge_data(gm_node, wf.get_node('pet_refmask_wf'))
+    assert ('out', 'inputnode.gm_probseg') in edge_prob['connect']
+
     assert any(name.startswith('pet_ref_tacs_wf') for name in wf.list_node_names())
     assert 'ds_ref_tacs' in wf.list_node_names()
     ds_tacs = wf.get_node('ds_ref_tacs')