Various documentation and formatting changes as well as optimizations during review

Author: dkuegler

CorpusCallosum/README.md

@@ -1,7 +1,7 @@
 # Corpus Callosum Pipeline
 
 A deep learning-based pipeline for automated segmentation, analysis, and shape analysis of the corpus callosum in brain MRI scans.
-Also segments the fornix, localizes the AC and PC and standardizes the orientation of the brain.
+Also segments the fornix, localizes the anterior and posterior commissure (AC and PC) and standardizes the orientation of the brain.
 
 ## Overview
 
@@ -15,7 +15,9 @@ This pipeline combines localization and segmentation deep learning models to:
 
 ## Quickstart
 
-``` python3 fastsurfer_cc.py --subject_dir /path/to/fastsurfer/output --verbose ```
+```bash
+python3 fastsurfer_cc.py --subject_dir /path/to/fastsurfer/output --verbose
+` ``
 
 Gives all standard outputs. Then corpus callosum morphometry can be found at `stats/callosum.CC.midslice.json`, including 100 thickness measurements and areas of sub-segments.
 Visualization will be placed in `/path/to/fastsurfer/output/qc_snapshots`. For more detailed info see the following sections.

CorpusCallosum/__init__.py

@@ -13,11 +13,8 @@
 # limitations under the License.
 
 __all__ = [
-    "config",
     "data",
-    "localization",
     "segmentation",
     "transforms",
     "utils",
-    "visualization",
 ]

CorpusCallosum/cc_visualization.py

@@ -1,5 +1,7 @@
 import argparse
+import sys
 from pathlib import Path
+from typing import Literal
 
 import numpy as np
 
@@ -9,8 +11,8 @@
 from CorpusCallosum.shape.cc_mesh import CC_Mesh
 
 
-def options_parse() -> argparse.Namespace:
-    """Parse command line arguments for the visualization pipeline."""
+def make_parser() -> argparse.ArgumentParser:
+    """Create a command line parser for the visualization pipeline."""
     parser = argparse.ArgumentParser(description="Visualize corpus callosum from template files.")
     parser.add_argument("--contours", type=str, required=False, help="Path to contours.txt file", default=None)
     parser.add_argument("--thickness", type=str, required=True, help="Path to thickness_values.txt file")
@@ -38,12 +40,17 @@ def options_parse() -> argparse.Namespace:
         nargs=2,
         default=None,
         metavar=("MIN", "MAX"),
-        help="Optional fixed range for the colorbar (min max)",
+        required=False,
+        help="Specify the range for the colorbar (2 values: min max). Defaults to automatic choice.",
     )
     parser.add_argument("--legend", type=str, default="Thickness (mm)", help="Legend for the colorbar")
     parser.add_argument("--twoD", action="store_true", help="Generate 2D visualization instead of 3D mesh")
 
-    args = parser.parse_args()
+    return parser
+
+def options_parse() -> argparse.Namespace:
+    """Parse command line arguments for the pipeline."""
+    args = make_parser().parse_args()
 
     # Create output directory if it doesn't exist
     Path(args.output_dir).mkdir(parents=True, exist_ok=True)
