From 6b4ecbdeeae433b2c8164eda39543effba5d91fc Mon Sep 17 00:00:00 2001 From: 36000 Date: Mon, 27 Jul 2026 22:47:20 -0700 Subject: [PATCH 1/2] [WIP/ENH] add babyseg and update babyAFQ --- AFQ/api/bundle_dict.py | 4 - AFQ/data/fetch.py | 19 +++++ AFQ/definitions/image.py | 23 ++++++ AFQ/nn/babyseg.py | 164 +++++++++++++++++++++++++++++++++++++ AFQ/nn/utils.py | 14 +++- AFQ/tasks/structural.py | 28 +++++++ AFQ/tasks/tissue.py | 32 +++++++- docs/source/references.bib | 20 +++++ 8 files changed, 296 insertions(+), 8 deletions(-) create mode 100644 AFQ/nn/babyseg.py diff --git a/AFQ/api/bundle_dict.py b/AFQ/api/bundle_dict.py index ca1cc7a2..9cf01a2c 100644 --- a/AFQ/api/bundle_dict.py +++ b/AFQ/api/bundle_dict.py @@ -916,7 +916,6 @@ def baby_bd(): "start": templates["pARC_L_start"], "end": templates["VOF_box_small_L"], "primary_axis": "I/S", - "primary_axis_percentage": 40, "cross_midline": False, "mahal": {"distance_threshold": 4}, }, @@ -926,7 +925,6 @@ def baby_bd(): "start": templates["pARC_R_start"], "end": templates["VOF_box_small_R"], "primary_axis": "I/S", - "primary_axis_percentage": 40, "cross_midline": False, "mahal": {"distance_threshold": 4}, }, @@ -934,7 +932,6 @@ def baby_bd(): "start": templates["VOF_L_start"], "end": templates["VOF_box_small_L"], "primary_axis": "I/S", - "primary_axis_percentage": 40, "cross_midline": False, "mahal": {"distance_threshold": 4}, }, @@ -942,7 +939,6 @@ def baby_bd(): "start": templates["VOF_R_start"], "end": templates["VOF_box_small_R"], "primary_axis": "I/S", - "primary_axis_percentage": 40, "cross_midline": False, "mahal": {"distance_threshold": 4}, }, diff --git a/AFQ/data/fetch.py b/AFQ/data/fetch.py index a7f4f30d..507de777 100644 --- a/AFQ/data/fetch.py +++ b/AFQ/data/fetch.py @@ -55,6 +55,7 @@ "fetch_brainchop_models", "fetch_multiaxial_models", "fetch_synthseg_models", + "fetch_babyseg_models", "fetch_templates", "read_templates", "fetch_stanford_hardi_tractography", @@ -263,6 +264,24 @@ def read_callosum_templates(as_img=True, resample_to=False): doc="Download ONNX SynthSeg models", ) +babyseg_remote_fnames = ["67035599"] + +babyseg_fnames = ["babyseg.onnx"] + +babyseg_md5_hashes = [ + "5f64b439dcfff36045448e0cc562f6e7", +] + +fetch_babyseg_models = _make_reusable_fetcher( + "fetch_babyseg_models", + op.join(afq_home, "babyseg_onnx"), + baseurl, + babyseg_remote_fnames, + babyseg_fnames, + md5_list=babyseg_md5_hashes, + doc="Download ONNX BabySeg models", +) + multiaxial_fnames = [ "sagittal_model.onnx", "axial_model.onnx", diff --git a/AFQ/definitions/image.py b/AFQ/definitions/image.py index 4d68fb8b..f5f28c43 100644 --- a/AFQ/definitions/image.py +++ b/AFQ/definitions/image.py @@ -12,6 +12,7 @@ __all__ = [ "ImageFile", "FullImage", + "BabyBrainMask", "RoiImage", "LabelledImageFile", "ThresholdedImageFile", @@ -312,6 +313,28 @@ def image_getter(data_imap): return image_getter +class BabyBrainMask(ImageDefinition): + def __init__(self): + pass + + def get_name(self): + return "babyseg_brain_mask" + + def get_image_getter(self, task_name): + def _image_getter_helper(babyseg_model): + predictions = nib.load(babyseg_model) + brain_mask = (predictions.get_fdata() > 0).astype(np.uint8) + brain_mask_img = nib.Nifti1Image(brain_mask, predictions.affine) + return brain_mask_img, dict(BabySegPredictions=babyseg_model) + + if task_name == "structural": + return _image_getter_helper + else: + raise ValueError( + "BabyBrainMask can only be used in the structural context." + ) + + class RoiImage(ImageDefinition): """ Define an image which is all include ROIs or'd together. diff --git a/AFQ/nn/babyseg.py b/AFQ/nn/babyseg.py new file mode 100644 index 00000000..59e7ef6c --- /dev/null +++ b/AFQ/nn/babyseg.py @@ -0,0 +1,164 @@ +import logging +import os.path as op +from enum import IntEnum +from time import time + +import numpy as np +from scipy.ndimage import gaussian_filter +from skimage.segmentation import find_boundaries + +from AFQ.data.fetch import afq_home, fetch_babyseg_models +from AFQ.nn.utils import prepare_t1_for_nn, resample_output + +logger = logging.getLogger("AFQ") + + +class BabySegLabels(IntEnum): + BACKGROUND = 0 + LEFT_CEREBRAL_WHITE_MATTER = 1 + LEFT_CEREBRAL_CORTEX = 2 + LEFT_LATERAL_VENTRICLE = 3 + LEFT_CEREBELLUM_CORTEX = 4 + LEFT_THALAMUS = 5 + LEFT_CAUDATE = 6 + BRAIN_STEM = 7 + LEFT_HIPPOCAMPUS = 8 + LEFT_AMYGDALA = 9 + LEFT_VENTRAL_DC = 10 + RIGHT_CEREBRAL_WHITE_MATTER = 11 + RIGHT_CEREBRAL_CORTEX = 12 + RIGHT_LATERAL_VENTRICLE = 13 + RIGHT_CEREBELLUM_CORTEX = 14 + RIGHT_THALAMUS = 15 + RIGHT_CAUDATE = 16 + RIGHT_HIPPOCAMPUS = 17 + RIGHT_AMYGDALA = 18 + RIGHT_VENTRAL_DC = 19 + LEFT_BASAL_GANGLIA = 20 + RIGHT_BASAL_GANGLIA = 21 + + +def _get_model(model_name): + model_dir = op.join(afq_home, "babyseg_onnx") + model_dictionary = { + "babyseg": "babyseg.onnx", + } + + model_fname = op.join(model_dir, model_dictionary[model_name]) + if not op.exists(model_fname): + fetch_babyseg_models() + + return model_fname + + +def run_babyseg( + ort, + t1_img, + onnx_kwargs, +): + """ + Run the BabySeg Model + + References + ---------- + .. [1] Hoffmann M, Zöllei L, Dalca AV. Deep infant brain segmentation from + multi-contrast MRI. Asilomar Conference on Signals, Systems, and + Computers, 2025, pp. 974-981. https://arxiv.org/abs/2512.05114 + .. [2] Hoffmann M. Domain-randomized deep learning for neuroimage analysis. + IEEE Signal Processing Magazine, 42(4):78-90, 2025. + https://arxiv.org/abs/2507.13458 + """ + model = _get_model("babyseg") + t1_data, conformed_affine = prepare_t1_for_nn( + t1_img, orientation="LIA", out_shape_dynamic=True + ) + + image = t1_data.astype(np.float32)[None, None, ...] + + logger.info("Running Babyseg...") + start_time = time() + sess = ort.InferenceSession(model, **onnx_kwargs) + input_name = sess.get_inputs()[0].name + output_name = sess.get_outputs()[0].name + output_channels = sess.run([output_name], {input_name: image})[0] + total_time = time() - start_time + logger.info((f"Finished Babyseg in {total_time:.2f} seconds.")) + + output = output_channels.argmax(axis=1)[0].astype(np.uint8) + + output_img = resample_output(output, conformed_affine, t1_img) + + return output_img + + +def pve_from_babyseg(babyseg_data): + """ + Compute partial volume estimates from BabySeg segmentation. + + Parameters + ---------- + babyseg_data : ndarray + The output segmentation from BabySeg. + + Returns + ------- + pve : ndarray + PVE data with CSF, GM, and WM segmentations. + """ + CSF_labels = [ + BabySegLabels.BACKGROUND, + BabySegLabels.LEFT_LATERAL_VENTRICLE, + BabySegLabels.RIGHT_LATERAL_VENTRICLE, + ] + + GM_labels = [ + BabySegLabels.LEFT_CEREBRAL_CORTEX, + BabySegLabels.LEFT_CEREBELLUM_CORTEX, + BabySegLabels.LEFT_THALAMUS, + BabySegLabels.LEFT_CAUDATE, + BabySegLabels.LEFT_HIPPOCAMPUS, + BabySegLabels.LEFT_AMYGDALA, + BabySegLabels.RIGHT_CEREBRAL_CORTEX, + BabySegLabels.RIGHT_CEREBELLUM_CORTEX, + BabySegLabels.RIGHT_THALAMUS, + BabySegLabels.RIGHT_CAUDATE, + BabySegLabels.RIGHT_HIPPOCAMPUS, + BabySegLabels.RIGHT_AMYGDALA, + BabySegLabels.LEFT_BASAL_GANGLIA, + BabySegLabels.RIGHT_BASAL_GANGLIA, + ] + + WM_labels = [ + BabySegLabels.LEFT_CEREBRAL_WHITE_MATTER, + BabySegLabels.RIGHT_CEREBRAL_WHITE_MATTER, + ] + + mixed_labels = [ + BabySegLabels.BRAIN_STEM, + BabySegLabels.LEFT_VENTRAL_DC, + BabySegLabels.RIGHT_VENTRAL_DC, + ] + + PVE = np.zeros(babyseg_data.shape + (3,), dtype=np.float32) + + PVE[np.isin(babyseg_data, CSF_labels), 0] = 1.0 + PVE[np.isin(babyseg_data, GM_labels), 1] = 1.0 + PVE[np.isin(babyseg_data, WM_labels), 2] = 1.0 + + # For mixed labels, we assume they are WM interior, GM exterior + # except on boundaries with wm, where we assume they are WM. + # We additionally set GM to 0.4 and WM to 0.6 + # This is a simplification, basically so they do not cause problems + # with ACT + wm_fuzzed = gaussian_filter(PVE[..., 2], 1) + nwm_fuzzed = gaussian_filter(PVE[..., 0] + PVE[..., 1], 1) + bs_exterior = np.logical_and( + find_boundaries(np.isin(babyseg_data, mixed_labels), mode="inner"), + nwm_fuzzed >= wm_fuzzed, + ) + PVE[np.isin(babyseg_data, mixed_labels), 1] = 0.4 + PVE[np.isin(babyseg_data, mixed_labels), 2] = 0.6 + PVE[bs_exterior, 1] = 1.0 + PVE[bs_exterior, 2] = 0.0 + + return PVE diff --git a/AFQ/nn/utils.py b/AFQ/nn/utils.py index 07f4dd65..96893122 100644 --- a/AFQ/nn/utils.py +++ b/AFQ/nn/utils.py @@ -23,12 +23,22 @@ def crop_to_nonzero(img): return nib.Nifti1Image(cropped_data, new_affine) -def prepare_t1_for_nn(t1_img, orientation="RAS"): +def prepare_t1_for_nn(t1_img, orientation="RAS", out_shape_dynamic=False): t1_img_cropped = crop_to_nonzero(t1_img) + if out_shape_dynamic: + divisor = 64 + min_shape = 128 + max_shape = 320 + s = np.array(t1_img_cropped.shape[:3]) + out_shape = tuple( + np.clip(np.ceil(s / divisor).astype(int) * divisor, min_shape, max_shape) + ) + else: + out_shape = (256, 256, 256) t1_img_conformed = nbp.conform( t1_img_cropped, - out_shape=(256, 256, 256), + out_shape=out_shape, voxel_size=(1.0, 1.0, 1.0), orientation=orientation, order=1, diff --git a/AFQ/tasks/structural.py b/AFQ/tasks/structural.py index aa85286d..d01380fd 100644 --- a/AFQ/tasks/structural.py +++ b/AFQ/tasks/structural.py @@ -6,6 +6,7 @@ from numba import get_num_threads from AFQ.definitions.utils import Definition +from AFQ.nn.babyseg import run_babyseg from AFQ.nn.brainchop import run_brainchop from AFQ.nn.multiaxial import run_multiaxial from AFQ.nn.synthseg import run_synthseg @@ -89,6 +90,32 @@ def onnx_kwargs( return {"onnx_kwargs": onnx_kwargs} +@immlib.calc("babyseg_model") +@as_file(suffix="_model-babyseg_probseg.nii.gz", subfolder="nn") +def babyseg_model(t1_file, citations, onnx_kwargs): + """ + full path to the babyseg model segmentations + + References + ---------- + [1] Hoffmann M, Zöllei L, Dalca AV. "Deep infant brain segmentation from + multi-contrast MRI." Asilomar Conference on Signals, Systems, and + Computers, 2025, pp. 974-981. https://arxiv.org/abs/2512.05114 + [2] Hoffmann M. "Domain-randomized deep learning for neuroimage analysis." + IEEE Signal Processing Magazine, 42(4):78-90, 2025. + https://arxiv.org/abs/2507.13458 + """ + citations.add("hoffmann2025deep") + citations.add("hoffmann2025domain") + ort = check_onnxruntime( + "BabySeg", + "Or, provide your own segmentations using PVEImage or PVEImages.", + ) + t1_img = nib.load(t1_file) + predictions = run_babyseg(ort, t1_img, onnx_kwargs) + return predictions, dict(T1w=t1_file) + + @immlib.calc("synthseg_model") @as_file(suffix="_model-synthseg2_probseg.nii.gz", subfolder="nn") def synthseg_model(t1_file, citations, onnx_kwargs): @@ -250,6 +277,7 @@ def get_structural_plan(kwargs): [ mx_model, synthseg_model, + babyseg_model, t1w_brain_mask, t1_subcortex, t1_masked, diff --git a/AFQ/tasks/tissue.py b/AFQ/tasks/tissue.py index 012f7bd2..cf6cbca4 100644 --- a/AFQ/tasks/tissue.py +++ b/AFQ/tasks/tissue.py @@ -22,6 +22,7 @@ from AFQ.models.msmt import MultiShellDeconvModel from AFQ.models.QBallTP import anisotropic_power from AFQ.models.wmgm_interface import fit_wm_gm_interface +from AFQ.nn.babyseg import pve_from_babyseg from AFQ.nn.brainchop import pve_from_subcortex from AFQ.nn.multiaxial import extract_pve from AFQ.nn.synthseg import pve_from_synthseg @@ -61,14 +62,33 @@ def pve_internal(structural_imap, pve="synthseg"): or a Definition object to import the PVE. Importing a PVE from software like Freesurfer or FSL FAST is recommended if they are available. - The built-in methods are "synthseg" or "multiaxial+brainchop". + The built-in methods are "synthseg", "multiaxial+brainchop", + "multiaxial+brainchop+synthseg", and "babyseg". "synthseg" uses SynthSeg2 [1] to get the PVE. "multiaxial+brainchop" uses MultiAxial [2] and BrainChop [3] segmentations to get the PVE. Note this requires downloading the pre-trained multi-axial model which is licensed with Creative Commons Attribution-NonCommercial-ShareAlike 4.0 - International. + International. "multiaxial+brainchop+synthseg" combines all + three. Finally, "babyseg" uses BabySeg [4] to get the PVE + for baby brain data. Default: "synthseg" + References + ---------- + [1] Billot, Benjamin, et al. "Robust machine learning segmentation + for large-scale analysis of heterogeneous clinical brain MRI + datasets." Proceedings of the National Academy of Sciences 120.9 + (2023): e2216399120. + [2] Birnbaum, Andrew M., et al. "Full-head segmentation of MRI + with abnormal brain anatomy: model and data release." Journal of + Medical Imaging 12.5 (2025): 054001-054001. + [3] Masoud, M., Hu, F., & Plis, S. (2023). Brainchop: In-browser MRI + volumetric segmentation and rendering. Journal of Open Source + Software, 8(83), 5098. + https://doi.org/10.21105/joss.05098 + [4] Hoffmann M, Zöllei L, Dalca AV. "Deep infant brain segmentation from + multi-contrast MRI." Asilomar Conference on Signals, Systems, and + Computers, 2025, pp. 974-981. https://arxiv.org/abs/2512.05114 """ if isinstance(pve, str): if pve == "synthseg": @@ -120,6 +140,14 @@ def pve_internal(structural_imap, pve="synthseg"): meta["SynthsegParcellation"] = structural_imap["synthseg_model"] return nib.Nifti1Image(PVE, t1_subcortex_img.affine), meta + elif pve == "babyseg": + babyseg_seg = nib.load(structural_imap["babyseg_model"]) + PVE = pve_from_babyseg(babyseg_seg.get_fdata()) + + return nib.Nifti1Image(PVE, babyseg_seg.affine), dict( + BabySegParcellation=structural_imap["babyseg_model"], + labels=["csf", "gm", "wm"], + ) raise ValueError( "pve must be a PVEImage, PVEImages, 'synthseg', or 'multiaxial+brainchop'" diff --git a/docs/source/references.bib b/docs/source/references.bib index 45295ce9..60c61b07 100644 --- a/docs/source/references.bib +++ b/docs/source/references.bib @@ -1,3 +1,23 @@ +@inproceedings{hoffmann2025deep, + title={{Deep infant brain segmentation from multi-contrast MRI}}, + author={Hoffmann, Malte and Z{\"o}llei, Lilla and Dalca, Adrian V}, + booktitle={{Asilomar Conference on Signals, Systems, and Computers}}, + pages={974--981}, + year={2025}, + publisher={IEEE} +} + +@article{hoffmann2025domain, + title={Domain-randomized deep learning for neuroimage analysis}, + author={Hoffmann, Malte}, + journal={IEEE Signal Processing Magazine}, + volume={42}, + number={4}, + pages={78--90}, + year={2025}, + publisher={IEEE} +} + @article{Grotheer2022, title={White matter myelination during early infancy is linked to spatial gradients and myelin content at birth}, author={Grotheer, Mareike and Rosenke, Mona and Wu, Hua and Kular, Holly and Querdasi, Francesca R and Natu, Vaidehi S and Yeatman, Jason D and Grill-Spector, Kalanit}, From 6e2805d4719e08ebadb429dddc333eae5ab973bf Mon Sep 17 00:00:00 2001 From: 36000 Date: Thu, 13 Aug 2026 12:18:41 -0700 Subject: [PATCH 2/2] fix ROI tolerance errors --- AFQ/recognition/criteria.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/AFQ/recognition/criteria.py b/AFQ/recognition/criteria.py index de0b17c5..3e8864c4 100644 --- a/AFQ/recognition/criteria.py +++ b/AFQ/recognition/criteria.py @@ -191,9 +191,9 @@ def include(b_sls, bundle_def, **kwargs): if "inc_addtol" in bundle_def: include_roi_tols = [] for inc_tol in bundle_def["inc_addtol"]: - include_roi_tols.append((inc_tol / kwargs["vox_dim"] + kwargs["tol"]) ** 2) - else: - include_roi_tols = [kwargs["tol"] ** 2] * len(bundle_def["include"]) + include_roi_tols.append(inc_tol / kwargs["vox_dim"] + kwargs["tol"]) + else: # TODO: should this be distance_to_corner / 2? + include_roi_tols = [kwargs["tol"]] * len(bundle_def["include"]) inc_results = abr.check_sls_with_inclusion( b_sls.get_selected_sls(), bundle_def["include"], include_roi_tols @@ -261,9 +261,9 @@ def exclude(b_sls, bundle_def, **kwargs): if "exc_addtol" in bundle_def: exclude_roi_tols = [] for exc_tol in bundle_def["exc_addtol"]: - exclude_roi_tols.append((exc_tol / kwargs["vox_dim"] + kwargs["tol"]) ** 2) + exclude_roi_tols.append(exc_tol / kwargs["vox_dim"] + kwargs["tol"]) else: - exclude_roi_tols = [kwargs["tol"] ** 2] * len(bundle_def["exclude"]) + exclude_roi_tols = [kwargs["tol"]] * len(bundle_def["exclude"]) for sl_idx, sl in enumerate(b_sls.get_selected_sls()): if abr.check_sl_with_exclusion(sl, bundle_def["exclude"], exclude_roi_tols): accept_idx[sl_idx] = 1