Modelpublic

gsstephenson/membrane2former-swinv2-large-fibsem-seed123

sourceHugging Faceapache-2.0updated 6mo agoView on Hugging Face

0likes

Model Card

Membrane2Former SwinV2-L FIB-SEM — Seed 123

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.4420 (epoch 315)
Seed	123

Tip: For an ensemble, combine predictions from seed42 (best_dice=0.4504) and this seed123 run — seeds differ only in random initialization, giving complementary errors.

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`	315
`best_dice`	0.4420

Setup (Longleaf / UNC)

1. Clone the repository

bash

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

2. Activate the environment

bash

micromamba activate csc

3. Compile the CUDA extension (one-time)

bash

pip install -e third_party/MSDeformAttn/

4. Download the checkpoint

python

from huggingface_hub import hf_hub_download

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

3D Volume Generation

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

bash

python scripts/visualization/generate_3d_volumes.py \
    --checkpoints checkpoints/best.pt \
    --names swin_seed123 \
    --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_seed123_<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:

bash

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.