pyAFQ 3.2__tar.gz → 3.4__tar.gz

{pyafq-3.2 → pyafq-3.4}/.codespellrc

@@ -1,3 +1,3 @@
 [codespell]
 skip = [setup.cfg] 
-ignore-words-list = Reson, DNE, ACI, FPT, sagital, saggital, abd, Joo, Mapp, Commun, vor, Claus, coo
+ignore-words-list = Reson, DNE, ACI, FPT, sagital, saggital, abd, Joo, Mapp, Commun, vor, Claus, coo, ans

{pyafq-3.2 → pyafq-3.4}/.github/workflows/docbuild.yml

@@ -44,7 +44,7 @@ jobs:
     - name: Build docs
       run: |
         cd docs
-        make html
+        make html SPHINXOPTS="-T"
     - name: Upload docs
       uses: actions/upload-artifact@v4
       with:

{pyafq-3.2 → pyafq-3.4}/.github/workflows/docker_pyafq.yml

@@ -54,3 +54,6 @@ jobs:
     - name: Push pyAFQ GPU docker image
       if: github.ref == 'refs/heads/master' || startsWith(github.event.ref, 'refs/tags')
       run: sh ./gpu_docker/docker-push.sh "nrdg" "11"
+
+    - name: Report pyAFQ image size
+      run: docker images --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}" | tee -a $GITHUB_STEP_SUMMARY

{pyafq-3.2 → pyafq-3.4}/.github/workflows/nightly_anisotropic_test.yml

@@ -27,7 +27,7 @@ jobs:
     - name: Install
       run: |
         python -m pip install --upgrade pip
-        pip install .[dev,fury,afqbrowser,plot,nn]
+        pip install .[dev,fury,afqbrowser,nn]
     - name: Lint
       run: |
         flake8 --ignore N802,N806,W503 --select W504 `find . -name \*.py | grep -v setup.py | grep -v version.py | grep -v __init__.py | grep -v /docs/`

{pyafq-3.2 → pyafq-3.4}/.github/workflows/nightly_basic_test.yml

@@ -27,7 +27,7 @@ jobs:
     - name: Install
       run: |
         python -m pip install --upgrade pip
-        pip install .[dev,fury,afqbrowser,plot,nn]
+        pip install .[dev,fury,afqbrowser,nn]
     - name: Lint
       run: |
         flake8 --ignore N802,N806,W503 --select W504 `find . -name \*.py | grep -v setup.py | grep -v version.py | grep -v __init__.py | grep -v /docs/`

{pyafq-3.2 → pyafq-3.4}/.github/workflows/nightly_reco80_test.yml

@@ -27,7 +27,7 @@ jobs:
     - name: Install
       run: |
         python -m pip install --upgrade pip
-        pip install .[dev,fury,afqbrowser,plot,nn]
+        pip install .[dev,fury,afqbrowser,nn]
     - name: Lint
       run: |
         flake8 --ignore N802,N806,W503 --select W504 `find . -name \*.py | grep -v setup.py | grep -v version.py | grep -v __init__.py | grep -v /docs/`

{pyafq-3.2 → pyafq-3.4}/.github/workflows/nightly_reco_test.yml

@@ -27,7 +27,7 @@ jobs:
     - name: Install
       run: |
         python -m pip install --upgrade pip
-        pip install .[dev,fury,afqbrowser,plot,nn]
+        pip install .[dev,fury,afqbrowser,nn]
     - name: Lint
       run: |
         flake8 --ignore N802,N806,W503 --select W504 `find . -name \*.py | grep -v setup.py | grep -v version.py | grep -v __init__.py | grep -v /docs/`

{pyafq-3.2 → pyafq-3.4}/.github/workflows/nightly_test.yml

@@ -27,7 +27,7 @@ jobs:
     - name: Install
       run: |
         python -m pip install --upgrade pip
-        pip install .[dev,fury,afqbrowser,plot,nn]
+        pip install .[dev,fury,afqbrowser,nn]
     - name: Lint
       run: |
         flake8 --ignore N802,N806,W503 --select W504 `find . -name \*.py | grep -v setup.py | grep -v version.py | grep -v __init__.py | grep -v /docs/`

