Spaces:
Running
Running
| from .base import BaseModel | |
| from .schema import DINOConfiguration | |
| import logging | |
| import torch | |
| import torch.nn as nn | |
| import sys | |
| import re | |
| import os | |
| from .dinov2.eval.depth.ops.wrappers import resize | |
| from .dinov2.hub.backbones import dinov2_vitb14_reg | |
| module_dir = os.path.dirname(os.path.abspath(__file__)) | |
| sys.path.append(module_dir) | |
| logger = logging.getLogger(__name__) | |
| class FeatureExtractor(BaseModel): | |
| mean = [0.485, 0.456, 0.406] | |
| std = [0.229, 0.224, 0.225] | |
| def build_encoder(self, conf: DINOConfiguration): | |
| BACKBONE_SIZE = "small" | |
| backbone_archs = { | |
| "small": "vits14", | |
| "base": "vitb14", # this one | |
| "large": "vitl14", | |
| "giant": "vitg14", | |
| } | |
| backbone_arch = backbone_archs[BACKBONE_SIZE] | |
| self.crop_size = int(re.search(r"\d+", backbone_arch).group()) | |
| backbone_name = f"dinov2_{backbone_arch}" | |
| self.backbone_model = dinov2_vitb14_reg( | |
| pretrained=conf.pretrained, drop_path_rate=0.1) | |
| if conf.frozen: | |
| for param in self.backbone_model.patch_embed.parameters(): | |
| param.requires_grad = False | |
| for i in range(0, 10): | |
| for param in self.backbone_model.blocks[i].parameters(): | |
| param.requires_grad = False | |
| self.backbone_model.blocks[i].drop_path1 = nn.Identity() | |
| self.backbone_model.blocks[i].drop_path2 = nn.Identity() | |
| self.feat_projection = torch.nn.Conv2d( | |
| 768, conf.output_dim, kernel_size=1) | |
| return self.backbone_model | |
| def _init(self, conf: DINOConfiguration): | |
| # Preprocessing | |
| self.register_buffer("mean_", torch.tensor( | |
| self.mean), persistent=False) | |
| self.register_buffer("std_", torch.tensor(self.std), persistent=False) | |
| self.build_encoder(conf) | |
| def _forward(self, data): | |
| _, _, h, w = data["image"].shape | |
| h_num_patches = h // self.crop_size | |
| w_num_patches = w // self.crop_size | |
| h_dino = h_num_patches * self.crop_size | |
| w_dino = w_num_patches * self.crop_size | |
| image = resize(data["image"], (h_dino, w_dino)) | |
| image = (image - self.mean_[:, None, None]) / self.std_[:, None, None] | |
| output = self.backbone_model.forward_features( | |
| image)['x_norm_patchtokens'] | |
| output = output.reshape(-1, h_num_patches, | |
| w_num_patches, output.shape[-1]) | |
| output = output.permute(0, 3, 1, 2) # channel first | |
| output = self.feat_projection(output) | |
| camera = data['camera'].to(data["image"].device, non_blocking=True) | |
| camera = camera.scale(output.shape[-1] / data["image"].shape[-1]) | |
| return output, camera | |