Spaces:
Sleeping
Sleeping
File size: 3,777 Bytes
2267636 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """torch.hub DINOv3/DINOv2 feature extractor (fallback for the HF extractor).
Used by `core.build_dino_extractor` only when a backbone name is not in the
HuggingFace repo map. Extracts intermediate transformer layers as spatial maps.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class DINOv3FeatureExtractor(nn.Module):
"""Extract intermediate ViT layers from a DINOv3/DINOv2 model loaded via torch.hub.
``take_indices`` selects the layers to return (a list of block indices, matching
the HuggingFace extractor); the backbone is frozen.
"""
def __init__(self, model_name="dinov3_vitb16", take_indices=(2, 5, 8, 11)):
super().__init__()
self.model_name = model_name
self.take_indices = list(take_indices)
if model_name.startswith('facebook/'):
hf_to_hub = {
'facebook/dinov2-base': 'dinov2_vitb14',
'facebook/dinov2-small': 'dinov2_vits14',
'facebook/dinov2-large': 'dinov2_vitl14',
'facebook/dinov2-giant': 'dinov2_vitg14',
'facebook/dinov3-vitb16-pretrain-lvd1689m': 'dinov3_vitb16',
'facebook/dinov3-vits16-pretrain-lvd1689m': 'dinov3_vits16',
'facebook/dinov3-vitl16-pretrain-lvd1689m': 'dinov3_vitl16',
}
model_name = hf_to_hub.get(model_name, 'dinov2_vitb14')
self.model_name = model_name
if 'dinov2' in model_name:
self.dino = torch.hub.load('facebookresearch/dinov2', model_name)
self.patch_size = 14
elif 'dinov3' in model_name:
try:
self.dino = torch.hub.load('facebookresearch/dinov3', model_name)
except Exception as e:
# torch.hub gated/unavailable -> fall back to HuggingFace weights.
print(f"[DINO] torch.hub failed ({e}); loading from HuggingFace")
from transformers import AutoModel
hf_map = {
'dinov3_vits16': 'facebook/dinov3-vits16-pretrain-lvd1689m',
'dinov3_vitb16': 'facebook/dinov3-vitb16-pretrain-lvd1689m',
'dinov3_vitl16': 'facebook/dinov3-vitl16-pretrain-lvd1689m',
}
self.dino = AutoModel.from_pretrained(
hf_map.get(model_name, 'facebook/dinov3-vitb16-pretrain-lvd1689m'),
trust_remote_code=True)
self.patch_size = 16
else:
raise ValueError(f"Unsupported model name: {model_name}. Use dinov2_* or dinov3_*")
self.dino.eval()
for p in self.dino.parameters():
p.requires_grad = False
@torch.no_grad()
def forward(self, images):
"""images: [B, 3, H, W] in [0, 1] -> list of feature maps [B, C, H//p, W//p]."""
h, w = images.shape[-2:]
if h % self.patch_size != 0 or w % self.patch_size != 0:
new_h = ((h + self.patch_size - 1) // self.patch_size) * self.patch_size
new_w = ((w + self.patch_size - 1) // self.patch_size) * self.patch_size
images = F.interpolate(images, size=(new_h, new_w), mode='bilinear', align_corners=False)
h, w = new_h, new_w
features = self.dino.get_intermediate_layers(images, self.take_indices, return_class_token=False)
patch_h, patch_w = h // self.patch_size, w // self.patch_size
outs = []
for feat in features:
B, N, C = feat.shape
outs.append(feat.permute(0, 2, 1).reshape(B, C, patch_h, patch_w).contiguous())
return outs
def create_dino_extractor(model_name="dinov3_vitb16", take_indices=(2, 5, 8, 11)):
return DINOv3FeatureExtractor(model_name=model_name, take_indices=take_indices)
|