MAESTRO

MAESTRO is an nnUNet-based model for stroke lesion segmentation using T1-weighted MRI scans. This repository provides scripts and instructions to easily set up the environment and run inference on your own data.

Requirements

• Python 3.9+
• nnUNetv2==2.8.1
• PyTorch (must be installed according to your system's CUDA/compute capabilities)
• NIFTI files corresponding to T1-weighted scans of brains

Installation and Initial Setup

1. Install nnUNet_v2

Follow the official instructions to install nnUNet_v2 from their GitHub repository: nnUNet_v2 Installation Guide

Please note the specific instructions to install PyTorch first before installing nnUNet_v2. MAESTRO was built and tested using nnUNet_v2 version 2.8.1.

2. Download Model Weights

1. Download this repository (git clone https://github.com/npnl/MAESTRO).
2. Download the pre-trained MAESTRO weights.
3. Unzip the file into this directory. The path should be MAESTRO/model_weights.

3. Using non-default paths

The config.sh file has been initialized to default paths within this repository. Edit the config.sh file if you wish to save the model weights or custom nnUNet trainer elsewhere. To initialize them for your session, run:

source config.sh

4. Custom Trainers

This model utilizes a custom trainer (nnUNet_PrimusV3S_NoDA_Trainer). We ship this custom trainer in the custom_trainers/maestro_pkg directory. The config.sh script automatically sets the nnUNet_extTrainer environment variable so that nnUNetv2 will discover it seamlessly!

Running Inference

To get predictions on your own T1-weighted MRI data, follow these steps:

Step 1: Organize Data

nnUNet expects input files to have specific suffixes (e.g., _0000.nii.gz for the primary T1 modality). Run the organization script to prepare your data:

./organize_data.sh <path_to_raw_nifti_folder> <path_to_organized_output_folder>

Alternatively, you are welcome to organize the data yourself. The requirements are that all niftis are in one folder, and they each end with _0000.nii.gz. Please note that you can use ln -s to link the files instead of copying them if you wish to save space. We also recommend each scan to have a unique identifier prepending to the filename (e.g., sub-1_ses-1_0000.nii.gz).

Step 2: Run Inference

Execute the inference script, passing the organized folder from Step 1 as the input. You can optionally specify the hardware method (gpu or cpu) and device ID (if using gpu):

# Run on GPU 0 (default)
./run_inference.sh -i <path_to_organized_output_folder> -o <path_to_predictions_folder> -m gpu -d 0

# Or run on CPU
./run_inference.sh -i <path_to_organized_output_folder> -o <path_to_predictions_folder> -m cpu

Based on our internal testing, we expect MAESTRO to take roughly 5-10 seconds per scan when running on GPU, and roughly 10-15 minutes per scan when running on CPU.

Step 3: Convert Output (Optional)

If you'd like to standardize the output names back from nnUNet's expected format (e.g., appending _pred.nii.gz), run:

./convert_output.sh <path_to_predictions_folder> <path_to_final_output_folder>

Literature

The preprint link will be available soon.

Documentation

For more detailed information on nnUNet, please refer to the official nnUNet repository.