{pyafq-3.2 → pyafq-3.4}/.gitignore

@@ -5,9 +5,6 @@ AFQ/version.py
 docs/_build
 docs/build
 docs/source/auto_examples/
-docs/source/tutorials/tutorial_examples/
-docs/source/howto/howto_examples/
-docs/source/reference/config.rst
 examples/**/*.nii.gz
 examples/**/*.trk
 examples/**/*.npy
@@ -17,5 +14,6 @@ dist
 .DS_Store
 **.sw?
 **Untitled.ipynb
+docs/source/examples/**/*.ipynb
 pip-wheel-metadata
 pyAFQ.egg-info

{pyafq-3.2 → pyafq-3.4}/AFQ/_fixes.py

@@ -3,6 +3,7 @@ import math
 import tempfile
 from math import radians
 
+import imageio
 import numpy as np
 from dipy.align import vector_fields as vfu
 from dipy.align.imwarp import DiffeomorphicMap, mult_aff
@@ -11,7 +12,7 @@ from dipy.tracking.streamline import set_number_of_points
 from PIL import Image
 from scipy.linalg import blas, pinvh
 from scipy.special import gammaln, lpmv
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 logger = logging.getLogger("AFQ")
 
@@ -271,17 +272,17 @@ def gaussian_weights(
     n_sls, n_nodes, _ = sls.shape
 
     def _weighting_failed():
-        weights = np.ones((n_sls, n_nodes))
         logger.warning(
             (
-                "Not enough streamlines for weight calculation, "
+                f"Not enough or too many streamlines (n_sls: {n_sls}) "
+                "for weight calculation, "
                 "weighting everything evenly"
             )
         )
         if return_mahalanobis:
             return np.full((n_sls, n_nodes), np.nan)
         else:
-            return weights / np.sum(weights, 0)
+            return np.full((n_sls, n_nodes), 1.0 / n_sls)
 
     if n_sls < 15:  # Cov^-1 unstable under this amount
         return _weighting_failed()
@@ -318,15 +319,20 @@ def gaussian_weights(
 
         flat_coords = sls.reshape(-1, 3)
         flat_groups = working_groups.reshape(-1)
-        unique_ids = np.unique(flat_groups)
+        flat_weights = weights.reshape(-1)
 
-        for gid in unique_ids:
-            mask = flat_groups == gid
-            group_data = flat_coords[mask]
+        order = np.argsort(flat_groups, kind="stable")
+        sorted_groups = flat_groups[order]
+        sorted_coords = flat_coords[order]
+        _, starts = np.unique(sorted_groups, return_index=True)
+        ends = np.append(starts[1:], len(sorted_groups))
 
-            if len(group_data) < 15:
+        for start, end in zip(starts, ends):
+            # small groups will have unstable covariance matrices
+            if end - start < 15:
                 continue
 
+            group_data = sorted_coords[start:end]
             mu = stat(group_data, axis=0)
             diff = group_data - mu
 
@@ -339,31 +345,49 @@ def gaussian_weights(
                 cov = eigenvectors @ np.diag(eigenvalues) @ eigenvectors.T
 
             if np.any(cov > 0):
-                weights.ravel()[mask] = np.sqrt(
-                    np.einsum("ij,jk,ik->i", diff, pinvh(cov), diff)
-                )
+                m = np.einsum("ij,jk,ik->i", diff, pinvh(cov), diff)
+                np.clip(m, 0, None, out=m)
+                flat_weights[order[start:end]] = np.sqrt(m)
 
     if return_mahalanobis:
         return weights
 
-    with np.errstate(divide="ignore"):
+    with np.errstate(divide="ignore", invalid="ignore"):
         w_inv = 1.0 / weights
-    w_inv[np.isinf(w_inv)] = 0
+    w_inv[~np.isfinite(w_inv)] = 0
+
+    denom = np.sum(w_inv, axis=0, dtype=np.float64)
+    w = np.divide(
+        w_inv,
+        denom,
+        out=np.zeros_like(w_inv, dtype=np.float64),
+        where=denom != 0,
+        dtype=np.float64,
+    )
 
-    denom = np.sum(w_inv, axis=0)
-    w = np.divide(w_inv, denom, out=np.zeros_like(w_inv), where=denom != 0)
-    col_sums = w.sum(axis=0)
-    if np.max(np.abs(col_sums - 1)) > 1e-3:
+    col_sums = w.sum(axis=0, dtype=np.float64)
+    if not np.all(np.abs(col_sums - 1) <= 1e-3):
         return _weighting_failed()
     else:
-        final_sums = w.sum(axis=0, keepdims=True)
-        np.divide(w, final_sums, out=w, where=final_sums != 0)
-        return w
+        final_sums = w.sum(axis=0, keepdims=True, dtype=np.float64)
+        np.divide(w, final_sums, out=w, where=final_sums != 0, dtype=np.float64)
+
+        weight_sum = np.sum(w, axis=0, dtype=np.float64)
+        if not np.allclose(weight_sum, np.ones(n_nodes)):
+            bad = ~np.isclose(weight_sum, 1.0, rtol=1e-5, atol=1e-8)
+            logger.warning(f"weights not normalized: {w.shape}, {w.dtype}")
+            logger.warning(f"n_sls: {n_sls}")
+            logger.warning(f"offending nodes: {np.where(bad)[0]}")
+            logger.warning(f"their sums: {weight_sum[bad]}")
+            logger.warning(f"any nan: {np.isnan(w).any()}")
+            del w, w_inv
+            return _weighting_failed()
+    return w
 
 
-def make_gif(show_m, out_path, n_frames=36, az_ang=-10, duration=150):
+def make_mp4(show_m, out_path, n_frames=720, az_ang=-0.5, fps=30, crf=35, verbose=True):
     """
-    Make a video from a Fury Show Manager.
+    Make an MP4 video from a Fury Show Manager with auto-cropping.
 
     Parameters
     ----------
@@ -371,33 +395,48 @@ def make_gif(show_m, out_path, n_frames=36, az_ang=-10, duration=150):
         The Fury Show Manager to use for rendering.
 
     out_path : str
-        The name of the output file.
+        The name of the output file
 
     n_frames : int
         The number of frames to render.
-        Default: 36
+        Default: 720
 
     az_ang : float
         The angle to rotate the camera around the
         z-axis for each frame, in degrees.
-        Default: -10
+        Default: -0.5
+
+    fps : float
+        The frames per second for the output video.
+        Default: 30
 
-    duration : int
-        The duration of each frame in the output GIF, in milliseconds.
-        Default: 150
+    crf : int
+        The Constant Rate Factor for the output video, which controls the
+        quality and file size. Lower values result in
+        higher quality and larger file sizes.
+        Default: 35 (very low quality, small file size)
+
+    verbose : bool
+        Whether to show a progress bar while generating the video.
+        Default: True
     """
+    if not out_path.lower().endswith(".mp4"):
+        out_path += ".mp4"
+
     video = []
 
     show_m.render()
     show_m.window.draw()
 
     with tempfile.TemporaryDirectory() as tmp_dir:
-        for ii in tqdm(range(n_frames), desc="Generating GIF", leave=False):
+        for ii in tqdm(
+            range(n_frames), desc="Generating MP4", leave=False, disable=not verbose
+        ):
             frame_fname = f"{tmp_dir}/{ii}.png"
             show_m.screens[0].controller.rotate((radians(az_ang), 0), None)
             show_m.render()
             show_m.window.draw()
-            show_m.snapshot(frame_fname)
+            show_m.snapshot(fname=frame_fname)
             video.append(Image.open(frame_fname).convert("RGB"))
 
         all_left, all_upper = float("inf"), float("inf")
@@ -406,7 +445,6 @@ def make_gif(show_m, out_path, n_frames=36, az_ang=-10, duration=150):
         for img in video:
             arr = np.array(img)
             bg_color = arr[0, 0]
-
             mask = np.any(arr != bg_color, axis=-1)
 
             if np.any(mask):
@@ -421,20 +459,42 @@ def make_gif(show_m, out_path, n_frames=36, az_ang=-10, duration=150):
                 all_lower = max(all_lower, ymax)
 
         if all_left < all_right:
-            crop_box = (
-                max(0, all_left),
-                max(0, all_upper),
-                min(video[0].width, all_right),
-                min(video[0].height, all_lower),
-            )
+
+            def align_down(x, multiple=16):
+                return (x // multiple) * multiple
+
+            def align_up(x, multiple=16):
+                return ((x + multiple - 1) // multiple) * multiple
+
+            left = align_up(max(0, all_left), 16)
+            upper = align_up(max(0, all_upper), 16)
+            right = align_down(min(video[0].width, all_right), 16)
+            lower = align_down(min(video[0].height, all_lower), 16)
+
+            crop_box = (left, upper, right, lower)
             cropped_video = [img.crop(crop_box) for img in video]
         else:
             cropped_video = video
 
-        cropped_video[0].save(
+        width, height = cropped_video[0].size
+        with imageio.get_writer(
             out_path,
-            save_all=True,
-            append_images=cropped_video[1:],
-            duration=duration,
-            loop=1,
-        )
+            fps=fps,
+            format="ffmpeg",
+            codec="libx264",
+            pixelformat="yuv420p",
+            output_params=[
+                "-crf",
+                f"{str(int(crf))}",
+                "-preset",
+                "veryslow",
+                "-movflags",
+                "+faststart",
+            ],
+        ) as writer:
+            for img in cropped_video:
+                if img.size != (width, height):
+                    img = img.crop((0, 0, width, height))
+
+                frame_arr = np.array(img)
+                writer.append_data(frame_arr)

{pyafq-3.2 → pyafq-3.4}/AFQ/api/bundle_dict.py

@@ -387,7 +387,7 @@ def default_bd():
                 "primary_axis": "I/S",
                 "ORG_spectral_subbundles": SpectralSubbundleDict(
                     {
-                        "Left V1V3": {
+                        "Left Early Visual": {
                             "cluster_IDs": [78],
                             "Left Optic Radiation": {
                                 "core": "Anterior",
@@ -449,7 +449,7 @@ def default_bd():
                 "primary_axis": "I/S",
                 "ORG_spectral_subbundles": SpectralSubbundleDict(
                     {
-                        "Right V1V3": {
+                        "Right Early Visual": {
                             "cluster_IDs": [78],
                             "Right Optic Radiation": {
                                 "core": "Anterior",
@@ -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},
             },

{pyafq-3.2 → pyafq-3.4}/AFQ/api/group.py

@@ -18,7 +18,7 @@ from dipy.io.stateful_tractogram import Space, StatefulTractogram
 from dipy.io.streamline import save_tractogram
 from dipy.utils.parallel import paramap
 from PIL import Image
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 import AFQ.api.bundle_dict as abd
 import AFQ.definitions.image as afm
@@ -555,7 +555,7 @@ class GroupAFQ(object):
                     ses
                 ]
                 seg_sft = aus.SegmentedSFT.fromfile(this_bundles_file, this_img)
-                seg_sft.sft.to_rasmm()
+                seg_sft.to_rasmm()
                 subses_info.append((seg_sft, this_mapping, this_img, this_reg_template))
 
             bundle_dict = self.export("bundle_dict", collapse=False)[
@@ -567,7 +567,7 @@ class GroupAFQ(object):
             for b in bundle_dict.bundle_names:
                 for i in range(len(self.valid_sub_list)):
                     seg_sft, mapping, img, reg_template = subses_info[i]
-                    idx = seg_sft.bundle_idxs[b]
+                    idx = seg_sft.get_bundle_idxs(b)
                     # use the first subses that works
                     # otherwise try each successive subses
                     if len(idx) == 0:
@@ -690,7 +690,14 @@ class GroupAFQ(object):
         for inputs in calcdata.calcs[idx].inputs:
             self.export(inputs)
 
-    def export_all(self, viz=True, afqbrowser=True, xforms=True, indiv=True):
+    def export_all(
+        self,
+        viz=True,
+        afqbrowser=True,
+        xforms=True,
+        indiv=True,
+        delete_full_tractogram=False,
+    ):
         """Exports all the possible outputs
 
         Parameters
@@ -714,6 +721,14 @@ class GroupAFQ(object):
             the AFQ segmentation algorithm, individual ROIs are also
             output.
             Default: True
+        delete_full_tractogram : bool
+            Whether to delete the full, unsegmented tractogram (and its
+            associated json) after all other outputs have been generated.
+            This can be useful for saving disk space, as the full
+            tractogram is typically over 90 percent of the total output size.
+            Only the full tractogram and its json are deleted; all other
+            outputs are left in place.
+            Default: False
         """
         start_time = time()
 
@@ -722,8 +737,19 @@ class GroupAFQ(object):
         self.combine_profiles()
         if afqbrowser:
             self.assemble_AFQ_browser()
+        if delete_full_tractogram:
+            self.delete_full_tractogram()
         self.logger.info(f"Time taken for export all: {str(time() - start_time)}")
 
+    def delete_full_tractogram(self):
+        """
+        Delete the full, unsegmented tractogram (and its associated json)
+        for every subject and session. Only the full tractogram and its
+        json are deleted; all other outputs are left in place.
+        """
+        for pAFQ in self.pAFQ_list:
+            pAFQ.delete_full_tractogram()
+
     def cmd_outputs(
         self, cmd="rm", dependent_on=None, up_to=None, exceptions=None, suffix=""
     ):
@@ -903,11 +929,6 @@ class GroupAFQ(object):
                     )
                 )
                 figure.write_image(this_fname)
-
-                # temporary fix for memory leak
-                import plotly.io as pio
-
-                pio.kaleido.scope._shutdown_kaleido()
             else:
                 from fury import window
 
@@ -917,7 +938,7 @@ class GroupAFQ(object):
                     scene=figure, window_type="offscreen", size=(600, 600)
                 )
                 scene_rotate_forward(show_m, figure)
-                show_m.snapshot(this_fname)
+                show_m.snapshot(fname=this_fname)
 
         def _save_file(curr_img, curr_file_num):
             save_path = op.abspath(
@@ -1427,9 +1448,9 @@ def combine_list_of_profiles(profile_fnames):
     dfs = []
     for fname in profile_fnames:
         profiles = pd.read_csv(fname)
-        profiles["subjectID"] = fname.split("sub-")[1].split("/")[0]
+        profiles["subjectID"] = fname.split("sub-")[1].split("/")[0].split("_")[0]
         if "ses-" in fname:
-            session_name = fname.split("ses-")[1].split("/")[0]
+            session_name = fname.split("ses-")[1].split("/")[0].split("_")[0]
         else:
             session_name = "unknown"
         profiles["sessionID"] = session_name

{pyafq-3.2 → pyafq-3.4}/AFQ/api/participant.py

@@ -1,5 +1,6 @@
 import logging
 import math
+import os
 import os.path as op
 import tempfile
 from math import radians
@@ -8,7 +9,7 @@ from time import time
 import nibabel as nib
 from dipy.align import resample
 from PIL import Image, ImageDraw, ImageFont
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 import AFQ.utils.streamlines as aus
 from AFQ.api.utils import (
@@ -30,7 +31,7 @@ from AFQ.tasks.tractography import get_tractography_plan
 from AFQ.tasks.utils import get_base_fname
 from AFQ.tasks.viz import get_viz_plan
 from AFQ.utils.bin import pyafq_str_to_val
-from AFQ.utils.path import apply_cmd_to_afq_derivs
+from AFQ.utils.path import apply_cmd_to_afq_derivs, drop_extension
 from AFQ.viz.utils import BEST_BUNDLE_ORIENTATIONS, get_eye, trim
 
 __all__ = ["ParticipantAFQ"]
@@ -240,7 +241,9 @@ class ParticipantAFQ(object):
         for inputs in calcdata.calcs[idx].inputs:
             self.export(inputs)
 
-    def export_all(self, viz=True, xforms=True, indiv=True):
+    def export_all(
+        self, viz=True, xforms=True, indiv=True, delete_full_tractogram=False
+    ):
         f""" Exports all the possible outputs
         {valid_exports_string}
 
@@ -262,11 +265,33 @@ class ParticipantAFQ(object):
             the AFQ segmentation algorithm, individual ROIs are also
             output.
             Default: True
+        delete_full_tractogram : bool
+            Whether to delete the full, unsegmented tractogram (and its
+            associated json) after all other outputs have been generated.
+            This can be useful for saving disk space, as the full
+            tractogram is typically over 90 percent of the total output size.
+            Only the full tractogram and its json are deleted;
+            all other outputs are left in place.
+            Default: False
         """
         start_time = time()
         export_all_helper(self, xforms, indiv, viz)
+        if delete_full_tractogram:
+            self.delete_full_tractogram()
         self.logger.info(f"Time taken for export all: {time() - start_time}")
 
+    def delete_full_tractogram(self):
+        """
+        Delete the full, unsegmented tractogram (and its associated json).
+        Only the full tractogram and its json are deleted; all other
+        outputs are left in place.
+        """
+        tractogram_file = self.export("streamlines")
+        for fname in (tractogram_file, drop_extension(tractogram_file) + ".json"):
+            if op.exists(fname):
+                os.remove(fname)
+                self.logger.info(f"Deleted {fname}")
+
     def participant_montage(self, images_per_row=3, anatomy=True, bundle_names=None):
         """
         Generate montage of all bundles for a given subject.
@@ -344,10 +369,6 @@ class ParticipantAFQ(object):
                         showlegend=False,
                     )
                     figure.write_image(this_fname, scale=4)
-                    # temporary fix for memory leak
-                    import plotly.io as pio
-
-                    pio.kaleido.scope._shutdown_kaleido()
                 else:
                     from fury import window
 
@@ -363,7 +384,7 @@ class ParticipantAFQ(object):
                         pass
                     show_m.render()
                     show_m.window.draw()
-                    show_m.snapshot(this_fname)
+                    show_m.snapshot(fname=this_fname)
 
         def _save_file(curr_img):
             save_path = op.abspath(

{pyafq-3.2 → pyafq-3.4}/AFQ/api/utils.py

@@ -7,7 +7,6 @@ from dipy.io.stateful_tractogram import set_sft_logger_level
 from pcollections._lazy import LazyError
 
 from AFQ.utils.docstring_parser import parse_numpy_docstring
-from AFQ.viz.utils import viz_import_msg_error
 
 __all__ = ["methods_descriptors", "kwargs_descriptors", "AFQclass_doc"]
 
@@ -135,6 +134,9 @@ def check_attribute(attr_name):
     if attr_name[:-5] in task_modules:
         return None
 
+    if attr_name in ["tracking_params", "segmentation_params"]:
+        return "segmentation_imap"
+
     if attr_name in methods_sections:
         return f"{methods_sections[attr_name]}_imap"
 
@@ -176,14 +178,7 @@ def export_all_helper(api_afq_object, xforms, indiv, viz):
     api_afq_object.export("profiles")
 
     if viz:
-        try:
-            import IPython  # noqa F401
-            import pingouin  # noqa F401
-            import seaborn  # noqa F401
-        except (ImportError, ModuleNotFoundError):
-            api_afq_object.logger.warning(viz_import_msg_error("plot"))
-        else:
-            api_afq_object.export("tract_profile_plots")
+        api_afq_object.export("tract_profile_plots")
         api_afq_object.export("all_bundles_figure")
         api_afq_object.export("indiv_bundles_figures")
 

{pyafq-3.2 → pyafq-3.4}/AFQ/data/fetch.py

@@ -22,7 +22,7 @@ from dipy.io.streamline import (
 from dipy.segment.clustering import QuickBundles
 from dipy.segment.featurespeed import ResampleFeature
 from dipy.segment.metric import AveragePointwiseEuclideanMetric
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 from AFQ._fixes import get_simplified_transform
 from AFQ.data.utils import aws_import_msg_error
@@ -55,6 +55,7 @@ __all__ = [
     "fetch_brainchop_models",
     "fetch_multiaxial_models",
     "fetch_synthseg_models",
+    "fetch_babyseg_models",
     "fetch_templates",
     "read_templates",
     "fetch_stanford_hardi_tractography",
@@ -263,6 +264,24 @@ fetch_synthseg_models = _make_reusable_fetcher(
     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",