spectralbrain.io.gpu_preprocess#

GPU-native preprocessing pipeline for brain MRI.

This module replaces the traditional CPU-bound FreeSurfer/ANTs chain with a fully GPU-accelerated PyTorch pipeline. Each step runs as a separate model load → inference → VRAM purge cycle to keep peak memory within consumer GPU limits (24 GB).

Pipeline stages#

Enhancement — BME-X (Sun et al., 2025, Nat Biomed Eng) or DeepN4 (Kanakaraj et al., 2024) for bias field correction, denoising, and optional super-resolution.
Skull stripping — HD-BET (Isensee et al., 2019) or SynthStrip (Hoopes et al., 2022) for brain extraction.
Registration — SynthMorph affine (Hoffmann et al., 2024) + uniGradICON deformable (Tian et al., 2024) for MNI normalization.
Tissue segmentation — SynthSeg+ (Billot et al., 2023) for GM/WM/CSF + DKT cortical parcellation in a single pass.
Parcellation — OpenMAP-T1 (Nishimaki et al., 2024) for 280 regions or BrainParc (Liu et al., 2026) for 106 lifespan regions.

Usage#

>>> from spectralbrain.io.gpu_preprocess import preprocess_gpu
>>> result = preprocess_gpu(
...     "sub-01_T1w.nii.gz",
...     output_dir="output/",
...     steps=["enhance", "skull_strip", "register", "segment"],
... )
>>> result.brain_path
PosixPath('output/sub-01_T1w_brain.nii.gz')
>>> result.segmentation_path
PosixPath('output/sub-01_T1w_synthseg.nii.gz')

Notes

Based on the gpu_preproc.py pipeline by Rodrigo Dalvit (rdneuro), refactored as a SpectralBrain library module with direct Python API integration where possible and subprocess fallback where required.

Functions

`discover_nifti`(input_path)	Find NIfTI files from a file or directory.
`enhance`(input_path, output_path, *[, method])	Enhance a T1w image using the best available method.
`enhance_bmex`(input_path, output_path, *[, ...])	Enhance a T1w image using BME-X (Sun et al., 2025).
`enhance_deepn4`(input_path, output_path)	GPU bias field correction via DeepN4 (DIPY).
`ensure_template`([seq, user_template])	Resolve or download the MNI registration template.
`parcellate_brainparc`(input_path, output_dir)	106-region lifespan parcellation via BrainParc (Liu et al., 2026).
`parcellate_openmap`(input_path, output_dir)	280-region parcellation via OpenMAP-T1 (Nishimaki et al., 2024).
`preprocess_gpu`(input_path, output_dir, *[, ...])	End-to-end GPU-native preprocessing pipeline.
`preprocess_gpu_batch`(input_paths, output_dir, *)	Process multiple subjects sequentially with VRAM isolation.
`purge_vram`()	Aggressively release all CUDA memory held by this process.
`register`(moving_path, fixed_path, ...[, ...])	Full registration pipeline: optional affine + deformable.
`register_synthmorph_affine`(moving_path, ...)	Affine registration via SynthMorph (Hoffmann et al., 2024).
`register_unigradicon`(moving_path, ...[, ...])	Deformable registration via uniGradICON (Tian et al., 2024).
`segment`(input_path, output_path, *[, method])	Tissue segmentation using the best available method.
`segment_fastsurfer`(input_path, output_dir, ...)	Tissue segmentation via FastSurferCNN (Henschel et al., 2020).
`segment_synthseg`(input_path, output_path, *)	Tissue segmentation + optional DKT parcellation via SynthSeg+.
`skull_strip`(input_path, output_path, *[, method])	Brain extraction using the best available method.
`skull_strip_hdbet`(input_path, output_path, *)	Brain extraction via HD-BET (Isensee et al., 2019).
`skull_strip_synthstrip`(input_path, output_path)	Brain extraction via SynthStrip (Hoopes et al., 2022).
`vram_info`()	Return current VRAM usage in GB.

Classes

`PreprocessResult`(input_path[, ...])	Container for GPU preprocessing pipeline outputs.
`Step`(*values)	Available preprocessing steps.

