Spaces:
Sleeping
Sleeping
| """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 | |
| 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) | |