Hugging Face's logo

rad-dit
/
jit-radimagenet

Unconditional Image Generation
PyTorch
diffusion
medical-imaging
radimagenet
Model card Files
xet
 Community

You need to agree to share your Hugging Face information to access this model

This repository is publicly accessible, but you have to accept the conditions to access its files and content.

These models were trained on RadImageNet and are released for non-commercial research only (CC BY-NC 4.0). If you use the models or images generated with them, you must cite our paper and the papers in the Citation section. The models are not for clinical use.

Log in or Sign Up to review the conditions and access this model content.

rad-dit: JiT models on RadImageNet

Pixel-space diffusion transformers (JiT), initialised with ImageNet weights and trained on RadImageNet (CT and MR, 256×256).

Available models

Folder Model Training data Epochs gFID ↓
B16 JiT-B/16 (131M) RadImageNet 80 36.48
L16 JiT-L/16 (459M) RadImageNet 80 32.34
H16 JiT-H/16 (952M) RadImageNet 40 32.70
L16-mix15 JiT-L/16 RadImageNet + 15% ImageNet 80 30.94
L16-mix25 JiT-L/16 RadImageNet + 25% ImageNet 80 31.01

gFID: 256×256, cfg 1.0, Heun with 50 steps (see the paper). Each checkpoint-last.pth holds model, model_ema1 and model_ema2, as the official JiT checkpoints. Sampling uses model_ema1. The optimizer state is not included.

Architecture: the checkpoints hold only the weights. The architecture comes from the command-line arguments in the usage below (--model, --img_size and the label arguments).

Classes

Label Class
1001 CT abdomen
1002 CT lung
1003 MR hip
1004 MR knee
1005 MR shoulder
1000 No class (classifier-free guidance row)
0–999 ImageNet classes (rows from the ImageNet initialisation; only the mixed models kept training on them)

Use --label_mode extended --label_offset 1001 --class_num 1005 --cfg_drop_index 1000.

Usage

1. Get the code and log in

git clone https://github.com/CAIROLab-Bern/rad-dit.git && cd rad-dit
# Install the environment: see environment/README.md
hf auth login   # accept the terms on this page first

2. Download a model

hf download rad-dit/jit-radimagenet --include "L16/*" --local-dir jit/checkpoints/radimagenet

Or in Python:

from huggingface_hub import hf_hub_download
hf_hub_download("rad-dit/jit-radimagenet", "L16/checkpoint-last.pth", local_dir="jit/checkpoints/radimagenet")

3. Generate images (no dataset needed)

cd jit
torchrun --nproc_per_node=1 main_jit.py \
    --model JiT-L/16 --img_size 256 --noise_scale 1.0 \
    --label_mode extended --label_offset 1001 --class_num 1005 \
    --cfg_drop_index 1000 --eval_class_num 5 \
    --evaluate_gen --resume checkpoints/radimagenet/L16 \
    --num_images 60 --gen_bsz 15 \
    --num_sampling_steps 50 --sampling_method heun \
    --cfg 1.0 --interval_min 0.1 --interval_max 1.0 \
    --output_dir samples/L16

The model makes the same number of images for each class, so --num_images must be a multiple of --eval_class_num. For one class only, for example MR knee, use --label_offset 1004 --eval_class_num 1.

For the other models, download their folder in step 2 and change --model and --resume:

Folder --model --resume
B16 JiT-B/16 checkpoints/radimagenet/B16
L16 JiT-L/16 checkpoints/radimagenet/L16
H16 JiT-H/16 checkpoints/radimagenet/H16
L16-mix15 JiT-L/16 checkpoints/radimagenet/L16-mix15
L16-mix25 JiT-L/16 checkpoints/radimagenet/L16-mix25

4. gFID: see jit/README.md in the code repository.

Terms of use

  • Use the models for non-commercial research only (CC BY-NC 4.0).
  • Do not use the models or their images for diagnosis or other clinical decisions.
  • Generated images are synthetic. They are not real patient data.

Citation

If you use these models or images generated with them in your research, please cite our paper and consider giving the code repository a ⭐:

@inproceedings{georgiou2027natural,
  title     = {From Natural Images to Medical Scans: Transferability of Diffusion Models for Medical Image Generation},
  author    = {Georgiou, Efthymios and Kakogeorgiou, Ioannis and Br{\"u}ningk, Sarah},
  booktitle = {2027 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)},
  year      = {2027}
}

If you use our pretrained models, you must also cite the RadImageNet paper, as the data terms require:

  • Mei et al., RadImageNet: An Open Radiologic Deep Learning Research Dataset for Effective Transfer Learning, Radiology: Artificial Intelligence, 2022. Paper
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
Unconditional Image Generation
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Paper for rad-dit/jit-radimagenet