@@ -62,7 +69,7 @@ def main(
     color_range: tuple[float, float] | None = None,
     legend: str | None = None,
     twoD: bool = False,
-) -> None:
+) -> Literal[0] | str:
     """Main function to visualize corpus callosum from template files.
 
     This function loads contours and thickness values from template files,
@@ -146,21 +153,24 @@ def main(
         cc_mesh.write_vtk(str(output_dir / "cc_mesh.vtk"))
         cc_mesh.write_fssurf(str(output_dir / "cc_mesh.fssurf"))
         cc_mesh.write_overlay(str(output_dir / "cc_mesh_overlay.curv"))
-        cc_mesh.snap_cc_picture(str(output_dir / "cc_mesh_snap.png"))
-
+        try:
+            cc_mesh.snap_cc_picture(str(output_dir / "cc_mesh_snap.png"))
+        except RuntimeError:
+            return ("The cc_visualization script requires whippersnappy>=1.3.1 to makes screenshots, install with "
+                    "`pip install whippersnappy>=1.3.1` !")
+    return 0
 
 if __name__ == "__main__":
-    options = options_parse()
-    main_args = {
-        "contours_path": options.contours,
-        "thickness_path": options.thickness,
-        "measurement_points_path": options.measurement_points,
-        "output_dir": options.output_dir,
-        "resolution": options.resolution,
-        "smoothing_window": options.smoothing_window,
-        "colormap": options.colormap,
-        "color_range": options.color_range,
-        "legend": options.legend,
-        "twoD": options.twoD,
-    }
-    main(**main_args)
+    options = make_parser().parse_args()
+    sys.exit(main(
+        contours_path=options.contours,
+        thickness_path=options.thickness,
+        measurement_points_path=options.measurement_points,
+        output_dir=options.output_dir,
+        resolution=options.resolution,
+        smooth_iterations=options.smooth_iterations,
+        colormap=options.colormap,
+        color_range=options.color_range,
+        legend=options.legend,
+        twoD=options.twoD,
+    ))

CorpusCallosum/data/constants.py

@@ -16,13 +16,13 @@
 from pathlib import Path
 
 ### Constants
-WEIGHTS_PATH = Path(__file__).parent.parent / "weights"
+WEIGHTS_PATH = Path(__file__).parents[1] / "weights"
 FSAVERAGE_CENTROIDS_PATH = Path(__file__).parent / "fsaverage_centroids.json"
 FSAVERAGE_DATA_PATH = Path(__file__).parent / "fsaverage_data.json"  # Contains both affine and header
 FSAVERAGE_MIDDLE = 128  # Middle slice index in fsaverage space
 CC_LABEL = 192          # Label value for corpus callosum in segmentation
 FORNIX_LABEL = 250      # Label value for fornix in segmentation
-SUBSEGEMNT_LABELS = [251, 252, 253, 254, 255] # labels for subsegments in segmentation
+SUBSEGMENT_LABELS = [251, 252, 253, 254, 255] # labels for subsegments in segmentation
 
 
 STANDARD_OUTPUT_PATHS = {

CorpusCallosum/data/fsaverage_cc_template.py

@@ -19,7 +19,9 @@
 import numpy as np
 from scipy import ndimage
 
+from CorpusCallosum.data import constants
 from CorpusCallosum.shape.cc_postprocessing import process_slice
+from FastSurferCNN.utils.brainvolstats import mask_in_array
 
 
 def smooth_contour(contour: tuple[np.ndarray, np.ndarray], window_size: int = 5) -> tuple[np.ndarray, np.ndarray]:
@@ -93,19 +95,15 @@ def load_fsaverage_cc_template() -> tuple[
 
     fsaverage_seg_path = freesurfer_home / 'subjects' / 'fsaverage' / 'mri' / 'aparc+aseg.mgz'
     fsaverage_seg = nib.load(fsaverage_seg_path)
-    segmentation = fsaverage_seg.get_fdata()
+    segmentation = np.asarray(fsaverage_seg.dataobj)
 
     PC = np.array([131, 99])
     AC = np.array([135, 130])
 
 
     midslice = segmentation.shape[0]//2 +1
 
-    cc_mask = segmentation[midslice] == 251
-    cc_mask |= segmentation[midslice] == 252
-    cc_mask |= segmentation[midslice] == 253
-    cc_mask |= segmentation[midslice] == 254
-    cc_mask |= segmentation[midslice] == 255
+    cc_mask = mask_in_array(segmentation[midslice], constants.SUBSEGMENT_LABELS)
 
     # Smooth the CC mask to reduce noise and irregularities
 
@@ -120,8 +118,7 @@ def load_fsaverage_cc_template() -> tuple[
     cc_mask_smoothed = cc_mask_smoothed > 0.5
 
     # Use the smoothed mask for further processing
-    cc_mask = cc_mask_smoothed.astype(int)
-    cc_mask[cc_mask > 0] = 192
+    cc_mask = cc_mask_smoothed.astype(int) * 192
 
     (_, contour_with_thickness, anterior_endpoint_idx,
      posterior_endpoint_idx) = process_slice(segmentation=cc_mask[None],

CorpusCallosum/data/generate_fsaverage_centroids.py

@@ -163,4 +163,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    main()

CorpusCallosum/data/read_write.py

@@ -15,6 +15,7 @@
 import json
 import multiprocessing
 from pathlib import Path
+from typing import overload
 
 import nibabel as nib
 import numpy as np
@@ -52,8 +53,15 @@ def run_in_background(function: callable, debug: bool = False, *args, **kwargs)
     return process
 
 
+@overload
+def get_centroids_from_nib(seg_img: nib.Nifti1Image, label_ids: None = None) -> dict[int, np.ndarray]:
+    ...
 
-def get_centroids_from_nib(seg_img: nib.Nifti1Image, label_ids: list[int] | None = None) -> dict[int, np.ndarray]:
+@overload
+def get_centroids_from_nib(seg_img: nib.Nifti1Image, label_ids: list[int]) -> tuple[dict[int, np.ndarray], list[int]]:
+    ...
+
+def get_centroids_from_nib(seg_img: nib.Nifti1Image, label_ids: list[int] | None = None):
     """Get centroids of segmentation labels in RAS coordinates.
 
     Parameters
@@ -187,12 +195,7 @@ def load_fsaverage_centroids(centroids_path: str | Path) -> dict[int, np.ndarray
         centroids_data = json.load(f)
 
     # Convert string keys back to integers and lists back to numpy arrays
-    centroids = {}
-    for label_str, centroid_list in centroids_data.items():
-        label_id = int(label_str)
-        centroids[label_id] = np.array(centroid_list)
-    
-    return centroids
+    return {int(label): np.array(centroid) for label, centroid in centroids_data.items()}
 
 
 def load_fsaverage_affine(affine_path: str | Path) -> np.ndarray:
@@ -270,8 +273,8 @@ def load_fsaverage_data(data_path: str | Path) -> tuple[np.ndarray, dict, np.nda
     if "header" not in data:
         raise ValueError("Required field 'header' missing from data file")
 
-    header_fields = ["dims", "delta", "Mdc", "Pxyz_c"]
-    for field in header_fields:
+    required_header_fields = ["dims", "delta", "Mdc", "Pxyz_c"]
+    for field in required_header_fields:
         if field not in data["header"]:
             raise ValueError(f"Required header field missing: {field}")