From 92b6097cfb71973498a6e583e9e180c1acb33777 Mon Sep 17 00:00:00 2001 From: Nick Tustison Date: Wed, 13 May 2026 11:53:59 -0700 Subject: [PATCH 1/2] ENH: Port NBM and CIT186 functionality to ANTsPyNet. --- antspynet/utilities/__init__.py | 2 + antspynet/utilities/cit168_labeling.py | 206 ++++++++++++++++++ antspynet/utilities/get_antsxnet_data.py | 8 +- antspynet/utilities/get_pretrained_network.py | 8 +- antspynet/utilities/nbm_labeling.py | 183 ++++++++++++++++ 5 files changed, 405 insertions(+), 2 deletions(-) create mode 100644 antspynet/utilities/cit168_labeling.py create mode 100644 antspynet/utilities/nbm_labeling.py diff --git a/antspynet/utilities/__init__.py b/antspynet/utilities/__init__.py index 13e224e..9f35863 100644 --- a/antspynet/utilities/__init__.py +++ b/antspynet/utilities/__init__.py @@ -89,6 +89,8 @@ from .desikan_killiany_tourville_labeling import desikan_killiany_tourville_labeling from .harvard_oxford_atlas_labeling import harvard_oxford_atlas_labeling from .cerebellum_morphology import cerebellum_morphology +from .nbm_labeling import nbm_labeling +from .cit168_labeling import cit168_labeling from .brain_age import brain_age from .mri_super_resolution import mri_super_resolution from .quality_assessment import tid_neural_image_assessment diff --git a/antspynet/utilities/cit168_labeling.py b/antspynet/utilities/cit168_labeling.py new file mode 100644 index 0000000..6ba6117 --- /dev/null +++ b/antspynet/utilities/cit168_labeling.py @@ -0,0 +1,206 @@ +import numpy as np +import ants + +import keras +from keras import layers, ops, Model + +def cit168_labeling(t1, verbose=False): + + """ + + Perform CIT168 segmentation in T1 images using Pauli atlas (CIT168) + labels described in https://pubmed.ncbi.nlm.nih.gov/29664465/ + + group_labels = [0,7,8,9,23,24,25,33,34] + + group_labels = [0,1,2,5,6,17,18,21,22] + + The labeling is as follows: + + Label 1: BN_STR_Pu_Left + Labeļ 2: BN_STR_Ca_Left + Label 3: BN_STR_NAC_Left. (not modeled) + Label 4: EXA_Left (not modeled) + Label 5: BN_GP_GPe_Left + Label 6: BN_GP_GPi_Left + Label 7: MTg_SN_SNc_Left + Label 8: MTg_RN_Left + Label 9: MTg_SN_SNc_Left + Label 10: MTg_VTR_PBP_Left (not modeled) + Label 11: MTg_VTR_VTA_Left (not modeled) + Label 12: BN_GP_VeP_Left (not modeled) + Label 13: THM_ETH_HN_Left (not modeled) + Label 14: Die_HTH_Left (not modeled) + Label 15: Die_HTH_MN_Left (not modeled) + Label 16: Die_STH_Left (not modeled) + Label 17: BN_STR_Pu_Right + Label 18: BN_STR_Ca_Right + Label 19: BN_STR_NAC_Right (not modeled) + Label 20: EXA_Right (not modeled) + Label 21: BN_GP_GPe_Right + Label 22: BN_GP_GPi_Right + Label 23: MTg_SN_SNc_Right + Label 24: MTg_RN_Right + Label 25: MTg_SN_SNr_Right + Label 26: MTg_VTR_PBP_Right (not modeled) + Label 27: MTg_VTR_VTA_Right (not modeled) + Label 28: BN_GP_VeP_Right (not modeled) + Label 29: THM_ETH_HN_Right (not modeled) + Label 30: Die_HTH_Right (not modeled) + Label 31: Die_HTH_MN_Right (not modeled) + Label 32: Die_STH_Right (not modeled) + Label 33: ReferenceRegion_Left + Label 34: ReferenceRegion_Right + + Preprocessing consists of: + * n4 bias correction and + * brain extraction + + Arguments + --------- + t1 : ANTsImage + input image + + verbose : boolean + Print progress to the screen. + + Returns + ------- + ANTs segmentation label and probability images. + + Example + ------- + >>> seg = cit168_labeling(t1) + """ + + from ..utilities import get_pretrained_network + from ..utilities import preprocess_brain_image + from ..utilities import get_antsxnet_data + from ..architectures import create_unet_model_3d + + if t1.dimension != 3: + raise ValueError( "Image dimension must be 3." ) + + template = ants.image_read(get_antsxnet_data("CIT168_T1w_700um_pad_adni")) + template_small = ants.resample_image(template, [2, 2, 2]) + template_large = ants.resample_image(template, [0.5, 0.5, 0.5]) + template_seg = ants.image_read(get_antsxnet_data("det_atlas_25_pad_LR_adni")) + + template_transform_type = "antsRegistrationSyNQuickRepro[s]" + + cropped_template_size = [160, 160, 112] + + ################################ + # + # Preprocess images + # + ################################ + + # clone input to float + t1 = ants.image_clone(t1, pixeltype="float") + t1 = ants.iMath(t1, "Normalize") + + t1_preprocessed = ants.n4_bias_field_correction(t1, verbose=verbose) + + reg = ants.registration(template_small, t1_preprocessed, type_of_transform=template_transform_type, verbose=verbose) + t1_preprocessed = ants.apply_transforms(fixed=template_large, moving=t1, + transformlist=reg['fwdtransforms'][1], + interpolator="linear", singleprecision=True, verbose=verbose) + template_priors = ants.apply_transforms(fixed=t1_preprocessed, moving=template_seg, + transformlist=reg['invtransforms'][1], + interpolator="nearestNeighbor", singleprecision=True, verbose=verbose) + + center_of_mass = list(ants.get_center_of_mass(ants.threshold_image(template_priors, 0, 0, 0, 1))) + center_of_mass[1] = center_of_mass[1] + 10.0 # same as we did in training + t1_cropped = ants.crop_image_from_center_point(t1_preprocessed, center_of_mass, cropped_template_size) + template_priors_cropped = ants.crop_image_from_center_point(template_priors, center_of_mass, cropped_template_size) + + ################################ + # + # Build model and load weights + # + ################################ + + + cit168_segmentation_image = t1 * 0 + for sn in [True,False]: + + if verbose: + print(" CIT168: Creating model and loading weights. (sn=" + str(sn) + ")") + + if sn: + group_labels = [0,7,8,9,23,24,25,33,34] + cit168_weights = get_pretrained_network("deepCIT168_sn") + else: + group_labels = [0,1,2,5,6,17,18,21,22] + cit168_weights = get_pretrained_network("deepCIT168") + + channel_size = 1 + number_of_outputs = len(group_labels) + number_of_channels = len(group_labels) + + unet0 = create_unet_model_3d([None, None, None, number_of_channels], + number_of_outputs=1, number_of_layers=4, + number_of_filters_at_base_layer=32, + convolution_kernel_size=3, + deconvolution_kernel_size=2, + pool_size=2, strides=2, dropout_rate=0.0, + weight_decay=0, + additional_options="nnUnetActivationStyle", + mode="sigmoid") + + unet1 = create_unet_model_3d([None, None, None, 2], + number_of_outputs=number_of_outputs, + number_of_filters=(32, 64, 96, 128, 256), + convolution_kernel_size=(3, 3, 3), + deconvolution_kernel_size=(2, 2, 2), + dropout_rate=0.0, weight_decay=0, + additional_options = "nnUnetActivationStyle", + mode="classification") + + splits = ops.split(unet0.inputs[0], 9, axis=-1) + newmult = layers.Concatenate(axis=-1)([splits[0], unet0.outputs[0]]) + unetonnet = unet1(newmult) + unet_model = Model(inputs=unet0.inputs, outputs=[unetonnet, unet0.outputs[0]]) + + unet_model.load_weights(cit168_weights) + + priors = ants.segmentation_to_one_hot(template_priors_cropped.numpy().astype('int'), + segmentation_labels=group_labels[1:len(group_labels)]) + batchX = np.concatenate((np.expand_dims(t1_cropped.numpy(), axis=-1), priors), axis=-1) + + if verbose: + print(" CIT168: Model prediction.") + + predicted_data = unet_model.predict(np.expand_dims(batchX, axis=0), verbose=verbose) + + probability_images = [] + for i in range(1, len(group_labels)): + probability_array = np.squeeze(predicted_data[0][0,:,:,:,i]) + probability_image = ants.from_numpy_like(probability_array, t1_cropped) + probability_image = ants.apply_transforms(fixed=t1, moving=probability_image, + transformlist=reg['invtransforms'][0], whichtoinvert=[True], + interpolator="linear", singleprecision=True, verbose=verbose) + probability_images.append(probability_image) + + predicted_mask = ants.from_numpy_like(np.squeeze(predicted_data[1][0,:,:,:,0]), t1_cropped) + predicted_mask = ants.apply_transforms(fixed=t1, moving=predicted_mask, + transformlist=reg['invtransforms'][0], whichtoinvert=[True], + interpolator="linear", singleprecision=True, verbose=verbose) + predicted_mask = ants.threshold_image(predicted_mask, 0.5, 1, 1, 0) + + image_matrix = ants.image_list_to_matrix(probability_images, predicted_mask) + segmentation_matrix = np.argmax(image_matrix, axis=0) + 1 + segmentation_image = ants.matrix_to_images( + np.expand_dims(segmentation_matrix, axis=0), predicted_mask)[0] + relabeled_segmentation_image = segmentation_image * 0 + for i in range(1, len(group_labels)): + single_label_image = ants.threshold_image(segmentation_image, i, i, 1, 0) + if group_labels[i] < 33: + single_label_image = ants.iMath_get_largest_component(single_label_image, 1) + relabeled_segmentation_image += single_label_image * group_labels[i] + + cit168_segmentation_image += relabeled_segmentation_image + + return cit168_segmentation_image + diff --git a/antspynet/utilities/get_antsxnet_data.py b/antspynet/utilities/get_antsxnet_data.py index bfa44b5..d5496be 100644 --- a/antspynet/utilities/get_antsxnet_data.py +++ b/antspynet/utilities/get_antsxnet_data.py @@ -117,7 +117,10 @@ def switch_data(argument): "DevCCF_P56_MRI-T2_50um_BrainParcellationNickMask": "https://ndownloader.figshare.com/files/44706238", "DevCCF_P56_MRI-T2_50um_BrainParcellationTctMask": "https://ndownloader.figshare.com/files/47214532", "DevCCF_P04_STPT_50um": "https://ndownloader.figshare.com/files/46711546", - "DevCCF_P04_STPT_50um_BrainParcellationJayMask": "https://ndownloader.figshare.com/files/46712656" + "DevCCF_P04_STPT_50um_BrainParcellationJayMask": "https://ndownloader.figshare.com/files/46712656", + "CIT168_T1w_700um_pad_adni": "https://ndownloader.figshare.com/files/64536261", + "CIT168_basal_forebrain_adni": "https://ndownloader.figshare.com/files/64536324", + "det_atlas_25_pad_LR_adni": "https://ndownloader.figshare.com/files/64536333", } return(switcher.get(argument, "Invalid argument.")) @@ -179,6 +182,9 @@ def switch_data(argument): "hcpinterFATemplate", "hcpinterTemplateBrainMask", "hcpinterTemplateBrainSegmentation", + "CIT168_T1w_700um_pad_adni", + "CIT168_basal_forebrain_adni", + "det_atlas_25_pad_LR_adni", "show") if not file_id in valid_list: diff --git a/antspynet/utilities/get_pretrained_network.py b/antspynet/utilities/get_pretrained_network.py index 95b3458..9f4654f 100644 --- a/antspynet/utilities/get_pretrained_network.py +++ b/antspynet/utilities/get_pretrained_network.py @@ -187,7 +187,10 @@ def switch_networks(argument): "sig_smallshort_train_2x2x2_1chan_featgraderL6_best_mdl": "https://ndownloader.figshare.com/files/49339858", "sig_smallshort_train_2x2x2_1chan_featvggL6_best_mdl": "https://ndownloader.figshare.com/files/49339861", "sig_smallshort_train_2x2x4_1chan_featgraderL6_best_mdl": "https://ndownloader.figshare.com/files/49339867", - "sig_smallshort_train_2x2x4_1chan_featvggL6_best_mdl": "https://ndownloader.figshare.com/files/49339864" + "sig_smallshort_train_2x2x4_1chan_featvggL6_best_mdl": "https://ndownloader.figshare.com/files/49339864", + "deep_nbm_rank": "https://ndownloader.figshare.com/files/64536405", + "deepCIT168": "https://ndownloader.figshare.com/files/64536618", + "deepCIT168_sn": "https://ndownloader.figshare.com/files/64537371" } return(switcher.get(argument, "Invalid argument.")) @@ -350,6 +353,9 @@ def switch_networks(argument): "sig_smallshort_train_2x2x2_1chan_featvggL6_best_mdl", "sig_smallshort_train_2x2x4_1chan_featgraderL6_best_mdl", "sig_smallshort_train_2x2x4_1chan_featvggL6_best_mdl", + "deep_nbm_rank", + "deepCIT168", + "deepCIT168_sn", "show") if not file_id in valid_list: diff --git a/antspynet/utilities/nbm_labeling.py b/antspynet/utilities/nbm_labeling.py new file mode 100644 index 0000000..6cc4fe3 --- /dev/null +++ b/antspynet/utilities/nbm_labeling.py @@ -0,0 +1,183 @@ +import numpy as np +import ants + +import keras +from keras import layers, Model + +def nbm_labeling(t1, verbose=False): + + """ + Perform CH13 and Nucleaus basalis of Meynert (NBM) segmentation in + T1 images using Avants labels. + + The labeling is as follows: + + Label 1: CH13_left + Label 2: CH13_right + Label 3: NBM_left_ant + Label 4: NBM_left_mid + Label 5: NBM_left_pos + Label 6: NBM_right_ant + Label 7: NBM_right_mid + Label 8: NBM_right_pos + + Preprocessing consists of: + * n4 bias correction and + * brain extraction + + Arguments + --------- + t1 : ANTsImage + input image + + verbose : boolean + Print progress to the screen. + + Returns + ------- + ANTs segmentation label and probability images. + + Example + ------- + >>> seg = nbm_labeling(t1) + """ + + from ..utilities import get_pretrained_network + from ..utilities import preprocess_brain_image + from ..utilities import get_antsxnet_data + from ..architectures import create_unet_model_3d + + if t1.dimension != 3: + raise ValueError( "Image dimension must be 3." ) + + which_template = "CIT168_T1w_700um_pad_adni" + template_transform_type = "antsRegistrationSyNReproQuick[a]" + template = ants.image_read(get_antsxnet_data(which_template)) + template = ants.rank_intensity(template) + + template_seg = ants.image_read(get_antsxnet_data("CIT168_basal_forebrain_adni")) + center_of_mass_labels = (1, 2, 3, 4) + + cropped_template_size = (144, 96, 64) + + ################################ + # + # Preprocess images + # + ################################ + + # clone input to float + t1 = ants.image_clone(t1, pixeltype="float") + t1 = ants.iMath(t1, "Normalize") + + t1_preprocessing = preprocess_brain_image(t1, + truncate_intensity=[1e-4, 0.999], + brain_extraction_modality="t1threetissue", + template=which_template, + template_transform_type=template_transform_type, + do_bias_correction=True, + do_denoising=False, + verbose=verbose) + t1_preprocessed = t1_preprocessing["preprocessed_image"] * t1_preprocessing['brain_mask'] + t1_preprocessed = ants.rank_intensity(t1_preprocessed) + + template_seg_masked = ants.mask_image(template_seg, template_seg, center_of_mass_labels, binarize=True) + center_of_mass = ants.get_center_of_mass(template_seg_masked) + + t1_cropped = ants.crop_image_from_center_point(t1_preprocessed, center_of_mass, cropped_template_size) + + + ################################ + # + # Build model and load weights + # + ################################ + + nbm_lateral_labels = (0,) + nbm_lateral_left_labels = (1, 3, 4, 5) + nbm_lateral_right_labels = (2, 6, 7, 8) + + labels = sorted((*nbm_lateral_labels, *nbm_lateral_left_labels, *nbm_lateral_right_labels)) + number_of_outputs = len(labels) + channel_size = 1 + + if verbose: + print(" NBM: Creating model and loading weights.") + + unet0 = create_unet_model_3d([None, None, None, channel_size], + number_of_outputs=1, number_of_layers=4, + number_of_filters_at_base_layer=32, + convolution_kernel_size=3, + deconvolution_kernel_size=2, + pool_size=2, strides=2, dropout_rate=0.0, + weight_decay=0, + additional_options="nnUnetActivationStyle", + mode="sigmoid") + + unet1 = create_unet_model_3d([None, None, None, 2], + number_of_outputs=number_of_outputs, + number_of_filters=(32, 64, 96, 128, 256), + convolution_kernel_size=(3, 3, 3), + deconvolution_kernel_size=(2, 2, 2), + dropout_rate=0.0, weight_decay=0, + additional_options = "nnUnetActivationStyle", + mode="classification") + + nextin = layers.Concatenate(axis=-1)([unet0.inputs[0], unet0.outputs[0]]) + + unet_model = Model(inputs=unet0.inputs, outputs=[unet1(nextin), unet0.outputs[0]]) + + nbm_weights = get_pretrained_network("deep_nbm_rank") + unet_model.load_weights(nbm_weights) + + ################################ + # + # Do prediction and normalize to native space + # + ################################ + + if verbose: + print("Model prediction using both the original and contralaterally flipped version.") + + batchX = np.zeros((2, *cropped_template_size, channel_size)) + batchX[0,:,:,:,0] = t1_cropped.numpy() + batchX[1,:,:,:,0] = np.flip(batchX[0,:,:,:,0], axis=0) + + predicted_data = unet_model.predict(batchX, verbose=verbose) + + nbm_labels = sorted((*nbm_lateral_labels, *nbm_lateral_left_labels, *nbm_lateral_right_labels)) + probability_images = [None] * len(labels) + + for i in range(len(nbm_labels)): + label = nbm_labels[i] + if label == 0: + probability_array = np.squeeze(predicted_data[0][0,:,:,:,0]) + probability_array_flipped = np.flip(np.squeeze(predicted_data[0][1,:,:,:,0]), axis=0) + else: + label_index = nbm_labels.index(label) + probability_array = np.squeeze(predicted_data[0][0,:,:,:,label_index]) + if label in nbm_lateral_left_labels: + left_index = nbm_lateral_left_labels.index(label) + right_label = nbm_lateral_right_labels[left_index] + label_flipped_index = nbm_labels.index(right_label) + else: + right_index = nbm_lateral_right_labels.index(label) + left_label = nbm_lateral_left_labels[right_index] + label_flipped_index = nbm_labels.index(left_label) + probability_array_flipped = np.flip(np.squeeze(predicted_data[0][1,:,:,:,label_flipped_index]), axis=0) + + probability_image = ants.from_numpy_like(0.5 * (probability_array + probability_array_flipped), t1_cropped) + probability_image = ants.apply_transforms(fixed=t1, + moving=probability_image, + transformlist=t1_preprocessing['template_transforms']['invtransforms'], + whichtoinvert=[True], interpolator="linear", singleprecision=True, verbose=verbose) + probability_images[i] = probability_image + + image_matrix = ants.image_list_to_matrix(probability_images, t1 * 0 + 1) + segmentation_matrix = np.argmax(image_matrix, axis=0) + segmentation_image = ants.matrix_to_images( + np.expand_dims(segmentation_matrix, axis=0), t1 * 0 + 1)[0] + + return_dict = {'segmentation_image' : segmentation_image, + 'probability_images' : probability_images} + return(return_dict) From f34fd6222c7ce074cc83186a4f6ed5f7740d2357 Mon Sep 17 00:00:00 2001 From: Nick Tustison Date: Wed, 13 May 2026 13:53:03 -0700 Subject: [PATCH 2/2] BUG: braine xtraction. --- antspynet/utilities/cit168_labeling.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/antspynet/utilities/cit168_labeling.py b/antspynet/utilities/cit168_labeling.py index 6ba6117..1fe28fe 100644 --- a/antspynet/utilities/cit168_labeling.py +++ b/antspynet/utilities/cit168_labeling.py @@ -98,14 +98,23 @@ def cit168_labeling(t1, verbose=False): # clone input to float t1 = ants.image_clone(t1, pixeltype="float") - t1 = ants.iMath(t1, "Normalize") - t1_preprocessed = ants.n4_bias_field_correction(t1, verbose=verbose) + t1_preprocessing = preprocess_brain_image(t1, + truncate_intensity=[1e-4, 0.999], + brain_extraction_modality="t1threetissue", + template=None, + do_bias_correction=True, + do_denoising=False, + intensity_normalization_type='01', + verbose=verbose) + + t1_preprocessed = t1_preprocessing['preprocessed_image'] * t1_preprocessing['brain_mask'] reg = ants.registration(template_small, t1_preprocessed, type_of_transform=template_transform_type, verbose=verbose) - t1_preprocessed = ants.apply_transforms(fixed=template_large, moving=t1, + t1_preprocessed = ants.apply_transforms(fixed=template_large, moving=t1_preprocessed, transformlist=reg['fwdtransforms'][1], interpolator="linear", singleprecision=True, verbose=verbose) + template_priors = ants.apply_transforms(fixed=t1_preprocessed, moving=template_seg, transformlist=reg['invtransforms'][1], interpolator="nearestNeighbor", singleprecision=True, verbose=verbose)