[arXiv]
- [2026.06.12] EffiCell-Seg has been accepted by MICCAI 2026 !
- [2026.02.22] We have released the code for EffiCell-Seg !
EffiCell-Seg is a prompt-guided framework for cell instance segmentation built on a frozen DINOv3 backbone. Class-token similarity and PCA of the patch features form a self-prompt that drives a SAM-style prompt encoder and a two-way transformer. The decoder has two branches: a binary foreground branch and a geometric branch, fused by cross-attention, and instances are recovered from the geometric output.
The geometric branch and its post-processing are switchable between four paradigms, so the same architecture can be trained as any of:
| Paradigm | Geometric target (channels) | Instance decoding |
|---|---|---|
tsfd |
normalized distance map (1) | peak detection + watershed |
hovernet |
horizontal / vertical maps (2) | gradient maps + watershed |
stardist |
radial distances (n_rays) |
star-convex polygons + greedy NMS |
cellpose |
flow field, dy/dx (2) | flow tracking + sink clustering |
Two running modes are supported:
online— run the DINOv3 backbone on the fly (optionally with LoRA / CLS-token PEFT).cached— precompute DINOv3 features once withextract_features.py, then train only the decoder. Recommended for very large backbones (e.g. the 7B DINOv3) since the backbone runs only a single time over the dataset.
git clone <GITHUB_URL>
cd EffiCell-Seg
conda create -n efficellseg python=3.10 -y
conda activate efficellseg
# Core dependencies
pip install torch torchvision # PyTorch 2.1+ (CUDA 12.x recommended)
pip install timm peft monai # DINOv3 (timm), LoRA (peft), losses/metrics (monai)
pip install opencv-python scikit-image scipy pandas numpy tqdm
Key requirements: CUDA 12.x, PyTorch 2.1+, a recent timm that ships the DINOv3 weights.
eval.pyimportsget_fast_pq,get_fast_aji,remap_labelfrommetrics. If they are unavailable, instance-level metrics (PQ/SQ/DQ/AJI) fall back to 0 while the binary metrics (Dice/IoU/NSD/HD95) are still computed.
Each dataset lives under ./datasets/<DATASET>/processed:
EffiCell-Seg
├── datasets
│ └── <DATASET>
│ └── processed
│ ├── images
│ │ ├── 0001.png
│ │ ├── ...
│ ├── npy # instance label maps (int), one .npy per image
│ │ ├── 0001.npy
│ │ ├── ...
│ └── ids.json
ids.json lists the split membership (file names or bare IDs):
{
"train": ["0001.png", "..."],
"valid": ["0002.png", "..."],
"test": ["0003.png", "..."]
}
train.pyuses thetrainsplit for training and thetestsplit for validation.extract_features.pyprocesses every split key present inids.json.
(Optional) precomputed features for cached mode are written to
./precomputed_feats/<DATASET>/<split>/<id>.pt.
python extract_features.py \
--input_dir ./datasets/<DATASET>/processed \
--json_path ./datasets/<DATASET>/processed/ids.json \
--output_dir ./precomputed_feats/<DATASET> \
--model_name vit_7b_patch16_dinov3.lvd1689m \
--img_size 512 --batch_size 4 --fp16Training uses DistributedDataParallel, so launch it with torchrun (single GPU works too):
# Single GPU, online mode, TSFD paradigm
torchrun --nproc_per_node=1 train.py \
--dataset <DATASET> --mode online \
--paradigm tsfd \
--backbone timm/vit_7b_patch16_dinov3.lvd1689m \
--img_size 512 --batch 4 --epoch 150 --lr 1e-4 \
--output_dir outputs/efficellseg_<DATASET># Multi-GPU (e.g. 4 GPUs), cached mode, Cellpose paradigm
torchrun --nproc_per_node=4 train.py \
--dataset <DATASET> --mode cached \
--feature_dir ./precomputed_feats \
--paradigm cellpose \
--output_dir outputs/efficellseg_<DATASET>Useful flags: --paradigm {tsfd,hovernet,stardist,cellpose}, --n_rays (stardist),
--geo_weight (override the geometric-loss weight), --peft_strategy {lora,cls_only},
--lora_rank, --num_workers, --sync_bn. Checkpoints best_loss_model.pth,
best_dice_model.pth, and latest_model.pth are saved to --output_dir.
python eval.py \
--dataset <DATASET> --mode online \
--paradigm tsfd \
--backbone vit_7b_patch16_dinov3.lvd1689m \
--model_path outputs/efficellseg_<DATASET>/best_loss_model.pth \
--post_processing instance \
--output_dir visual_test/efficellseg_<DATASET>--paradigmmust match the paradigm the checkpoint was trained with.--post_processing instanceuses the paradigm-aware decoder;simplelabels the binary map with connected components.- Per-image metrics and visualizations (instance maps, boundaries, geometric maps, prompt /
similarity / PCA maps) are written under
--output_dir.
eval.py reports Dice, IoU, NSD, PQ, SQ, DQ, F1, Precision, Recall, AJI, HD95, saving a
per-image CSV and a mean/std summary.
EffiCell-Seg
├── model.py # EffiCell-Seg model: DINOv3 wrapper, self-prompt, dual-branch decoder, per-paradigm post-processing
├── utils.py # shared building blocks (LayerNorm2d, MLP, CrossAttention) + base class
├── dataloader.py # dataset + paradigm-specific geometric-target generation
├── train.py # DDP training with paradigm-aware losses
├── eval.py # inference, metrics, and visualization
└── extract_features.py # optional DINOv3 feature precomputation (for cached mode)
If you find this work helpful, please consider citing:
