nnU-Net Axon/Myelin Segmentation Pipeline

Pipeline for training and evaluating nnU-Net models on axon and myelin segmentation from electron microscopy images.

Repository Structure

mikula_nnunet/
├── scripts/              # Main pipeline scripts (REQUIRED)
│   ├── prepare_dataset.py    # PNG → NIfTI conversion for nnU-Net
│   ├── eval_dice.py          # Dice score evaluation (MONAI-based)
│   └── train_nnunet.ps1      # Complete training pipeline (PowerShell)
│
├── utils/                # Conversion utilities (OPTIONAL)
│   ├── nii_to_png_images.py  # NIfTI → PNG (images)
│   ├── nii_to_png_masks.py   # NIfTI → PNG (masks with color mapping)
│   ├── extract_tif_layers.py # Extract layers from TIFF stacks
│   └── centered_crop.py       # Crop images to 421x421
│
└── validation/           # Quality check scripts (OPTIONAL)
    ├── check_pair_alignment.py     # Validate image-mask pairs
    └── check_pred_gt_alignment.py  # Validate predictions vs GT

Quick Start

1. Environment Setup

Set nnU-Net environment variables:

$env:nnUNet_raw = "C:\path\to\nnUNet_raw"
$env:nnUNet_preprocessed = "C:\path\to\nnUNet_preprocessed"
$env:nnUNet_results = "C:\path\to\nnUNet_results"

2. Data Preparation

Convert your PNG images and masks to nnU-Net format (nii.gz):

python scripts/prepare_dataset.py

3. Training Pipeline

Run the complete training and evaluation pipeline:

# Edit train_nnunet.ps1 to adjust paths, then run:
.\scripts\train_nnunet.ps1

Or run steps individually (see train_nnunet.ps1 for details):

• Verify dataset integrity
• Plan and preprocess
• Train model (default: 1000 epochs, auto-determined by nnU-Net)
• Generate predictions
• Evaluate with Dice metric

4. Manual Evaluation

Compute Dice scores using MONAI:

python scripts/eval_dice.py <pred_dir> <gt_dir> --mode 3c

Modes:

• 3c: 3-class segmentation (background, myelin, axon)
• bin: Binary segmentation (background, axon)

Example Results

Training Set (36 images):
- Background: 0.9987 ± 0.0047
- Myelin:     0.9499 ± 0.0632
- Axon:       0.9765 ± 0.0338

Validation Set (16 images):
- Background: 0.9681 ± 0.0416
- Myelin:     0.7639 ± 0.0702
- Axon:       0.8016 ± 0.0677

Requirements

Python packages:

nnunetv2
monai
torch
numpy
SimpleITK
Pillow

Notes

• Dice scores computed using MONAI's compute_dice for standardization
• Compatible with nnU-Net v2 format
• Supports 2D segmentation tasks