class spectralbrain.io.gpu_preprocess.PreprocessResult(input_path, enhanced_path=None, brain_path=None, brain_mask_path=None, registered_paths=None, segmentation_path=None, parcellation_path=None, template_path=None, timings=<factory>, methods=<factory>)[source]#

Bases: object

Container for GPU preprocessing pipeline outputs.

Parameters:

input_path (Path)
enhanced_path (Path | None)
brain_path (Path | None)
brain_mask_path (Path | None)
registered_paths (dict[str, Path] | None)
segmentation_path (Path | None)
parcellation_path (Path | None)
template_path (Path | None)
timings (dict[str, float])
methods (dict[str, str])

input_path#

Original input NIfTI.

Type:: Path

enhanced_path#

Bias-corrected / enhanced image.

Type:: Path or None

brain_path#

Skull-stripped brain image.

Type:: Path or None

brain_mask_path#

Brain mask.

Type:: Path or None

registered_paths#

Registration outputs (affine, warped, transforms).

Type:: dict or None

segmentation_path#

Tissue segmentation (aseg labels).

Type:: Path or None

parcellation_path#

Full brain parcellation.

Type:: Path or None

template_path#

Registration template used.

Type:: Path or None

timings#

Per-step wall-clock timings in seconds.

Type:: dict

methods#

Which method was used for each step.

Type:: dict

summary()[source]#

Human-readable summary of the pipeline run.

Return type:: str

brain_mask_path: Path | None = None#

brain_path: Path | None = None#

enhanced_path: Path | None = None#

input_path: Path#

methods: dict[str, str]#

parcellation_path: Path | None = None#

registered_paths: dict[str, Path] | None = None#

segmentation_path: Path | None = None#

template_path: Path | None = None#

timings: dict[str, float]#

property total_time: float#: Total processing time in seconds.

class spectralbrain.io.gpu_preprocess.Step(*values)[source]#

Bases: Enum

Available preprocessing steps.

ENHANCE = 1#

PARCELLATE = 5#

REGISTER = 3#

SEGMENT = 4#

SKULL_STRIP = 2#

spectralbrain.io.gpu_preprocess.discover_nifti(input_path)[source]#

Find NIfTI files from a file or directory.

If a file, returns [file]. If a directory, searches root and first-level subdirectories for .nii.gz and .nii.

Parameters:: input_path (PathLike) – File or directory path.
Returns:: list of Path – Sorted list of NIfTI paths.
Raises:: FileNotFoundError – If path doesn’t exist or no NIfTI files found.
Return type:: list[Path]

spectralbrain.io.gpu_preprocess.enhance(input_path, output_path, *, method='auto', **kwargs)[source]#

Enhance a T1w image using the best available method.

Tries BME-X first (foundation model with bias+denoise+harmonize), falls back to DeepN4 (bias correction only).

Parameters:

input_path (PathLike) – Raw T1w NIfTI.
output_path (PathLike) – Enhanced output.
method ({'bmex', 'deepn4', 'auto'}) – 'auto' tries BME-X, falls back to DeepN4.
**kwargs – Passed to the chosen method.

Returns:

Path – Path to the enhanced image.

Return type:

Path

spectralbrain.io.gpu_preprocess.enhance_bmex(input_path, output_path, *, age_group='adult', mode='enhance')[source]#

Enhance a T1w image using BME-X (Sun et al., 2025).

BME-X performs bias field correction, denoising, and optional super-resolution in a single forward pass using a tissue-aware foundation model.

Parameters:

input_path (PathLike) – Raw T1w NIfTI.
output_path (PathLike) – Enhanced output NIfTI.
age_group (str) – BME-X age model: 'fetal', 'infant', 'adult' (24+ months, valid through 100 years).
mode (str) – Enhancement mode: 'enhance' (bias+denoise), 'super_resolution', 'harmonize'.

Returns:

Path – Path to the enhanced image.

Return type:

Path

Notes

Requires the brain-mri-enhancement package: pip install brain-mri-enhancement

spectralbrain.io.gpu_preprocess.enhance_deepn4(input_path, output_path)[source]#

GPU bias field correction via DeepN4 (DIPY).

