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