| 7 | from transformers import AutoImageProcessor, AutoModel |
| 8 | import math |
| 9 | class DinoFeatureModule(nn.Module): |
| 10 | def __init__(self, dino_model='dinov2_giant', pretrained_path=r'pretrained_model/facebookdinov2_giant',img_size = 128): |
| 11 | super(DinoFeatureModule, self).__init__() |
| 12 | |
| 13 | self.dino = AutoModel.from_pretrained( |
| 14 | pretrained_path, |
| 15 | local_files_only=False, |
| 16 | torch_dtype=torch.float16 |
| 17 | ) |
| 18 | |
| 19 | |
| 20 | self.dino.eval() |
| 21 | for param in self.dino.parameters(): |
| 22 | param.requires_grad = False |
| 23 | |
| 24 | |
| 25 | frozen = all(not p.requires_grad for p in self.dino.parameters()) |
| 26 | assert frozen, "DINOv2 model parameters are not completely frozen!" |
| 27 | |
| 28 | |
| 29 | self.shallow_dim = 1536 |
| 30 | self.mid_dim = 1536 |
| 31 | self.deep_dim = 1536 |
| 32 | |
| 33 | def get_dino_features(self, x): |
| 34 | with torch.no_grad(): |
| 35 | outputs = self.dino(x, output_hidden_states=True) |
| 36 | hidden_states = outputs.hidden_states |
| 37 | |
| 38 | _, _, H, W = x.shape |
| 39 | aspect_ratio = W / H |
| 40 | |
| 41 | shallow_feat1 = hidden_states[7] |
| 42 | shallow_feat2 = hidden_states[15] |
| 43 | mid_feat1 = hidden_states[20] |
| 44 | mid_feat2 = hidden_states[22] |
| 45 | deep_feat1 = hidden_states[33] |
| 46 | deep_feat2 = hidden_states[39] |
| 47 | |
| 48 | def reshape_features(feat): |
| 49 | feat = feat[:, 1:, :] |
| 50 | B, N, C = feat.shape |
| 51 | |
| 52 | h = int(math.sqrt(N / aspect_ratio)) |
| 53 | w = int(N / h) |
| 54 | |
| 55 | |
| 56 | if(aspect_ratio > 1): |
| 57 | if h * w > N: |
| 58 | h -= 1 |
| 59 | w = N // h |
| 60 | if h * w < N: |
| 61 | h += 1 |
| 62 | w = N // h |
| 63 | else: |
| 64 | if h * w > N: |
| 65 | w -= 1 |
| 66 | h = N // w |