[WIP/ENH] add babyseg and update babyAFQ

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,23 +925,20 @@ 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},
             },
             "Left Vertical Occipital": {
                 "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},
             },
             "Right Vertical Occipital": {
                 "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},
             },

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",

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.

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

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,

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

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,