Automatic segmentation of neural structures in electron microscopy images 🔬🧠

PyTorch Connectomics (PyTC) helps neuroscientists and biologists:

• ✅ Initial run: Ready-to-use segmentation models for strong initial segmentation
• ✅ Development: Easy to build and adapt models to your annotated data
• ✅ Deployment: Scales from a single GPU to large-scale clusters

Built on: PyTorch Lightning + MONAI + nnU-Net for modern, scalable deep learning.

Choose your preferred method:

curl -fsSL https://raw.githubusercontent.com/zudi-lin/pytorch_connectomics/refs/heads/master/quickstart.sh | bash
conda activate pytc

git clone https://github.com/PytorchConnectomics/pytorch_connectomics.git
cd pytorch_connectomics
python install.py
conda activate pytc

conda create -n pytc python=3.10 -y
conda activate pytc
conda install -c conda-forge numpy h5py cython connected-components-3d -y
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
git clone https://github.com/PytorchConnectomics/pytorch_connectomics.git
cd pytorch_connectomics
pip install -e . --no-build-isolation

📖 Detailed instructions: INSTALLATION.md | 🚀 Quick start: QUICKSTART.md

Verify your installation with a 30-second demo:

python scripts/main.py --demo

Expected output:

🎯 PyTorch Connectomics Demo Mode
...
✅ DEMO COMPLETED SUCCESSFULLY!

# Download tutorial data (~50 MB)
just download lucchi++

• Ours: 0.913 Jaccard index (Foreground-IoU)
• Prior art: 0.907 (systematic comparsions in Casser et.al 20)

# Download pretrained checkpoint
mkdir -p checkpoints
curl -L "https://huggingface.co/pytc/tutorial2.0/resolve/main/mito_lucchi%2B%2B_15k.ckpt?download=true" \
  -o checkpoints/mito_lucchi_pp_15k.ckpt

# Run inference
just test mito_lucchi++ checkpoints/mito_lucchi_pp_15k.ckpt

• current settings: num_gpus=4, num_cpus=8, batch_size=16
• change settings in tutorials/mito_lucchi++.yaml

# Quick test (1 batch)
just train mito_lucchi++ --fast-dev-run

# Full training
just train mito_lucchi++

Monitor progress:

just tensorboard lucchi++

Run inference from your trained checkpoint:

just test lucchi++ outputs/lucchi++/$EXPERIMENT_DATE/checkpoints/best.ckpt

• PyTorch Lightning: Automatic distributed training, mixed precision, callbacks
• MONAI: Medical imaging models, transforms, losses optimized for 3D volumes
• Hydra/OmegaConf: Type-safe configurations with CLI overrides
• Extensible: Easy to add custom models, losses, and transforms

• MONAI Models: BasicUNet3D, UNet, UNETR, Swin UNETR
• MedNeXt (MICCAI 2023): ConvNeXt-based architecture for medical imaging
• Custom Models: Easily integrate your own architectures

• Distributed Training: Automatic multi-GPU with DDP
• Mixed Precision: FP16/BF16 training for 2x speedup
• Efficient Data Loading: Pre-loaded caching, MONAI transforms
• Gradient Accumulation: Train with large effective batch sizes

• Waterz Decoder: Watershed + hierarchical agglomeration on affinity graphs with tunable scoring functions (aff50_his256, aff85_his256, etc.)
• Dust Merge: Zwatershed-style size+affinity cleanup — merges small fragments into their best-affinity neighbor instead of dropping to background (C++/Cython)
• Optuna Tuning: Batch threshold sweep in a single waterz call (watershed computed once) for efficient hyperparameter search
• Experiment Auto-Logging: Every decode+eval run auto-appends parameters and metrics to decode_experiments.tsv

• TensorBoard: Training curves, images, metrics
• Weights & Biases: Experiment tracking (optional)
• Early Stopping: Automatic stopping when training plateaus
• Checkpointing: Save best models automatically

• 🚀 Quick Start Guide - Get running in 5 minutes
• 📦 Installation Guide - Detailed setup instructions
• 📚 Full Documentation - Complete reference
• 🎯 Tutorials - Example configurations
• 🔧 Troubleshooting - Common issues and solutions
• 👨‍💻 Developer Guide - Contributing and architecture

