| |
| |
| |
| import os |
| current_dir_path = os.path.dirname(__file__) |
| import torch |
| from torch import nn |
|
|
| class Dinov2Backbone(nn.Module): |
| def __init__(self, name='dinov2_vitb14', pretrained=False, *args, **kwargs): |
| super().__init__() |
| self.name = name |
| self.encoder = torch.hub.load(current_dir_path+'/../dinov2', self.name, pretrained=pretrained, source='local') |
| self.patch_size = self.encoder.patch_size |
| self.embed_dim = self.encoder.embed_dim |
|
|
| def forward(self, x): |
| """ |
| Encode a RGB image using a ViT-backbone |
| Args: |
| - x: torch.Tensor of shape [bs,3,w,h] |
| Return: |
| - y: torch.Tensor of shape [bs,k,d] - image in patchified mode |
| """ |
| assert len(x.shape) == 4 |
| y = self.encoder.get_intermediate_layers(x)[0] |
| return y |
|
|
|
|