Includes the monkey-patch for DIPY’s upstream CUDA bug where __predict calls .numpy() on a CUDA tensor without .cpu() first.

Parameters:

input_path (PathLike) – Raw T1w NIfTI.
output_path (PathLike) – Bias-corrected output NIfTI.

Returns:

Path – Path to the corrected image.

Return type:

Path

spectralbrain.io.gpu_preprocess.ensure_template(seq='t1', user_template=None)[source]#

Resolve or download the MNI registration template.

Parameters:

seq (str) – MRI sequence type ('t1', 't2').
user_template (path or None) – Explicit template path; skips download if provided.

Returns:

Path – Absolute path to the template NIfTI.

Return type:

Path

spectralbrain.io.gpu_preprocess.parcellate_brainparc(input_path, output_dir)[source]#

106-region lifespan parcellation via BrainParc (Liu et al., 2026).

Edge-guided progressive parcellation that works consistently across neonates to elderly without retraining.

Parameters:

input_path (PathLike) – Raw or enhanced T1w NIfTI.
output_dir (PathLike) – Output directory.

Returns:

Path – Path to the parcellation NIfTI (106 labels).

Return type:

Path

spectralbrain.io.gpu_preprocess.parcellate_openmap(input_path, output_dir)[source]#

280-region parcellation via OpenMAP-T1 (Nishimaki et al., 2024).

Covers cortical and subcortical gray matter plus white matter tracts (JHU-MNI atlas). Includes internal skull-stripping, cropping, and hemispheric segmentation.

Parameters:

input_path (PathLike) – Raw or enhanced T1w NIfTI.
output_dir (PathLike) – Output directory.

Returns:

Path – Path to the parcellation NIfTI (280 labels).

Return type:

Path

spectralbrain.io.gpu_preprocess.preprocess_gpu(input_path, output_dir, *, steps=None, enhance_method='auto', strip_method='auto', segment_method='auto', parcellate_method=None, template=None, io_iterations=None, affine_pre=True, device='cuda:0', skip_existing=True)[source]#

End-to-end GPU-native preprocessing pipeline.

Chains enhancement → skull stripping → registration → segmentation (→ parcellation) with VRAM purge between each step. Replaces the traditional FreeSurfer recon-all + ANTs pipeline with a fully GPU-accelerated chain that runs in 2–4 minutes per subject on a 24 GB consumer GPU.

Parameters:

input_path (PathLike) – Raw T1-weighted NIfTI file.
output_dir (PathLike) – Output directory for all products.
steps (list of str or None) – Steps to run: ['enhance', 'skull_strip', 'register', 'segment', 'parcellate']. None = default (all except parcellate).
enhance_method ({'bmex', 'deepn4', 'auto'}) – Enhancement method.
strip_method ({'hdbet', 'synthstrip', 'auto'}) – Skull stripping method.
segment_method ({'synthseg', 'fastsurfer', 'auto'}) – Tissue segmentation method.
parcellate_method ({'openmap', 'brainparc'} or None) – Full brain parcellation method. None = skip.
template (PathLike or None) – Registration template. None = auto-download MNI.
io_iterations (int or None) – uniGradICON instance optimization iterations.
affine_pre (bool) – Run SynthMorph affine before uniGradICON.
device (str) – CUDA device for HD-BET / FastSurfer.
skip_existing (bool) – Skip steps whose outputs already exist.

Returns:

PreprocessResult – Container with all output paths and timings.

Return type:

PreprocessResult

Examples

>>> result = preprocess_gpu(
...     "sub-01_T1w.nii.gz", "output/",
...     steps=["enhance", "skull_strip", "register", "segment"],
... )
>>> print(result.summary())
PreprocessResult for sub-01_T1w.nii.gz
  Total time: 142.3s
  enhance: 12.5s (bmex)
  skull_strip: 8.2s (hdbet)
  register: 18.7s (synthmorph_affine+unigradicon)
  segment: 14.3s (synthseg)

spectralbrain.io.gpu_preprocess.preprocess_gpu_batch(input_paths, output_dir, *, steps=None, enhance_method='auto', strip_method='auto', segment_method='auto', parcellate_method=None, template=None, io_iterations=None, affine_pre=True, device='cuda:0')[source]#

