Merge pull request #1194 from AntoineTheb/atheb/bundleparc

ENH: BundleParc - continuous and distance-based parcellation

Author: arnaudbore

src/scilpy/cli/scil_fodf_bundleparc.py

@@ -5,7 +5,7 @@
 
 This method takes as input fODF maps and outputs 71 bundle label maps. These maps can then be used to perform tractometry/tract profiling/radiomics. The bundle definitions follow TractSeg's minus the whole CC.
 
-Inputs are presumed to come from Tractoflow and must be BET and cropped. fODFs must be of basis descoteaux07 and can be of order < 8 but accuracy may be reduced.
+**IMPORTANT**: fODF inputs are presumed to come from Tractoflow, must have stride of -1,2,3,4 and must be BET and cropped. fODFs must be in SH format, basis descoteaux07_legacy and can be of order < 8 but accuracy may be reduced.
 
 Model weights will be downloaded the first time the script is run, which will require an internet connection at runtime. Otherwise they can be manually downloaded from zenodo [1] and by specifying --checkpoint.
 
@@ -19,15 +19,16 @@
 
 The default value of 50 for --min_blob_size was found empirically on adult brains at a resolution of 1mm^3. The best value for your dataset may differ.
 
-This script requires a GPU with ~6GB of available memory. If you use half-precision (float16) inference, you may be able to run it with ~3GB of GPU memory available. Otherwise, install the CPU version of PyTorch.
+This script requires a GPU with ~6GB of available memory. If you use half-precision (float16) inference, you may be able to run it with ~3GB of GPU memory available. Otherwise, install the CPU version of PyTorch. Execution on MacOS is not supported for now.
 
 Parts of the implementation are based on or lifted from:
     SAM-Med3D: https://github.com/uni-medical/SAM-Med3D
     Multidimensional Positional Encoding: https://github.com/tatp22/multidim-positional-encoding
 
 To cite: Antoine Théberge, Zineb El Yamani, François Rheault, Maxime Descoteaux, Pierre-Marc Jodoin (2025). LabelSeg. ISMRM Workshop on 40 Years of Diffusion: Past, Present & Future Perspectives, Kyoto, Japan.
 
