Views
No views yet
vit_base_patch14_dinov2 on a subset of the ImageNet-1k dataset.
The classification head has been trained for ImageNet-1k class classification (1000 classes).1import torch
2import timm
3from PIL import Image
4from huggingface_hub import hf_hub_download
5
6# 1. Create Base Model
7# Using the base model ID compatible with timm
8model_id = "vit_base_patch14_dinov2"
9model = timm.create_model(model_id, pretrained=True) # Load pretrained backbone
10
11# 2. Modify Head for 1000 classes
12# DINOv2 usually has 768 dim for Base.
13model.head = torch.nn.Linear(768, 1000, bias=True)
14
15# 3. Load Fine-tuned Weights
16# Download the checkpoint from this repo
17checkpoint_path = hf_hub_download(repo_id="SasikaA073/vit_base_patch14_dinov2_sp_ft_in1k", filename="vit_base_patch14_dinov2_sp_ft_in1k.pth")
18state_dict = torch.load(checkpoint_path, map_location='cpu')
19
20# Load state dict
21# Note: The model was saved as state_dict only
22model.load_state_dict(state_dict)
23model.eval()
24
25# 4. Inference
26data_config = timm.data.resolve_data_config(model.pretrained_cfg)
27transforms = timm.data.create_transform(**data_config, is_training=False)
28
29image = Image.open("your_image.jpg").convert('RGB')
30input_tensor = transforms(image).unsqueeze(0)
31
32with torch.no_grad():
33 output = model(input_tensor)
34 probabilities = torch.nn.functional.softmax(output[0], dim=0)
35
36print(f"Top class index: {probabilities.argmax().item()}")1@article{oquab2023dinov2,
2 title={DINOv2: Learning Robust Visual Features without Supervision},
3 author={Oquab, Maxime and Darcet, Timothée and Moutakanni, Theo and Vo, Huy V. and Szafraniec, Marc and Khalidov, Vasil and Fernandez, Pierre and Haziza, Daniel and Massa, Francisco and El-Nouby, Alaaeldin and Howes, Russell and Huang, Po-Yao and Xu, Hu and Sharma, Vasu and Li, Shang-Wen and Galuba, Wojciech and Rabbat, Mike and Assran, Mido and Ballas, Nicolas and Synnaeve, Gabriel and Misra, Ishan and Jegou, Herve and Mairal, Julien and Labatut, Patrick and Joulin, Armand and Bojanowski, Piotr},
4 journal={arXiv:2304.07193},
5 year={2023}
6}