Process multiple subjects sequentially with VRAM isolation.

Parameters:

input_paths (list of PathLike) – NIfTI files to process.
output_dir (PathLike) – Base output directory (per-subject subdirs created).
steps (list[str] | None)
enhance_method (Literal['bmex', 'deepn4', 'auto'])
strip_method (Literal['hdbet', 'synthstrip', 'auto'])
segment_method (Literal['synthseg', 'fastsurfer', 'auto'])
parcellate_method (Literal['openmap', 'brainparc'] | None)
template (str | PathLike | None)
io_iterations (int | None)
affine_pre (bool)
device (str)

Return type:

dict[str, PreprocessResult]

:param : :param parcellate_method: Passed to preprocess_gpu() for each subject. :param template: Passed to preprocess_gpu() for each subject. :param io_iterations: Passed to preprocess_gpu() for each subject. :param affine_pre: Passed to preprocess_gpu() for each subject. :param device: Passed to preprocess_gpu() for each subject.

Returns:

dict of {subject_stem (PreprocessResult}) – Results keyed by filename stem.

Parameters:

input_paths (list[str | PathLike])
output_dir (str | PathLike)
steps (list[str] | None)
enhance_method (Literal['bmex', 'deepn4', 'auto'])
strip_method (Literal['hdbet', 'synthstrip', 'auto'])
segment_method (Literal['synthseg', 'fastsurfer', 'auto'])
parcellate_method (Literal['openmap', 'brainparc'] | None)
template (str | PathLike | None)
io_iterations (int | None)
affine_pre (bool)
device (str)

Return type:

dict[str, PreprocessResult]

spectralbrain.io.gpu_preprocess.purge_vram()[source]#

Aggressively release all CUDA memory held by this process.

Runs two garbage-collection passes to break reference cycles, empties the PyTorch caching allocator, and collects IPC handles. Call between GPU-heavy pipeline steps so each model gets full VRAM access.

Return type:: None

spectralbrain.io.gpu_preprocess.register(moving_path, fixed_path, output_dir, stem, *, affine_method='synthmorph', io_iterations=None)[source]#

Full registration pipeline: optional affine + deformable.

Parameters:

moving_path (PathLike) – Brain-extracted moving image.
fixed_path (PathLike) – Template (fixed) image.
output_dir (PathLike) – Output directory.
stem (str) – Filename stem for outputs.
affine_method ({'synthmorph', 'none'}) – Whether to run SynthMorph affine pre-alignment.
io_iterations (int or None) – uniGradICON IO iterations.

Returns:

dict – Keys: 'affine', 'affine_xfm', 'warped', 'deformable_xfm'. Values are Paths (None if skipped).

Return type:

dict[str, Path]

spectralbrain.io.gpu_preprocess.register_synthmorph_affine(moving_path, fixed_path, output_path, transform_path=None)[source]#

Affine registration via SynthMorph (Hoffmann et al., 2024).

Contrast-invariant affine alignment trained with domain randomization. Ideal as a robust pre-step before deformable registration, especially for large initial misalignments.

Parameters:

moving_path (PathLike) – Moving (subject) brain image.
fixed_path (PathLike) – Fixed (template) image.
output_path (PathLike) – Affine-registered output.
transform_path (PathLike or None) – Output transform file (.lta or .txt).

Returns:

(Path, Path or None) – Registered image and transform file.

Return type:

tuple[Path, Path | None]

spectralbrain.io.gpu_preprocess.register_unigradicon(moving_path, fixed_path, warped_path, transform_path, *, io_iterations=None)[source]#

Deformable registration via uniGradICON (Tian et al., 2024).

Foundation model for medical image registration. Produces a dense displacement field stored as HDF5.

Parameters:

moving_path (PathLike) – Moving (subject) brain image (ideally affine-pre-aligned).
fixed_path (PathLike) – Fixed (template) image.
warped_path (PathLike) – Deformably warped output.
transform_path (PathLike) – Output HDF5 transform.
io_iterations (int or None) – Instance optimization iterations. None = feedforward only (fast); 50 = gradient refinement (more accurate).

Returns:

(Path, Path) – Warped image and transform paths.