-[1]: https://zenodo.org/records/15579498
+[1]: Descoteaux, M., Deriche, R., Knösche, T. R., & Anwander, A. (2007). Deterministic and probabilistic tractography based on complex fibre orientation distributions. IEEE Transactions on Medical Imaging, 26(11), 1464-1477.
+[2]: https://zenodo.org/records/15579498
 """  # noqa
 
 import argparse
@@ -37,15 +38,18 @@
 import os
 
 from argparse import RawTextHelpFormatter
+from functools import partial
 
 from scilpy.io.utils import (
     assert_inputs_exist, assert_output_dirs_exist_and_empty,
     add_overwrite_arg, add_verbose_arg)
 from scilpy.image.volume_operations import resample_volume
 
 from scilpy.ml.bundleparc.predict import predict
+from scilpy.ml.bundleparc.labels import post_process_labels_discrete, \
+    post_process_labels_mm, post_process_labels_continuous
 from scilpy.ml.bundleparc.utils import DEFAULT_BUNDLES, \
-    download_weights, get_model
+     download_weights, get_model
 from scilpy.ml.utils import get_device, IMPORT_ERROR_MSG
 from scilpy import SCILPY_HOME
 
@@ -62,20 +66,11 @@ def _build_arg_parser():
         formatter_class=RawTextHelpFormatter)
 
     parser.add_argument('in_fodf',
-                        help='fODF input.')
+                        help='Input fODF volume in nifti format. ')
     parser.add_argument('--out_prefix', default='',
                         help='Output file prefix. Default is nothing. ')
     parser.add_argument('--out_folder', default='bundleparc',
                         help='Output destination. Default is [%(default)s].')
-    parser.add_argument('--nb_pts', type=int, default=50,
-                        help='Number of divisions per bundle. '
-                             'Default is [%(default)s].')
-    parser.add_argument('--min_blob_size', type=int, default=50,
-                        help='Minimum blob size (in voxels) to keep. Smaller '
-                             'blobs will be removed. Default is '
-                             '[%(default)s].')
-    parser.add_argument('--keep_biggest_blob', action='store_true',
-                        help='If set, only keep the biggest blob predicted.')
     parser.add_argument('--half_precision', action='store_true',
                         help='Use half precision (float16) for inference. '
                              'This reduces memory usage but may lead to '
@@ -88,6 +83,24 @@ def _build_arg_parser():
                              'and weights of model. Default is '
                              '[%(default)s]. If the file does not exist, it '
                              'will be downloaded.')
+    parcel_group = parser.add_mutually_exclusive_group()
+    parcel_group.add_argument('--nb_pts', type=int, default=10,
+                              help='Number of divisions per bundle. Default is'
+                                   ' [%(default)s].')
+    parcel_group.add_argument('--mm', type=float,
+                              help='If set, bundles will be split in sections '
+                                   'roughly X mm wide.')
+    parcel_group.add_argument('--continuous', action='store_true',
+                              help='If set, the output label maps will be '
+                                   'continuous ∈ [0, 1].')
+    blob_group = parser.add_mutually_exclusive_group()
+    blob_group.add_argument('--min_blob_size', type=int, default=50,
+                            help='Minimum blob size (in voxels) to keep. '
+                                 'Smaller blobs will be removed. Default is '
+                                 '[%(default)s].')
+    blob_group.add_argument('--keep_biggest_blob', action='store_true',
+                            help='Only keep the biggest blob predicted.')
+
     add_overwrite_arg(parser)
     add_verbose_arg(parser)
 
@@ -99,7 +112,7 @@ def main():
     parser = _build_arg_parser()
     args = parser.parse_args()
 
-    assert_inputs_exist(parser, [args.in_fodf])
+    assert_inputs_exist(parser, [args.in_fodf], [])
     assert_output_dirs_exist_and_empty(parser, args, args.out_folder,
                                        create_dir=True)
 
@@ -135,19 +148,33 @@ def main():
                                     interp='lin',
                                     enforce_dimensions=False)
 
+    # Get the voxel size of the input fODF after resampling
+    # Presuming isotropic resampling
+    voxel_size = np.mean(resampled_img.header.get_zooms()[:3])
+
+    # Get the label function to use for post-processing
+    if args.continuous:
+        label_function = post_process_labels_continuous
+    elif args.mm is not None:
+        label_function = partial(post_process_labels_mm, args.mm, voxel_size)
+    else:
+        label_function = partial(post_process_labels_discrete,
+                                 args.nb_pts)
+
     # Predict label maps. `predict` is a generator
     # yielding one label map per bundle and its name.
     for y_hat_label, b_name in predict(
-        model, resampled_img.get_fdata(dtype=np.float32), n_coefs, args.nb_pts,
-        args.bundles, args.min_blob_size, args.keep_biggest_blob,
-        args.half_precision, logging.getLogger().getEffectiveLevel() <
-        logging.WARNING
+        model, resampled_img.get_fdata(dtype=np.float32), n_coefs,
+        label_function, DEFAULT_BUNDLES, args.keep_biggest_blob,
+        args.half_precision,
+        logging.getLogger().getEffectiveLevel() < logging.WARNING
     ):
 
         # Format the output as a nifti image
         label_img = nib.Nifti1Image(y_hat_label,
                                     resampled_img.affine,
-                                    resampled_img.header, dtype=np.uint16)
+                                    resampled_img.header,
+                                    dtype=y_hat_label.dtype)
 
         # Resampling volume to fit the original image size
         resampled_label = resample_volume(label_img, ref_img=None,

src/scilpy/cli/tests/test_fodf_bundleparc.py

@@ -53,3 +53,23 @@ def test_execution_invalid_bundle(script_runner, monkeypatch):
     ret = script_runner.run(['scil_fodf_bundleparc', in_fodf,
                              '-f', '--bundles', 'CC'])
     assert not ret.success
+
+
+@pytest.mark.ml
+def test_execution_mm(script_runner, monkeypatch):
+    in_fodf = os.path.join(SCILPY_HOME, 'tracking', 'fodf.nii.gz')
+
+    ret = script_runner.run('scil_fodf_bundleparc.py', in_fodf,
+                            '--mm', '10',
+                            '--bundles', 'IFO_right', '-f')
+    assert ret.success
+
+
+@pytest.mark.ml
+def test_execution_cont(script_runner, monkeypatch):
+    in_fodf = os.path.join(SCILPY_HOME, 'tracking', 'fodf.nii.gz')
+
+    ret = script_runner.run('scil_fodf_bundleparc.py', in_fodf,
+                            '--continuous',
+                            '--bundles', 'IFO_right', '-f')
+    assert ret.success

src/scilpy/ml/bundleparc/labels.py

@@ -0,0 +1,156 @@
+import logging
+import numpy as np
+
+from dipy.tracking.streamline import length, set_number_of_points
+
+from scilpy.tractograms.streamline_operations import smooth_line_gaussian
+
+
+def post_process_labels_discrete(
+    nb_labels, bundle_label, bundle_mask, bundle_name
+):
+    """ Discretize the labels and apply a mask to the bundle. Labels are
+    discretized to integers in the range [1, nb_labels] uniformly.
+
+    Parameters
+    ----------
+    nb_labels : int
+        Number of labels to discretize to.
+    bundle_label : np.ndarray
+        Predicted continuous labels for the bundle.
+    bundle_mask : np.ndarray
+        Binary mask of the bundle.
+    bundle_name : str
+        Name of the bundle, used for logging.
+
+    Returns
+    -------
+    bundle_label : np.ndarray
+        Predicted labels for the bundle.
+    """
+
+    # Determine the output type based on the number of labels
+    # In scilpy/MI-Brain, uint16 is used for labels, uint8 for binary masks.
+    out_type = np.uint16 if nb_labels > 1 else np.uint8
+
+    # Label masking
+    discrete_labels = bundle_label[bundle_mask.astype(bool)]
+
+    # Label dicretizing
+    discrete_labels = np.ceil(discrete_labels * nb_labels)
+    bundle_label[bundle_mask.astype(bool)] = discrete_labels
+    bundle_label[~bundle_mask.astype(bool)] = 0
+
+    return bundle_label.astype(out_type)
+
+
+def post_process_labels_mm(
+    labels_mm, voxel_size, bundle_label, bundle_mask, bundle_name
+):
+    """ Discretize the labels and apply a mask to the bundle. Labels are
+    discritezed to integers so that each section is roughly `labels_mm` mm long
+    To do so, the barycenter of each label is computed to form a centroid
+    streamline. Then, the centroid is resampled to have a number of points such
+    that the step-size is roughly `labels_mm` mm. Finally, the labels are
+    reassigned to the closest point in the resampled centroid.
+
+    Parameters
+    ----------
+    labels_mm : float
+        Length of each section in mm.
+    voxel_size : np.ndarray
+        Voxel size of the bundle image.
+    bundle_label : np.ndarray
+        Predicted continuous labels for the bundle.
+    bundle_mask : np.ndarray
+        Binary mask of the bundle.
+    bundle_name : str
+        Name of the bundle, used for logging.
+
+    Returns
+    -------
+    bundle_label : np.ndarray
+        Predicted labels for the bundle.
+    """
+
+    # Label masking
+    bundle_label[~bundle_mask.astype(bool)] = 0
+
+    ref_labels = np.ceil(bundle_label * 50)
+    unique = np.unique(ref_labels)
+
+    # Get the 3D coordinates of the barycenter of each label
+    barycenters = np.zeros((len(unique) - 1, 3), dtype=np.float32)
+    for i, label in enumerate(unique[1:]):
+        coords = np.argwhere(ref_labels == label)
+        barycenters[i] = np.mean(coords, axis=0)
+
+    # Form the barycenters into a single streamline
+    centroid = np.asarray(barycenters)
+    centroid = smooth_line_gaussian(centroid, 5)
+
+    # Resampling
+    c_length = length(centroid * voxel_size)
+    # Calculate the number of points to resample to
+    nb_points = np.round(c_length / labels_mm).astype(int)
+    if nb_points < 2:
+        logging.warning(f"{bundle_name} is shorter than the section length.")
+        nb_points = 2
+
+    # Resample the centroid to have `nb_points` points
+    # Adding 2 points so they can be excluded from the labels. Sort of a
+    # reverse signpost problem. Otherwize, the first and last labels
+    # and no other would be assigned to the first and last point of the
+    # centroid.
+    resampled_centroid = set_number_of_points(centroid, nb_points + 2)
+
+    # Re-discretizing the labels based on the resampled centroid
+    discrete_labels = np.zeros_like(bundle_label, dtype=np.float32)
+    for i, label in enumerate(unique[1:]):
+        # Find the closest label in the resampled centroid
+        c = centroid[i]
+        # Calculate the distances from the centroid to the resampled centroid
+        # Exclude the first and last points of the resampled centroid (see
+        # above)
+        distances = np.linalg.norm(
+            c - resampled_centroid[None, 1:-1], axis=-1)
+        # Get the index of the closest label
+        closest_index = np.argmin(distances)
+        # Assign the label to the closest index in the resampled centroid
+        discrete_labels[ref_labels == label] = closest_index + 1
+
+    # Determine the output type based on the number of labels
+    out_type = np.uint16 if nb_points > 1 else np.uint8
+
+    return discrete_labels.astype(out_type)
+
+
+def post_process_labels_continuous(
+    bundle_label, bundle_mask, bundle_name
+):
+    """ Don't discretize the labels, just apply a mask to the bundle.
+
+    Parameters
+    ----------
+    bundle_label : np.ndarray
+        Predicted continuous labels for the bundle.
+    bundle_mask : np.ndarray
+        Binary mask of the bundle.
+    bundle_name : str
+        Name of the bundle, used for logging.
+
+    Returns
+    -------
+    bundle_label : np.ndarray
+        Predicted labels for the bundle.
+    """
+
+    # Determine the output type based on the number of labels
+    # In this case, we assume the labels are continuous and
+    # can be represented as floats.
+    out_type = float
+
+    # Label masking
+    bundle_label[~bundle_mask.astype(bool)] = 0
+
+    return bundle_label.astype(out_type)