|
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", |
|
"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): |
|
|
|
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) |
|
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 |
|
|