Return type:

tuple[Path, Path]

spectralbrain.io.gpu_preprocess.segment(input_path, output_path, *, method='auto', **kwargs)[source]#

Tissue segmentation using the best available method.

Parameters:

input_path (PathLike) – T1w NIfTI.
output_path (PathLike) – Segmentation output.
method ({'synthseg', 'fastsurfer', 'auto'}) – 'auto' tries SynthSeg first.
**kwargs – Passed to the chosen method.

Returns:

Path – Segmentation output path.

Return type:

Path

spectralbrain.io.gpu_preprocess.segment_fastsurfer(input_path, output_dir, subject_id, *, device='cuda:0')[source]#

Tissue segmentation via FastSurferCNN (Henschel et al., 2020).

Produces FreeSurfer-compatible aseg + DKT in < 1 minute on GPU.

Parameters:

input_path (PathLike) – T1w NIfTI.
output_dir (PathLike) – FastSurfer output directory (SUBJECTS_DIR-like).
subject_id (str) – Subject ID for the output directory structure.
device (str) – CUDA device.

Returns:

Path – Path to the output directory containing mri/aparc.DKTatlas+aseg.mgz.

Return type:

Path

spectralbrain.io.gpu_preprocess.segment_synthseg(input_path, output_path, *, parc=True, robust=False, vol_path=None, qc_path=None)[source]#

Tissue segmentation + optional DKT parcellation via SynthSeg+.

Contrast-agnostic segmentation producing FreeSurfer-compatible aseg labels (32 structures) with optional DKT cortical parcellation (31 labels per hemisphere).

Parameters:

input_path (PathLike) – T1w NIfTI (raw or preprocessed — SynthSeg is contrast-agnostic).
output_path (PathLike) – Segmentation output NIfTI.
parc (bool) – Include DKT cortical parcellation (aparc+DKTatlas).
robust (bool) – Use SynthSeg-robust mode (slower, better for clinical scans).
vol_path (PathLike or None) – Output CSV with structure volumes.
qc_path (PathLike or None) – Output CSV with QC scores.

Returns:

Path – Path to the segmentation volume.

Return type:

Path

spectralbrain.io.gpu_preprocess.skull_strip(input_path, output_path, *, method='auto', **kwargs)[source]#

Brain extraction using the best available method.

Parameters:

input_path (PathLike) – Enhanced or raw T1w NIfTI.
output_path (PathLike) – Brain-extracted output.
method ({'hdbet', 'synthstrip', 'auto'}) – 'auto' tries HD-BET, falls back to SynthStrip.
**kwargs – Passed to the chosen method.

Returns:

(Path, Path or None) – Brain image and mask paths.

Return type:

tuple[Path, Path | None]

spectralbrain.io.gpu_preprocess.skull_strip_hdbet(input_path, output_path, *, device='cuda:0', save_mask=True)[source]#

Brain extraction via HD-BET (Isensee et al., 2019).

Parameters:

input_path (PathLike) – Enhanced T1w NIfTI.
output_path (PathLike) – Brain-extracted output NIfTI.
device (str) – CUDA device string.
save_mask (bool) – Whether to save the brain mask.

Returns:

(Path, Path or None) – Paths to the brain image and mask (if saved).

Return type:

tuple[Path, Path | None]

spectralbrain.io.gpu_preprocess.skull_strip_synthstrip(input_path, output_path, mask_path=None)[source]#

Brain extraction via SynthStrip (Hoopes et al., 2022).

Contrast-agnostic skull stripping trained with domain randomization. Works on T1, T2, FLAIR, EPI, DWI, CT.

Parameters:

input_path (PathLike) – Input NIfTI (any contrast).
output_path (PathLike) – Brain-extracted output.
mask_path (PathLike or None) – Brain mask output path.

Returns:

(Path, Path or None) – Brain image and mask paths.

Return type:

tuple[Path, Path | None]

spectralbrain.io.gpu_preprocess.vram_info()[source]#

Return current VRAM usage in GB.

Returns:: dict – Keys: 'allocated', 'reserved', 'total', 'free'. All values in GB. Returns zeros if CUDA is unavailable.
Return type:: dict[str, float]