Skip to content

PytorchConnectomics/pytorch_connectomics

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1,502 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

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.


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, LiConn-MIT neuron_snemi/*, neuron_nisb/*, neuron_liconn_mit
Cell-scale Synapses CREMI syn_cremi
Cell-scale Fibers Linghu26 fiber_linghu26
Organelle-scale Mitochondria Lucchi++, MitoEM, MitoLab, BetaSeg mito_lucchi++, mitoEM/*, mito_mitolab, mito_betaseg
Organelle-scale Vesicles XM vesicle_xm

Sample data downloads in <1 min for lucchi++, snemi, mitoem, cremi (just download <name>).


Benchmarks

Headline metric per public benchmark. Full tables, training curves, and pretrained checkpoints live in docs/benchmarks/.

Dataset Task Architecture Metric Score
Lucchi++ Mito — semantic Swin UNETR Jaccard ↑
MitoEM-R Mito — instance MedNeXt-L + waterz AP ↑
SNEMI3D Neurons — instance RSUNet + waterz adapted Rand ↓
BANIS Neurons — instance MedNeXt-M + abiss NERL ↑
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/zudi-lin/pytorch_connectomics/refs/heads/master/quickstart.sh | bash
cd pytorch_connectomics
conda activate pytc

b) LLM-assisted installation

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

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

c) Manual installation

git clone https://github.com/zudi-lin/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/main.py --demo

Recipes

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

just download lucchi++              # ~50 MB
just train     mito_lucchi++        # train from scratch
just test      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++.yaml tutorials/my_mito.yaml
# edit data.{train,val}.{image,label} paths inside my_mito.yaml
just train my_mito

…or override on the CLI without copying:

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

3. Fine-tune from a published checkpoint:

just resume my_mito <pretrained.ckpt>

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

just test my_mito <ckpt> evaluation.enabled=false

5. Sweep decode params with Optuna:

python scripts/main.py --config tutorials/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> data.dataloader.batch_size=4 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.

Packages

 
 
 

Contributors

Languages