PyTorch Connectomics

Modern deep learning for 2D / 3D connectomics. Train, run inference, decode, and evaluate segmentation pipelines on large EM volumes — from a single GPU to multi-node clusters.

---

Research workflows are developed separately and are not part of this distribution.

What you can segment

Scale	Structure	Datasets	Tutorial config
Tissue-scale	Blood vessel	—	(coming)
Tissue-scale	Nuclei	NucMM-Z	nuc_nucmm-z
Cell-scale	Neurons	SNEMI3D, BANIS	neuron_snemi/*, neuron_nisb/*
Cell-scale	Synapses	CREMI	syn_cremi
Organelle-scale	Mitochondria	Lucchi++, MitoEM, MitoLab, BetaSeg	mito_lucchi++/mito_lucchi++, mitoEM/*, mito_mitolab, mito_betaseg
Organelle-scale	Vesicles	XM	vesicle_xm

Download sample data for lucchi++, snemi, mitoem, or cremi with just download <name>.

---

Benchmarks

Headline metric per public benchmark. Full tables, training curves, and pretrained checkpoints are linked in the table below.

Dataset	Task	Architecture	Metric	Score
Lucchi++	Mito — semantic	MedNeXt-S	Jaccard ↑	0.935
MitoEM-R	Mito — instance	MedNeXt-L + waterz	AP ↑	—
SNEMI3D	Neurons — instance	MedNeXt-S + waterz	adapted Rand ↓	—
BANIS	Neurons — instance	MedNeXt-L + cc3d	NERL ↑	0.601
CREMI	Synapse — semantic	MONAI U-Net	F1 ↑	—

---

Quick Start [more details]

Pick one install path, then verify.

a) Auto installation

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

b) LLM-assisted installation

git clone https://github.com/PytorchConnectomics/pytorch_connectomics.git
cd pytorch_connectomics
just install-claude          # or: just install-codex
conda activate pytc

Reads prompts/INSTALL.md and lets the agent drive install.py. Requires an authenticated claude or codex CLI.

c) Manual installation

git clone https://github.com/PytorchConnectomics/pytorch_connectomics.git
cd pytorch_connectomics
python install.py --install-type basic --python 3.11 --env-name pytc
conda activate pytc

See INSTALLATION.md for CUDA versions, extras, and troubleshooting.

Verify

python scripts/check_install.py   # imports + GPU kernel check; must print PASS
python scripts/main.py --demo     # ~30 s CPU training on synthetic data

---

Recipes

1. Train + evaluate on a built-in benchmark — sample data is tiny:

just download lucchi++                             # ~211 MiB
just train     mito_lucchi++/mito_lucchi++        # train from scratch
just test      mito_lucchi++/mito_lucchi++ <ckpt> # infer + decode + evaluate
just tensorboard mito_lucchi++                     # monitor

2. Train on your own EM volume — copy the closest tutorial, point at your data:

cp tutorials/mito_lucchi++/mito_lucchi++.yaml tutorials/<name>.yaml
# edit train.data.{train,val}.{image,label} paths inside the new YAML
just train <name>

…or override on the CLI without copying:

python scripts/main.py --config tutorials/mito_lucchi++/mito_lucchi++.yaml \
    train.data.train.image=/path/to/train.h5 \
    train.data.train.label=/path/to/label.h5

3. Fine-tune from a published checkpoint:

just resume <name> <pretrained.ckpt>

4. Predict on a new volume (no labels needed):

just test <name> <ckpt> test.evaluation.enabled=false

5. Sweep decode params with Optuna:

python scripts/main.py --config tutorials/mito_lucchi++/mito_lucchi++.yaml \
    --mode tune --checkpoint <ckpt>

6. Drive a workflow with a coding agent — agent reads a prompt and runs the steps for you (interactive):

just add-dataset-claude       # or: just add-dataset-codex
just add-arch-claude          # or: just add-arch-codex
just debug-tutorial-claude    # or: just debug-tutorial-codex

Requires an authenticated claude or codex CLI.

---

Under the hood

Five composable stages, each with its own config section and entry point. Data flows left-to-right; each stage owns one transformation:

EM volume ─[train]─▶ checkpoint ─[infer]─▶ representation ─[decode]─▶ instances ─[evaluate]─▶ metrics
                                            (aff / dist / flow)       (cc3d / waterz / abiss)   (Rand / VOI / NERL)
                                                                              ▲
                                                                        [tune]┘  Optuna over decode params

Single CLI dispatches all of them; any field is overridable from the command line:

python scripts/main.py --config <cfg> train.data.dataloader.batch_size=4 train.optimization.max_epochs=200

---

Acknowledgments

PyTC stands on the work of many open-source projects. Organized by what they enable here:

• High-performance pipeline — PyTorch Lightning, MONAI, Zarr, HDF5
• State-of-the-art models — Swin UNETR, MedNeXt, nnU-Net, Cellpose, BANIS
• Effective decoding — cc3d, waterz, abiss
• Reproducible experiments — Hydra, Optuna
• Agent-driven development — Claude Code, Codex, orchestrated via ccc-agent-flow

---

Citation

@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}
}

Actively developed and maintained by Donglai Wei's group at Boston College. Supported by NSF IIS-1835231, IIS-2124179, IIS-2239688.

MIT License — see LICENSE.