• tutorials/*.yaml: runnable experiment configs
• connectomics/config/profiles/*.yaml: section-level registries selected by *.profile
• connectomics/config/templates/*.yaml: explicit list-item templates, currently used for inference.decoding

Merge rule:

• Profile payloads are merged into the target section as the base config.
• Explicit keys in the tutorial/config override profile keys.
• Explicit lists replace profile lists; they are not additive.
• Canonical decoding syntax is explicit list expansion, for example inference.decoding: [{template: decoding_waterz}].

Create a config file (my_config.yaml):

system:
  training:
    num_gpus: 1
    num_cpus: 4
    batch_size: 2

model:
  architecture: monai_basic_unet3d
  in_channels: 1
  out_channels: 2
  loss_functions: [DiceLoss]

data:
  train_image: "path/to/train_image.h5"
  train_label: "path/to/train_label.h5"
  patch_size: [128, 128, 128]

optimization:
  max_epochs: 100
  precision: "16-mixed"  # Mixed precision for speed

optimizer:
  name: AdamW
  lr: 1e-4

Train:

python scripts/main.py --config my_config.yaml

Override from CLI:

python scripts/main.py --config my_config.yaml data.dataloader.batch_size=4 optimization.max_epochs=200

Configure instance segmentation decoding in your YAML config:

inference:
  decoding:
    - name: decode_waterz
      kwargs:
        thresholds: 0.4                  # agglomeration threshold
        merge_function: aff85_his256     # 85th percentile affinity scoring
        aff_threshold: [0.1, 0.99]       # watershed low/high thresholds
        channel_order: xyz               # affinity channel order (auto-transpose to zyx)
        dust_merge_size: 100             # merge dust < 100 voxels into best neighbor
        dust_merge_affinity: 0.0         # min affinity for dust merge
        dust_remove_size: 50             # remove remaining orphans < 50 voxels

Run inference with decoding:

python scripts/main.py --config tutorials/neuron_snemi.yaml --mode test --checkpoint path/to/best.ckpt

Tune decode parameters with Optuna:

python scripts/main.py --config tutorials/neuron_snemi.yaml --mode tune --checkpoint path/to/best.ckpt

Results are auto-logged to outputs/<experiment>/results/decode_experiments.tsv with all parameters and metrics.

• BasicUNet3D - Fast, simple 3D U-Net (recommended for beginners)
• UNet - U-Net with residual units
• UNETR - Transformer-based architecture
• Swin UNETR - Swin Transformer U-Net

• MedNeXt-S - 5.6M parameters (fast)
• MedNeXt-B - 10.5M parameters (balanced)
• MedNeXt-M - 17.6M parameters (accurate)
• MedNeXt-L - 61.8M parameters (state-of-the-art)

See: .claude/MEDNEXT.md for MedNeXt integration guide

• HDF5 (.h5) - Primary format (recommended)
• TIFF (.tif, .tiff) - Multi-page TIFF stacks
• Zarr - For large-scale datasets
• NumPy - Direct array loading

Input shape: (batch, channels, depth, height, width)

• 💬 Slack: Join our community (friendly and helpful!)
• 🐛 Issues: GitHub Issues
• 📧 Contact: See lab website
• 📄 Paper: arXiv:2112.05754

If PyTorch Connectomics helps your research, please cite:

@article{lin2021pytorch,
  title={PyTorch Connectomics: A Scalable and Flexible Segmentation Framework for EM Connectomics},
  author={Lin, Zudi and Wei, Donglai and Lichtman, Jeff and Pfister, Hanspeter},
  journal={arXiv preprint arXiv:2112.05754},
  year={2021}
}

Powered by:

• PyTorch Lightning - Lightning AI Team
• MONAI - MONAI Consortium
• MedNeXt - DKFZ Medical Image Computing

Supported by:

• NSF awards IIS-1835231, IIS-2124179, IIS-2239688

MIT License - See LICENSE for details.

Copyright © PyTorch Connectomics Contributors

• v2.0 (2025) - Complete rewrite with PyTorch Lightning + MONAI
• v1.0 (2021) - Initial release

See RELEASE_NOTES.md for detailed release notes.