Commit Message to fix

Author: dkuegler

CorpusCallosum/cc_visualization.py

@@ -93,22 +93,19 @@ def make_parser() -> argparse.ArgumentParser:
         default=0,
         help="Enable verbose (pass twice for debug-output).",
     )
-
     return parser
 
+
 def options_parse() -> argparse.Namespace:
     """Parse command line arguments for the pipeline."""
     parser = make_parser()
     args = parser.parse_args()
 
     # Create output directory if it doesn't exist
     Path(args.output_dir).mkdir(parents=True, exist_ok=True)
-
     return args
 
 
-
-
 def load_contours_from_template_dir(
     template_dir: Path, resolution: float, smoothing_window: int
 ) -> list[CCContour]:
@@ -122,7 +119,6 @@ def load_contours_from_template_dir(
         )
 
     fsaverage_contour = None
-
     contours: list[CCContour] = []
     for thickness_file in thickness_files:
         try:
@@ -176,14 +172,11 @@ def main(
     _, _, vox2ras_tkr = load_fsaverage_data(FSAVERAGE_DATA_PATH)
 
     contours = load_contours_from_template_dir(
-        Path(template_dir), resolution=resolution, smoothing_window=smoothing_window
+        Path(template_dir), resolution=resolution, smoothing_window=smoothing_window,
     )
 
     # 2D visualization
     mid_contour = contours[len(contours) // 2]
-    
-    
-    
 
     # for now, we only support thickness visualization, this is preparing to plot also p-values and icc values
     mode = "thickness"
@@ -215,8 +208,7 @@ def main(
     cc_mesh.plot_mesh(**plot_kwargs)
     cc_mesh.plot_mesh(output_path=str(output_dir / "cc_mesh.html"), **plot_kwargs)
 
-
-    cc_mesh.to_fs_coordinates(vox2ras_tkr=vox2ras_tkr)
+    cc_mesh = cc_mesh.to_fs_coordinates(vox2ras_tkr=vox2ras_tkr)
     logger.info(f"Writing vtk file to {output_dir / 'cc_mesh.vtk'}")
     cc_mesh.write_vtk(str(output_dir / "cc_mesh.vtk"))
     logger.info(f"Writing freesurfer surface file to {output_dir / 'cc_mesh.fssurf'}")

CorpusCallosum/fastsurfer_cc.py

@@ -61,11 +61,11 @@
 from CorpusCallosum.utils.types import CCMeasuresDict, SliceSelection, SubdivisionMethod
 from FastSurferCNN.data_loader.conform import conform, is_conform
 from FastSurferCNN.segstats import HelpFormatter
-from FastSurferCNN.utils import AffineMatrix4x4, Image3d, Mask3d, Shape3d, Vector2d, logging, nibabelImage
+from FastSurferCNN.utils import AffineMatrix4x4, Image3d, Mask3d, Shape3d, Vector2d, logging, nibabelImage, Image4d
 from FastSurferCNN.utils.arg_types import path_or_none
 from FastSurferCNN.utils.common import SubjectDirectory, find_device
 from FastSurferCNN.utils.lta import write_lta
-from FastSurferCNN.utils.parallel import shutdown_executors, thread_executor
+from FastSurferCNN.utils.parallel import shutdown_executors, thread_executor, get_num_threads, serial_executor
 from FastSurferCNN.utils.parser_defaults import modify_argument
 from recon_surf.align_points import find_rigid
 
@@ -163,7 +163,7 @@ def _set_help_sid(action):
              "cost of precision.",
     )
     def _slice_selection(a: str) -> SliceSelection:
-        if b := a.lower() in ("middle", "all"):
+        if (b := a.lower()) in ("middle", "all"):
             return b
         return int(a)
     parser.add_argument(
@@ -494,7 +494,7 @@ def segment_cc(
     pc_coords: Vector2d,
     aseg_nib: nibabelImage,
     model_segmentation: "torch.nn.Module",
-) -> tuple[Mask3d, Image3d]:
+) -> tuple[Mask3d, Image4d]:
     """Segment the corpus callosum using a trained model.
 
     Performs corpus callosum segmentation on mid-sagittal slices using a trained model, with AC-PC points as anatomical
@@ -518,7 +518,7 @@ def segment_cc(
     cc_seg_labels : np.ndarray
         Binary cc_seg_labels of the corpus callosum.
     cc_softlabels : np.ndarray
-        Soft cc_seg_labels probabilities.
+        Soft cc_seg_labels probabilities of shape (H, W, D, C=3).
     """
     pre_clean_segmentation, inputs, cc_softlabels = segmentation_inference.run_inference_on_slice(
         model_segmentation,
@@ -732,7 +732,7 @@ def main(
         sys.exit(1)
 
     logger.info("Performing centroid registration to fsaverage space")
-    orig2fsavg_vox2vox, orig2fsavg_ras2ras, fsavg_vox2ras, fsavg_header, fsavg_vox2ras_tkr = (
+    orig2fsavg_vox2vox, orig2fsavg_ras2ras, fsavg_vox2ras, fsavg_header, _ = (
         register_centroids_to_fsavg(aseg_img)
     )
 
@@ -754,11 +754,11 @@ def main(
     affine_x_offset = partial(create_sag_slice_vox2vox, fsaverage_middle=FSAVERAGE_MIDDLE / vox_size[0])
     fsavg2midslab_in_vox2vox: AffineMatrix4x4 = affine_x_offset(slices_to_analyze // 2)
     # first, midslice->fsaverage in vox2vox, then vox2ras in fsaverage space
-    midslab_vox2ras: AffineMatrix4x4 = fsavg_vox2ras @ np.linalg.inv(fsavg2midslab_in_vox2vox)
+    fsaverage_midslab_vox2ras: AffineMatrix4x4 = fsavg_vox2ras @ np.linalg.inv(fsavg2midslab_in_vox2vox)
 
     # calculate vox2vox for input resampling volumes
-    def _orig2midslab_vox2vox(extra_slices: int) -> AffineMatrix4x4:
-        fsavg2midslab = affine_x_offset(slices_to_analyze // 2 + extra_slices //  2)
+    def _orig2midslab_vox2vox(additional_context: int = 0) -> AffineMatrix4x4:
+        fsavg2midslab = affine_x_offset(slices_to_analyze // 2 + additional_context // 2)
         # first, orig->fsaverage in vox2vox, then fsaverage->midslab in vox2vox
         return fsavg2midslab @ orig2fsavg_vox2vox
 
@@ -769,16 +769,16 @@ def _orig2midslab_vox2vox(extra_slices: int) -> AffineMatrix4x4:
     ac_coords, pc_coords = localize_ac_pc(
         np.asarray(orig.dataobj),
         aseg_img,
-        _orig2midslab_vox2vox(extra_slices=2),
+        _orig2midslab_vox2vox(additional_context=2),
         _model_localization.result(),
         target_shape,
     )
     logger.info("Starting corpus callosum segmentation")
-    # "+ 8" in x-direction for context slices
-    target_shape: Shape3d = (slices_to_analyze + 8, fsavg_header["dims"][1], fsavg_header["dims"][2])
+    extra_slices = 8  # 8 extra in x-direction for context slices
+    target_shape: Shape3d = (slices_to_analyze + extra_slices, fsavg_header["dims"][1], fsavg_header["dims"][2])
     midslices: Image3d = affine_transform(
         np.asarray(orig.dataobj),
-        np.linalg.inv(_orig2midslab_vox2vox(extra_slices=8)), # inverse is required for affine_transform
+        np.linalg.inv(_orig2midslab_vox2vox(additional_context=extra_slices)), # inverse is required for affine_transform
         output_shape=target_shape,
         order=2, # @ClePol unclear, why this is not order=3
         mode="constant",
@@ -799,7 +799,7 @@ def _orig2midslab_vox2vox(extra_slices: int) -> AffineMatrix4x4:
             logger.info(f"Saving {name} softlabels to {sd.filename_by_attribute(f'cc_softlabels_{attr}')}")
             io_futures.append(thread_executor().submit(
                 nib.save,
-                nib.MGHImage(cc_fn_softlabels[..., i], midslab_vox2ras, orig.header),
+                nib.MGHImage(cc_fn_softlabels[..., i], fsaverage_midslab_vox2ras, orig.header),
                 sd.filename_by_attribute(f"cc_softlabels_{attr}"),
             ))
 
@@ -819,7 +819,6 @@ def _orig2midslab_vox2vox(extra_slices: int) -> AffineMatrix4x4:
         subdivision_method=subdivision_method,
         contour_smoothing=contour_smoothing,
         vox_size=vox_size,
-        vox2ras_tkr=fsavg_vox2ras_tkr,
         subject_dir=sd,
     )
     io_futures.extend(slice_io_futures)
@@ -837,7 +836,7 @@ def _orig2midslab_vox2vox(extra_slices: int) -> AffineMatrix4x4:
         cc_subseg_midslice = make_subdivision_mask(
             (cc_fn_seg_labels.shape[1], cc_fn_seg_labels.shape[2]),
             middle_slice_result["split_contours"],
-            vox_size[1:],
+            vox_size[1:3],
         )
     else:
         logger.warning("Too many subsegments for lookup table, skipping sub-division of output segmentation.")
@@ -847,19 +846,21 @@ def _orig2midslab_vox2vox(extra_slices: int) -> AffineMatrix4x4:
     if sd.has_attribute("cc_segmentation"):
         io_futures.append(thread_executor().submit(
             nib.save,
-            nib.MGHImage(cc_fn_seg_labels, midslab_vox2ras, orig.header),
+            nib.MGHImage(cc_fn_seg_labels, fsaverage_midslab_vox2ras, orig.header),
             sd.filename_by_attribute("cc_segmentation"),
         ))
     # map soft labels to original space (in parallel because this takes a while, and we only do it to save the labels)
     if sd.has_attribute("cc_orig_segfile"):
-        io_futures.append(thread_executor().submit(
+        # if num_threads is not large enough (>1), this might be blocking ; serial_executor runs the function in submit
+        executor = thread_executor() if get_num_threads() > 2 else serial_executor()
+        io_futures.append(executor.submit(
             map_softlabels_to_orig,
             cc_fn_softlabels=cc_fn_softlabels,
-            orig_fsaverage_vox2vox=orig2fsavg_vox2vox,
             orig=orig,
             orig_space_segmentation_path=sd.filename_by_attribute("cc_orig_segfile"),
-            fsaverage_middle=FSAVERAGE_MIDDLE,
+            orig2slab_vox2vox=_orig2midslab_vox2vox(),
             cc_subseg_midslice=cc_subseg_midslice,
+            orig2midslice_vox2vox=affine_x_offset(0) @ orig2fsavg_vox2vox,  # orig2fsavg, then full2midslice
         ))
 
     METRICS = [
@@ -1004,6 +1005,7 @@ def save_cc_measures_json(cc_mid_measure_file: Path, metrics: dict[str, object])
         conf_name=options.conf_name,
         aseg_name=options.aseg_name,
         subject_dir=options.subject_dir,
+        #FIXME: slice_selection is True/bool
         slice_selection=options.slice_selection,
         num_thickness_points=options.num_thickness_points,
         subdivisions=list(options.subdivisions),

CorpusCallosum/localization/inference.py

@@ -188,7 +188,7 @@ def predict(
         outputs = model(inputs) * torch.as_tensor([PATCH_SIZE + PATCH_SIZE], device=device)
 
     t_crops = [(t_dict['crop_left'] + t_dict['crop_top']) * 2]
-    outs: np.ndarray[tuple[int, Literal[4]], np.dtype[float]] = outputs.cpu().numpy() + np.asarray(t_crops, dtype=float)
+    outs: np.ndarray[tuple[int, Literal[4]], np.dtype[np.float_]] = outputs.cpu().numpy() + np.asarray(t_crops, dtype=float)
     crop_offsets: tuple[int, int] = (t_dict["crop_left"][0], t_dict["crop_top"][0])
     return outs[:, :2], outs[:, 2:], crop_offsets
 

CorpusCallosum/segmentation/inference.py

@@ -92,7 +92,7 @@ def run_inference(
     voxel_size: tuple[float, float],
     device: torch.device | None = None,
     transform: transforms.Transform | None = None
-) -> tuple[np.ndarray[Shape4d, np.dtype[int]], Image4d, Image4d]:
+) -> tuple[np.ndarray[Shape4d, np.dtype[np.int_]], Image4d, Image4d]:
     """Run inference on a single image slice.
 
     Parameters
@@ -229,15 +229,15 @@ def _load(label_path: str | Path) -> int:
     return images, ac_centers, pc_centers, label_widths, labels, subj_ids
 
 @overload
-def one_hot_to_label(one_hot: Image4d, label_ids: list[int] | None = None) -> np.ndarray[Shape3d, np.dtype[int]]: ...
+def one_hot_to_label(one_hot: Image4d, label_ids: list[int] | None = None) -> np.ndarray[Shape3d, np.dtype[np.int_]]: ...
 
 @overload
-def one_hot_to_label(one_hot: Image3d, label_ids: list[int] | None = None) -> np.ndarray[Shape2d, np.dtype[int]]: ...
+def one_hot_to_label(one_hot: Image3d, label_ids: list[int] | None = None) -> np.ndarray[Shape2d, np.dtype[np.int_]]: ...
 
 def one_hot_to_label(
-    one_hot: np.ndarray[tuple[int, ...], np.dtype[bool]],
+    one_hot: np.ndarray[tuple[int, ...], np.dtype[np.bool_]],
     label_ids: list[int] | None = None,
-) -> np.ndarray[tuple[int, ...], np.dtype[int]]:
+) -> np.ndarray[tuple[int, ...], np.dtype[np.int_]]:
     """Convert one-hot encoded segmentation to label map.
 
     Converts a one-hot encoded segmentation array to discrete labels by taking
@@ -273,7 +273,7 @@ def run_inference_on_slice(
         ac_center: Vector2d,
         pc_center: Vector2d,
         voxel_size: tuple[float, float],
-) -> tuple[np.ndarray[Shape3d, np.dtype[int]], Image4d, Image4d]:
+) -> tuple[np.ndarray[Shape3d, np.dtype[np.int_]], Image4d, Image4d]:
     """Run inference on a single slice.
 
     Parameters

CorpusCallosum/segmentation/segmentation_postprocessing.py

@@ -500,9 +500,9 @@ def extract_largest_connected_component(
 
 
 def clean_cc_segmentation(
-    seg_arr: np.ndarray[Shape3d, np.dtype[int]],
+    seg_arr: np.ndarray[Shape3d, np.dtype[np.int_]],
     max_connection_distance: float = 3.0,
-) -> tuple[np.ndarray[Shape3d, np.dtype[int]], Mask3d]:
+) -> tuple[np.ndarray[Shape3d, np.dtype[np.int_]], Mask3d]:
     """Clean corpus callosum segmentation by removing non-connected components.
 
     Parameters

CorpusCallosum/shape/contour.py

@@ -73,8 +73,8 @@ class CCContour:
 
     def __init__(
         self,
-        contour: np.ndarray[tuple[Literal["N", 2]], np.dtype[float]],
-        thickness_values: np.ndarray[tuple[Literal["N"]], np.dtype[float]],
+        contour: np.ndarray[tuple[Literal["N", 2]], np.dtype[np.float_]],
+        thickness_values: np.ndarray[tuple[Literal["N"]], np.dtype[np.float_]],
         endpoint_idxs: tuple[int, int] | None = None,
         resolution: float = 1.0
     ):
@@ -86,8 +86,10 @@ def __init__(
             Array of shape (N, 2) containing 2D contour points.
         thickness_values : np.ndarray
             Array of thickness measurements for each contour point.
-        endpoint_idxs : tuple[int, int]
+        endpoint_idxs : tuple[int, int], optional
             Tuple containing start and end indices for the contour.
+        resolution : float, default=1.0
+            The left-right spacing.
         """
         self.contour = contour
         if self.contour.shape[1] != 2:

CorpusCallosum/shape/endpoint_heuristic.py

@@ -19,7 +19,7 @@
 import skimage.measure
 from scipy.ndimage import label
 
-from FastSurferCNN.utils import Vector2d
+from FastSurferCNN.utils import Vector2d, Mask2d
 
 
 def smooth_contour(x: np.ndarray, y: np.ndarray, window_size: int) -> tuple[np.ndarray, np.ndarray]:
@@ -169,7 +169,7 @@ def extract_cc_contour(cc_mask: np.ndarray, contour_smoothing: int = 5) -> np.nd
 
 @overload
 def get_endpoints(
-    cc_mask: np.ndarray[tuple[int, int], np.dtype[bool]],
+    cc_mask: Mask2d,
     ac_2d: Vector2d,
     pc_2d: Vector2d,
     resolution: tuple[float, float],
@@ -180,7 +180,7 @@ def get_endpoints(
 
 @overload
 def get_endpoints(
-    cc_mask: np.ndarray[tuple[int, int], np.dtype[bool]],
+    cc_mask: Mask2d,
     ac_2d: Vector2d,
     pc_2d: Vector2d,
     resolution: tuple[float, float],
@@ -190,7 +190,7 @@ def get_endpoints(
 
 
 def get_endpoints(
-    cc_mask: np.ndarray[tuple[int, int], np.dtype[bool]],
+    cc_mask: Mask2d,
     ac_2d: Vector2d,
     pc_2d: Vector2d,
     resolution: tuple[float, float],