Views
No views yet
| Model | Feature Dim | Parameters | Checkpoint |
|---|---|---|---|
| DinoBloom-S | 384 | 22M | pytorch_model_s.bin |
| DinoBloom-B | 768 | 86M | pytorch_model_b.bin |
| DinoBloom-L | 1024 | 304M | pytorch_model_l.bin |
| DinoBloom-G | 1536 | 1136M | pytorch_model_g.bin |
1from huggingface_hub import hf_hub_download
2import torch
3import torch.nn as nn
4
5device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6
7# Choose variant: "s", "b", "l", or "g"
8variant = "b"
9
10# Configuration
11variant_config = {
12 "s": ("dinov2_vits14", 384),
13 "b": ("dinov2_vitb14", 768),
14 "l": ("dinov2_vitl14", 1024),
15 "g": ("dinov2_vitg14", 1536),
16}
17
18dinov2_model, embed_dim = variant_config[variant]
19
20# Load base DINOv2 model
21model = torch.hub.load("facebookresearch/dinov2", dinov2_model)
22
23# Download DinoBloom weights
24ckpt_path = hf_hub_download(
25 repo_id="MarrLab/DinoBloom",
26 filename=f"pytorch_model_{variant}.bin"
27)
28ckpt = torch.load(ckpt_path, map_location="cpu")
29
30num_tokens = int(1 + (224 / 14) ** 2)
31model.pos_embed = nn.Parameter(torch.zeros(1, num_tokens, embed_dim))
32model.load_state_dict(ckpt, strict=True)
33model.to(device)
34model.eval()
35
36# Get transforms
37from torchvision import transforms
38transform = transforms.Compose([
39 transforms.Resize((224,224)),
40 transforms.ToTensor(),
41 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
42])
43
44# Apply to image
45from PIL import Image
46img = Image.open("path/to/cell_image")
47img_tensor = transform(img).unsqueeze(0).to(device)
48
49# Get features
50with torch.no_grad():
51 features = model(img_tensor)
52
53print(f"Features shape: {features.shape}") # [1, 768] for DinoBloom-Bpip install torch torchvision huggingface_hub1@inproceedings{koch2024dinobloom,
2 title={DinoBloom: a foundation model for generalizable cell embeddings in hematology},
3 author={Koch, Valentin and Wagner, Sophia J and Kazeminia, Salome and Sancar, Ece and Hehr, Matthias and Schnabel, Julia A and Peng, Tingying and Marr, Carsten},
4 booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
5 pages={520--530},
6 year={2024},
7 organization={Springer}
8}