Skip to content

Repository files navigation

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

One tutorial per dataset (tutorials/<config>). "PyTC" is the released checkpoint (linked to its weights on Hugging Face), "SOTA" the best published score on the same data and metric; the better of the two is bold. EMNeuron: mean VI over the three out-of-distribution volumes with DECODER=agglo, the best of the three decoders (added after scoring; 0.900 with the pre-specified default; see the tutorial README).

Structure Dataset Tutorial config PyTC Prior SOTA
Neurons EMNeuron (SegNeuron) neuron_segneuron/* VI 0.821 ↓ VI 0.887 (SegNeuron)
Neurons NISB-base neuron_nisb/* NERL 0.601 ↑ Ours
Neurons SNEMI3D neuron_snemi/* — —
Mitochondria Lucchi++ mito_lucchi++/mito_lucchi++ Jaccard 0.935 ↑ Ours
Mitochondria MitoEM mitoEM/* — —
Mitochondria MitoLab, BetaSeg mito_mitolab, mito_betaseg — —
Synapses CREMI syn_cremi — —
Vesicles XM vesicle_xm — —
Nuclei NucMM-Z nuc_nucmm-z — —
Blood vessel — (coming) — —

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

Neurons in your own EM volume, no training: the pretrained EMNeuron model runs zero-shot with tutorials/neuron_segneuron/predict.sh and fine-tunes with tutorials/neuron_segneuron/finetune.sh; see its README.


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:


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.

Releases

Packages

Used by

Contributors

Languages