README.md

---
license: cc-by-nc-nd-4.0
language:
- en
tags:
- spatial-proteomics
- multiplex imaging
- pathology
- vision
- pytorch
- self-supervised
extra_gated_prompt: >-
  This model and associated code are released under the CC-BY-NC-ND 4.0 license
  and may only be used for non-commercial, academic research purposes with
  proper attribution. Any commercial use, sale, or other monetization of the
  KRONOS2 model and its derivatives, which include models trained on outputs from
  the KRONOS2 model or datasets created from the KRONOS2 model, is prohibited and
  requires prior approval. Please note that the primary email used to sign up
  for your Hugging Face account must match your institutional email to receive
  approval. By downloading the model, you attest that all information
  (affiliation, research use) is correct and up-to-date. Downloading the model
  requires prior registration on Hugging Face and agreeing to the terms of use.
  By downloading this model, you agree not to distribute, publish or reproduce a
  copy of the model. If another user within your organization wishes to use the
  KRONOS2 model, they must register as an individual user and agree to comply
  with the terms of use. Users may not attempt to re-identify the deidentified
  data used to develop the underlying model. If you are a commercial entity,
  please contact the corresponding author.
extra_gated_fields:
  Full name (first and last): text
  Current affiliation (no abbreviations): text
  Type of Affiliation:
    type: select
    options:
    - Academia
    - Industry
    - label: Other
      value: other
  Current and official institutional email (**this must match your primary email in your Hugging Face account, @gmail/@hotmail/@qq email domains will be denied**): text
  Please explain your intended research use in detail: text
  I agree to all terms outlined above: checkbox
  I agree to use this model for non-commercial, academic research purposes only: checkbox
  I agree not to distribute the model, if another user within your organization wishes to use the KRONOS2 model, they must register as an individual user: checkbox
metrics:
- accuracy
pipeline_tag: image-feature-extraction
library_name: kronos2
---
# KRONOS2

> [!NOTE]
> **For end-to-end spatial proteomics workflows use [CORAL](https://github.com/mahmoodlab/CORAL), which operationalizes KRONOS2 and other foundation models for cell phenotyping, patient prognostication, and many more analyses.** 

> [!NOTE]
> [Get CORAL and KRONOS updates](https://forms.gle/8oQWmhCajxv9j7C78)

KRONOS2 is a marker-aware vision foundation model for multiplex immunofluorescence spatial proteomics images, including CODEX/PhenoCycler, CyCIF, and related imaging platforms.

Given a stack of single-marker image patches and the corresponding marker names, KRONOS2 produces one feature embedding per multiplexed patch, conditioned on *which* markers are present. So, a single model handles arbitrary, heterogeneous antibody panels.

KRONOS2 is the next-generation successor to [KRONOS](https://huggingface.co/MahmoodLab/KRONOS).

- **Architecture:** Marker-aware DINOv2 ViT-B/16
- **Embedding dimension:** 768
- **Reference tissue-patch size:** 256 × 256 pixels
- **Pretraining marker vocabulary:** 268 markers
- **Inference:** Evaluation-mode inference is performed in fp32 and is deterministic when run with the same inputs, software environment, and hardware

<p align="center">
  <img src="SPM56M.png" width="600px" alt="Overview of the KRONOS2 pretraining dataset"/>
  <br>
  <em>KRONOS2 pretraining dataset</em>
</p>

## Requesting Access

KRONOS2 is a gated model. To request access, you must agree to the terms of use shown in the Hugging Face access form.

The primary email associated with your Hugging Face account must match your institutional email address. Requests submitted from personal email addresses will be denied.

## Installation

```bash
pip install -r requirements.txt   # pip
# or, with uv (installs the locked environment from pyproject.toml + uv.lock):
uv sync
```

## Example script
The primary usage of KRONOS2 is through companion repo [CORAL](https://github.com/mahmoodlab/CORAL)

The script below shows how KRONOS2 is loaded and used for feature extraction.<br>
It can also be run with `uv run python test.py`

```python
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModel
from sp_image import SPImage   # standalone loader

# 1) load the model
model = AutoModel.from_pretrained("MahmoodLab/KRONOS2", trust_remote_code=True)

# 2) fetch the demo image + its channel names (pulled on demand — the model
#    load itself does NOT download demo_image/)
tiff_path  = hf_hub_download("MahmoodLab/KRONOS2", "demo_image/core.ome.tiff")
names_path = hf_hub_download("MahmoodLab/KRONOS2", "demo_image/channel_names.txt")
with open(names_path) as f:
    channel_names = [line.strip() for line in f if line.strip()]

# 3) load + patch a multiplex image (SPImage only loads / patches / scales)
sp = SPImage(
    tiff_path,                     # (C,H,W) or (cycles,channels,H,W)
    markers=channel_names,         # one name per channel
    mpp=0.37,                      # microns/pixel
)
patches, marker_names, coords = sp.to_patches(
    patch_size=256,
    marker_subset=channel_names,   # the markers to encode (default: all)
)

# 4) marker-aware normalization
patches = model.preprocess(patches, marker_names, preferred_dapi="DRAQ5")

# 5) extract CLS features
feats = []
with torch.inference_mode():
    for k in range(0, len(patches), 16):
        x = torch.from_numpy(patches[k:k + 16]).cuda()
        feats.append(model(x, marker_names).cpu())
features = torch.cat(feats)         # (n_patches, 768)
```

## Marker handling

The marker passed as `preferred_dapi` is treated as the nuclear stain. `model.preprocess` matches each input channel to the pretraining vocabulary 
(268 markers) by an **exact, separator-insensitive key** (`marker_match_key`: lower-case, then
drop separators so `CD-8` = `CD_8` = `CD8`), and assigns per-marker
normalization stats:

- **Matched markers** use their pretraining mean/std. Matching is exact on the
  key — there is **no fuzzy / alias step**, so distinct markers are never
  silently conflated (e.g. `CD24` is not mapped to `CD4`).
- **Unmatched markers** Add each new marker to marker_metadata.csv,
  complete the required metadata fields, and populate the mean and
  std columns using values calculated directly from your dataset.


## Issues

- The preferred mode of communication is via HuggingFace issues.
- If HuggingFace issues are inappropriate, email asong2@mdanderson.org and avaidya@mit.edu.
- Immediate response to minor issues may not be available.

## License & Terms of Use

This model and associated code are released under the CC-BY-NC-ND 4.0 license and may only be used for non-commercial, academic research purposes with proper attribution. Any commercial use, sale, or other monetization of the KRONOS2 model and its derivatives, which include models trained on outputs from the KRONOS2 model or datasets created from the KRONOS2 model, is prohibited and requires prior approval. Downloading the model requires prior registration on Hugging Face and agreeing to the terms of use. By downloading this model, you agree not to distribute, publish or reproduce a copy of the model. If another user within your organization wishes to use the KRONOS2 model, they must register as an individual user and agree to comply with the terms of use. Users may not attempt to re-identify the deidentified data used to develop the underlying model. If you are a commercial entity, please contact the corresponding author.


## Citation

```bibtex
@article{shaban2025foundation,
  title={A foundation model for spatial proteomics},
  author={Shaban, Muhammad and Chang, Yuzhou and Qiu, Huaying and Yeo, Yao Yu and Song, Andrew H and Jaume, Guillaume and Wang, Yuchen and Weishaupt, Luca L and Ding, Tong and Vaidya, Anurag and others},
  journal={arXiv preprint arXiv:2506.03373},
  year={2025}
}
```