| | |
| | |
| | |
| | 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 |
| |
|
| |
|
