Segmentation features needed for intracranial mask computation

.github/workflows/setup-data.yml

@@ -22,7 +22,7 @@ jobs:
           path: data/  # The folder you want to cache
           # The key determines if we have a match.
           # Change 'v1' to 'v2' manually to force a re-download in the future.
-          key: test-data-v7
+          key: test-data-v9
 
 
         # 2. DOWNLOAD ONLY IF CACHE MISS

src/mritk/cli.py

@@ -9,7 +9,7 @@
 from rich.logging import RichHandler
 from rich_argparse import RichHelpFormatter
 
-from . import concentration, datasets, hybrid, info, looklocker, masks, mixed, napari, r1, show, statistics
+from . import concentration, datasets, hybrid, info, looklocker, masks, mixed, napari, r1, segmentation, show, statistics
 
 
 def version_info():
@@ -75,6 +75,9 @@ def setup_parser():
     napari_parser = subparsers.add_parser("napari", help="Show MRI data using napari", formatter_class=parser.formatter_class)
     napari.add_arguments(napari_parser)
 
+    segmentation_parser = subparsers.add_parser("seg", help="Perform segmentation tasks", formatter_class=parser.formatter_class)
+    segmentation.add_arguments(segmentation_parser, extra_args_cb=add_extra_arguments)
+
     looklocker_parser = subparsers.add_parser(
         "looklocker", help="Process Look-Locker data", formatter_class=parser.formatter_class
     )
@@ -142,6 +145,8 @@ def dispatch(parser: argparse.ArgumentParser, argv: Optional[Sequence[str]] = No
             show.dispatch(args)
         elif command == "napari":
             napari.dispatch(args)
+        elif command == "seg":
+            segmentation.dispatch(args)
         elif command == "looklocker":
             looklocker.dispatch(args)
         elif command == "mask":

src/mritk/masks.py

@@ -12,6 +12,7 @@
 import skimage
 
 from .data import MRIData
+from .segmentation import CSFSegmentation, Segmentation
 from .testing import assert_same_space
 
 
@@ -81,12 +82,12 @@ def compute_csf_mask_array(
     return binary
 
 
-def csf_mask(input: Path, connectivity: int | None = 2, use_li: bool = False) -> MRIData:
+def csf_mask(input: MRIData, connectivity: int | None = 2, use_li: bool = False) -> MRIData:
     """
     I/O wrapper for generating and saving a CSF mask from a NIfTI file.
 
     Args:
-        input (Path): Path to the input NIfTI image.
+        input (MRIData): An MRIData object containing the input volume (typically T2-weighted or Spin-Echo).
         connectivity (Optional[int], optional): Connectivity distance. Defaults to 2.
         use_li (bool, optional): If True, uses Li thresholding. Defaults to False.
         output (Optional[Path], optional): Path to save the resulting mask. Defaults to None.
@@ -97,12 +98,10 @@ def csf_mask(input: Path, connectivity: int | None = 2, use_li: bool = False) ->
     Raises:
         AssertionError: If the resulting mask contains no voxels.
     """
-    input_vol = MRIData.from_file(input, dtype=np.single)
-    mask = compute_csf_mask_array(input_vol.data, connectivity, use_li)
-
+    mask = compute_csf_mask_array(input.data, connectivity, use_li)
     assert np.max(mask) > 0, "Masking failed, no voxels in mask"
 
-    mri_data = MRIData(data=mask, affine=input_vol.affine)
+    mri_data = MRIData(data=mask, affine=input.affine)
 
     return mri_data
 
@@ -134,29 +133,28 @@ def compute_intracranial_mask_array(csf_mask_array: np.ndarray, segmentation_arr
     return ~opened_background
 
 
-def intracranial_mask(csf_segmentation_path: Path, segmentation_path: Path) -> MRIData:
+def intracranial_mask(segmentation: Segmentation, csf_mask: MRIData) -> MRIData:
     """
     I/O wrapper for generating and saving an intracranial mask from NIfTI files.
 
     Loads the masks, verifies they share the same physical coordinate space, and
     delegates the array computation.
 
     Args:
-        csf_segmentation_path (Path): Path to the CSF segmentation NIfTI file.
-        segmentation_path (Path): Path to the brain segmentation NIfTI file.
-        output (Optional[Path], optional): Path to save the resulting mask. Defaults to None.
+        segmentation (MRIData): The refined segmentation (MRIData), \
+            generated by the segmentation refinement module.
+        csf_mask (MRIData): The CSF mask (MRIData), generated by the csf mask module.
 
     Returns:
         MRIData: An MRIData object containing the intracranial mask.
     """
-    input_csf_mask = MRIData.from_file(csf_segmentation_path, dtype=bool)
-    segmentation_data = MRIData.from_file(segmentation_path, dtype=bool)
+    csf_seg = CSFSegmentation(segmentation, csf_mask).to_csf_segmentation()
 
     # Validate spatial alignment before array operations
-    assert_same_space(input_csf_mask, segmentation_data)
+    assert_same_space(csf_seg, segmentation.mri)
 
-    mask_data = compute_intracranial_mask_array(input_csf_mask.data, segmentation_data.data)
-    mri_data = MRIData(data=mask_data, affine=segmentation_data.affine)
+    mask_data = compute_intracranial_mask_array(csf_seg.data, segmentation.mri.data)
+    mri_data = MRIData(data=mask_data, affine=segmentation.mri.affine)
 
     return mri_data
 
@@ -183,8 +181,18 @@ def add_arguments(
     intracranial_mask_parser = subparser.add_parser(
         "intracranial", help="Compute intracranial mask", formatter_class=parser.formatter_class
     )
-    intracranial_mask_parser.add_argument("--csf-segmentation-path", type=Path, help="Path to the CSF segmentation NIfTI file")
-    intracranial_mask_parser.add_argument("--segmentation-path", type=Path, help="Path to the brain segmentation NIfTI file")
+    intracranial_mask_parser.add_argument(
+        "--segmentation-path",
+        type=Path,
+        help="Path to refined segmentation file, generated by \
+        the segmentation refinement module, i.e. mritk seg refine",
+    )
+    intracranial_mask_parser.add_argument(
+        "--csf-mask-path",
+        type=Path,
+        help="Path to the CSF mask NIfTI file, generated by \
+        the csf mask module, i.e. mritk mask csf",
+    )
     intracranial_mask_parser.add_argument("-o", "--output", type=Path, help="Desired output path for the resulting mask")
 
     if extra_args_cb is not None:
@@ -195,11 +203,14 @@ def add_arguments(
 def dispatch(args):
     command = args.pop("mask-command")
     if command == "csf":
-        csf_mask_data = csf_mask(input=args.pop("input"), connectivity=args.pop("connectivity"), use_li=args.pop("use_li"))
+        csf_mask_data = csf_mask(
+            input=MRIData.from_file(args.pop("input")), connectivity=args.pop("connectivity"), use_li=args.pop("use_li")
+        )
         csf_mask_data.save(args.pop("output"), dtype=np.uint8)
     elif command == "intracranial":
         intracranial_mask_data = intracranial_mask(
-            csf_segmentation_path=args.pop("csf_segmentation_path"), segmentation_path=args.pop("segmentation_path")
+            segmentation=MRIData.from_file(args.pop("segmentation_path"), dtype=np.single),
+            csf_mask=MRIData.from_file(args.pop("csf_mask_path"), dtype=np.single),
         )
         intracranial_mask_data.save(args.pop("output"), dtype=np.uint8)
     else: