gsstephenson/membrane2former-swinv2-large-fibsem-seed42

Membrane2Former SwinV2-L FIB-SEM — Seed 42

35-class organelle segmentation from FIB-SEM electron microscopy images.

Membrane2Former architecture with SwinV2-Large backbone, trained on FIB-SEM volumes from the OpenOrganelle / CellMap dataset.

Architecture	Membrane2Former + SwinV2-L-384
Classes	35 atomic organelle classes
Input	2.5D — 3 adjacent Z-slices at 384×384 px
Best val Dice	0.4504 (epoch 300)
Seed	42

---

Classes (35 atomic)

ecs, pm, mito_mem, mito_lum, mito_ribo, golgi_mem, golgi_lum, ves_mem, ves_lum, endo_mem, endo_lum, er_mem, er_lum, nuc, lyso_mem, lyso_lum, ld_mem, ld_lum, eres_mem, eres_lum, ne_mem, ne_lum, np_out, np_in, hchrom, echrom, nucpl, mt_out, cyto, mt_in, perox_mem, perox_lum, nhchrom, nechrom, nucleo

---

Checkpoint Contents

The best.pt file is a full training checkpoint:

Key	Description
ema_state_dict	EMA model weights — use for inference
model_state_dict	Raw training weights (lower quality than EMA)
optimizer_state_dict	Adam optimizer state — enables fine-tuning resumption
scheduler_state_dict	LR schedule state
scaler_state_dict	AMP scaler state
config	Full training config dict
logit_adj	Logit adjustment tensor (class imbalance correction)
epoch	300
best_dice	0.4504

---

Setup (Longleaf / UNC)

1. Clone the repository

git clone https://github.com/gsstephenson/OrganelleSeg.git
cd OrganelleSeg

2. Activate the environment

micromamba activate csc

3. Compile the CUDA extension (one-time)

pip install -e third_party/MSDeformAttn/

4. Download the checkpoint

from huggingface_hub import hf_hub_download

path = hf_hub_download(
    repo_id="gsstephenson/membrane2former-swinv2-large-fibsem-seed42",
    filename="best.pt",
    local_dir="checkpoints/",
)
print("Downloaded to:", path)

---

3D Volume Generation

Generate EM, GT, and prediction TIFF volumes for a zarr crop:

python scripts/visualization/generate_3d_volumes.py \
    --checkpoints checkpoints/best.pt \
    --names swin_seed42 \
    --dataset jrc_hela-2 \
    --crop crop28 \
    --output-dir /path/to/output \
    --device cuda

Output files:

• em_<dataset>_<crop>.tif — raw EM volume (ZYX, uint8)
• gt_<dataset>_<crop>.tif — ground truth argmax labels (ZYX, uint8)
• pred_swin_seed42_<dataset>_<crop>.tif — model predictions (ZYX, uint8)
• organelle_35cls.lut — ImageJ lookup table
• class_legend.csv — label index → class name + colour

Open in Fiji:

1. File → Open each TIFF
2. Image → Lookup Tables → Load LUT → select organelle_35cls.lut
3. Analyze → Tools → Synchronize Windows to link slices

CLI options:

Flag	Default	Description
--crop-size	auto (384)	Tile size for inference
--overlap	0.5	Tile overlap fraction
--save-probs	off	Also save per-class probability volumes
--no-gaussian	off	Use uniform instead of Gaussian tile blending
--device	cuda	cuda or cpu

---

Multiple Models (Ensemble)

Compare seed42 and seed123 side by side:

python scripts/visualization/generate_3d_volumes.py \
    --checkpoints checkpoints/seed42/best.pt checkpoints/seed123/best.pt \
    --names swin_seed42 swin_seed123 \
    --dataset jrc_hela-2 --crop crop28 \
    --output-dir /path/to/output \
    --device cuda

---

Citation

This model is part of an ongoing study targeting publication in Nature Methods. Citation details will be added upon publication.