DeepvEM

Code accompanying a manuscript describing DeepvEM: a 3D foundation model for volume electron microscopy (vEM), built by self-supervised masked-autoencoder (MAE) pre-training on a large corpus of vEM volumes, then adapted to downstream image-restoration and analysis tasks by fine-tuning task-specific decoders on the shared pre-trained encoder.

Schematic of the DeepvEM workflow: self-supervised 3D masked-autoencoder pretraining on DeepvEM-1M, evaluation of the resulting 3D image embeddings by clustering and classification, and fine-tuning of task-specific decoders (denoising, super-resolution, segmentation) on the shared pretrained encoder.

---

Demo notebooks

The quickest way to check that the code actually runs end-to-end — installation through a trained checkpoint to a visualized result — is the runnable demo notebook in each fine-tuning task, without needing our cluster or full training runs:

• DeepvEM_DN/deepvem_DN_demo.ipynb — denoising
• DeepvEM_SR/deepvem_SR_demo.ipynb — super-resolution
• DeepvEM_Seg/deepvem_Seg_demo.ipynb — segmentation

Each notebook: loads a pre-trained checkpoint into the task network and sanity-checks the forward pass, converts a small set of volumes into the expected data layout, shows the exact training and inference commands, and renders the model's output next to the input (and ground truth, where applicable).

---

Repository map

Each stage of the pipeline lives in its own self-contained directory, so a reader who only cares about one stage does not need to pull in the others.

Directory	Stage	What's in it
nnssl/	1. Pre-training	Fork of nnssl — self-supervised MAE pre-training of the shared 3D encoder on the DeepvEM-1M corpus. SLURM scripts, dataset conversion, multi-node training.
classification/	2. Embedding evaluation	Extracts embeddings from a pre-trained encoder and evaluates them with k-NN / linear-probe classification and data-efficiency curves against organism / tissue / modality labels.
DeepvEM_DN/	3. Fine-tuning — denoising	Self-supervised denoising decoder (synthetic Poisson-Gaussian noise, residual on a Gaussian baseline), built on nnU-Net.
DeepvEM_SR/	3. Fine-tuning — super-resolution	Anisotropic super-resolution decoder (synthetic per-axis degradation), built on nnU-Net.
DeepvEM_Seg/	3. Fine-tuning — segmentation	Fully-supervised semantic segmentation, standard nnU-Net, optionally warm-started from the pre-trained encoder.
nnUNet/	(shared dependency)	Fork of nnU-Net v2 holding the denoising and super-resolution trainers/inference scripts used by DeepvEM_DN and DeepvEM_SR, plus stock nnU-Net for DeepvEM_Seg.

nnssl and nnUNet are the two frameworks the rest of the tree depends on; classification, DeepvEM_DN, DeepvEM_SR and DeepvEM_Seg are the four things a reader actually runs. Each directory has its own README.md/readme.md with installation, data-layout and usage instructions specific to that stage — start there for anything beyond this overview.

---

Workflow

1. Pre-train the shared 3D encoder with a masked-autoencoder objective on DeepvEM-1M (nnssl/).
2. Evaluate the resulting embeddings directly — no fine-tuning — via clustering and a k-NN / linear-probe classification benchmark against dataset metadata (organism, tissue, imaging modality) (classification/).
3. Fine-tune a task-specific decoder on top of the frozen or warm-started encoder for each downstream task: denoising (DeepvEM_DN/), super-resolution (DeepvEM_SR/), and segmentation (DeepvEM_Seg/).

Every hyperparameter that affects a reported number in any of these stages is a required, explicit setting — nothing is hard-coded, so every run traces back to a value the reader chose (see each stage's README for the exact configuration surface).

---

License

This tree combines forks of two upstream frameworks under their original, differing licenses — check which applies before reusing code from a given directory:

Directory	License	Inherited from
nnssl/	Creative Commons Attribution-NonCommercial 4.0 (non-commercial)	MIC-DKFZ/nnssl
nnUNet/	Apache License 2.0	MIC-DKFZ/nnUNet

classification/, DeepvEM_DN/, DeepvEM_SR/ and DeepvEM_Seg/ are original DeepvEM code with no license file of their own; their use is bound by whichever framework they import:

• classification/extract_embeddings.py imports nnssl directly (to build the network and load a pre-trained checkpoint) and so falls under nnssl's non-commercial license.
• classification/classification_benchmark.py has no framework dependency — it only reads the .npz embeddings extract_embeddings.py produces.
• DeepvEM_DN/, DeepvEM_SR/ and DeepvEM_Seg/ depend only on nnUNet/ (Apache 2.0); they do not import nnssl. Warm-starting their encoder from an nnssl-produced checkpoint consumes a weights file, not nnssl code, and does not itself invoke the non-commercial license.

Full attribution and a file-by-file list of changes relative to each upstream framework are in nnssl/CHANGES_VS_UPSTREAM.md.