CytoFormer
CytoFormer assigns a cell type to every nucleus on a routine H&E slide. It was trained on 15.4 million cells from 81 paired Xenium / H&E sections across 16 organs, with the cell types derived from the spatial transcriptome of the same section rather than from manual annotation.
Inference code, usage examples and an online sample viewer: https://github.com/zhihuanglab/CytoFormer
Files
| file | what it is |
|---|---|
checkpoint.pth | the model checkpoint |
organ_celltype_map.json | which cell types each organ can predict |
Model
| Input | one 56 µm field of view centred on the nucleus, resized to 224×224 px |
| Encoder | ViT-giant, patch 14, 1536-d embedding, initialised from UNI2-h and fine-tuned end-to-end |
| Head | 16 linear classifiers, one per organ, each over that organ's cell types only (103 outputs in total); the organ identifier selects which head is used and is not given to the encoder |
| Output | one of 23 global cell types, restricted to the cell types that occur in the given organ |
The organ is a routing signal, not an image feature: the same representation is produced whatever
organ is declared, so the encoder can also be used as a general cell-level feature extractor
(extract_features returns the 1536-d cell embedding).
Usage
git clone https://github.com/zhihuanglab/CytoFormer && cd CytoFormer
pip install -r requirements.txt
mkdir -p checkpoints && cp cytoformer/organ_celltype_map.json checkpoints/
hf download zhihuanglab/CytoFormer checkpoint.pth --local-dir checkpoints
import torch
from cytoformer import CellClassifier, ORGAN_IDX
net = CellClassifier()
net.load_state_dict(torch.load("checkpoints/checkpoint.pth", map_location="cpu"))
net.eval()
x = torch.randn(2, 3, 224, 224) # ImageNet-normalised 56 µm patches
logits = net(x, torch.tensor([ORGAN_IDX["skin"]] * 2)) # (2, 23), out-of-organ classes masked
emb = net.extract_features(x) # (2, 1536) cell embedding
Intended use
Research use only.