"""DINOv2 特征提取器 — vit_small_patch14_reg4_dinov2 (带寄存器令牌)""" import torch import torch.nn.functional as F from torchvision import transforms from PIL import Image import numpy as np from module.config import DEVICE class DINOv2Extractor: """DINOv2 全局语义特征提取器 使用 vit_small_patch14_reg4_dinov2 模型（4个寄存器令牌）， 相比无寄存器版本，CLS token 特征质量更高，对背景噪声更鲁棒。 输出 384 维归一化特征向量。 支持提取 patch 特征用于可视化匹配。 """ TIMM_MODEL_NAME = 'vit_small_patch14_reg4_dinov2' PATCH_SIZE = 14 def __init__(self, device=None): self.device = device or DEVICE self.model = None self._use_timm = False # DINOv2 标准预处理：518x518，ImageNet 归一化 self.transform = transforms.Compose([ transforms.Resize(518, interpolation=transforms.InterpolationMode.BICUBIC), transforms.CenterCrop(518), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) try: import timm print(f"通过 timm 加载 DINOv2 (reg4): {self.TIMM_MODEL_NAME}") self.model = timm.create_model(self.TIMM_MODEL_NAME, pretrained=True) self._use_timm = True self.model = self.model.to(self.device).eval() print("DINOv2 reg4 模型加载完成") except Exception as e: print(f"timm 加载 DINOv2 失败: {e}") self.model = None @property def is_available(self): return self.model is not None @torch.no_grad() def _get_cls_feature(self, img_tensor): """从 forward_features 提取 CLS token，适配 timm 的 dict 返回格式""" features = self.model.forward_features(img_tensor) # timm 0.9+ 返回 dict if isinstance(features, dict): if 'x_norm_clstoken' in features: return features['x_norm_clstoken'] # 兜底：取第一个值的 CLS 位置 for v in features.values(): if isinstance(v, torch.Tensor) and v.dim() >= 2: return v[:, 0] # 旧版 timm 返回 tensor (B, N, C) if isinstance(features, torch.Tensor): return features[:, 0] raise ValueError(f"无法从 forward_features 提取 CLS token: {type(features)}") @torch.no_grad() def extract_feature(self, img_path): """提取单张图片的 DINOv2 CLS 全局特征向量（384维，已归一化）""" if self.model is None: return None img = Image.open(img_path).convert('RGB') tensor = self.transform(img).unsqueeze(0).to(self.device) cls_token = self._get_cls_feature(tensor) return F.normalize(cls_token, dim=-1).squeeze(0) @torch.no_grad() def compute_similarity(self, img_path1, img_path2): """计算两张图片的 DINOv2 CLS 特征余弦相似度（返回 0~1）""" feat1 = self.extract_feature(img_path1) feat2 = self.extract_feature(img_path2) if feat1 is None or feat2 is None: return None sim = F.cosine_similarity(feat1.unsqueeze(0), feat2.unsqueeze(0)) return float(sim.item()) @torch.no_grad() def extract_patch_features(self, img_path): """提取图片的 patch 级别特征（用于可视化匹配） Returns: dict with keys: - 'cls': CLS token (384维) - 'patches': patch token 特征 (n_patches, 384) - 'patch_grid': (n_h, n_w) patch 网格尺寸 - 'image_size': 原始图片尺寸 (w, h) """ if self.model is None: return None img_pil = Image.open(img_path).convert('RGB') orig_w, orig_h = img_pil.size tensor = self.transform(img_pil).unsqueeze(0).to(self.device) features = self.model.forward_features(tensor) if isinstance(features, dict): cls_token = features.get('x_norm_clstoken', None) patch_tokens = features.get('x_norm_patchtokens', None) # 如果没有 x_norm_patchtokens，从 x_norm_tokens 中提取 if patch_tokens is None and 'x_norm_tokens' in features: all_tokens = features['x_norm_tokens'] # 第0个是 CLS，后面是 patches + registers # reg4 模型有4个 register tokens 在末尾 patch_tokens = all_tokens[:, 1:-4] if all_tokens.shape[1] > 1 else None if cls_token is None and 'x_norm_tokens' in features: cls_token = features['x_norm_tokens'][:, 0] elif isinstance(features, torch.Tensor): # 旧版 timm: (B, N, C)，N = 1(cls) + n_patches + n_registers cls_token = features[:, 0] # reg4 有4个 register tokens 在末尾 patch_tokens = features[:, 1:-4] if features.shape[1] > 5 else features[:, 1:] else: return None if cls_token is None or patch_tokens is None: return None cls_feat = F.normalize(cls_token.squeeze(0), dim=-1) patch_feats = F.normalize(patch_tokens.squeeze(0), dim=-1) # DINOv2 预处理：Resize(518) → CenterCrop(518) # 需要记录 CenterCrop 的偏移量，用于可视化时将 patch 坐标映射回原图 input_size = 518 # CenterCrop 尺寸 n_patches = patch_feats.shape[0] n_side = int(n_patches ** 0.5) if n_side * n_side != n_patches: n_h = input_size // self.PATCH_SIZE n_w = n_h else: n_h = n_w = n_side # 计算 CenterCrop 在 Resize 后图像上的偏移 # Resize(518) 将短边缩放到518，长边按比例缩放 resize_w, resize_h = orig_w, orig_h if orig_w <= orig_h: resize_w = 518 resize_h = int(orig_h * 518 / orig_w) else: resize_h = 518 resize_w = int(orig_w * 518 / orig_h) crop_x = (resize_w - input_size) // 2 crop_y = (resize_h - input_size) // 2 return { 'cls': cls_feat.cpu().numpy(), 'patches': patch_feats.cpu().numpy(), 'patch_grid': (n_h, n_w), 'image_size': (orig_w, orig_h), 'crop_offset': (crop_x, crop_y), # CenterCrop 偏移 'resize_size': (resize_w, resize_h), # Resize 后尺寸 } def compute_patch_matches(self, img_path1, img_path2, top_k=50): """计算两张图片之间的 patch 级别匹配（用于可视化） Returns: dict with keys: - 'cls_similarity': CLS token 余弦相似度 - 'matches': list of (query_idx, hist_idx, similarity) 最优匹配 - 'query_grid': (n_h, n_w) - 'hist_grid': (n_h, n_w) - 'query_image_size': (w, h) - 'hist_image_size': (w, h) """ feat1 = self.extract_patch_features(img_path1) feat2 = self.extract_patch_features(img_path2) if feat1 is None or feat2 is None: return None # CLS 相似度 cls_sim = float(np.dot(feat1['cls'], feat2['cls'])) # Patch 间余弦相似度矩阵 sim_matrix = feat1['patches'] @ feat2['patches'].T # (n1, n2) # 双向最近邻匹配（类似 MASt3R 的互为最近邻） matches = [] q_best = sim_matrix.argmax(axis=1) # 每个 query patch 的最佳匹配 h_best = sim_matrix.argmax(axis=0) # 每个 hist patch 的最佳匹配 for qi, hi in enumerate(q_best): if h_best[hi] == qi: # 互为最近邻 sim = float(sim_matrix[qi, hi]) matches.append((qi, hi, sim)) # 按相似度排序，取 top_k matches.sort(key=lambda x: x[2], reverse=True) matches = matches[:top_k] return { 'cls_similarity': round(cls_sim, 4), 'matches': [(q, h, round(s, 4)) for q, h, s in matches], 'query_grid': feat1['patch_grid'], 'hist_grid': feat2['patch_grid'], 'query_image_size': feat1['image_size'], 'hist_image_size': feat2['image_size'], 'query_crop_offset': feat1.get('crop_offset', (0, 0)), 'hist_crop_offset': feat2.get('crop_offset', (0, 0)), 'query_resize_size': feat1.get('resize_size', (518, 518)), 'hist_resize_size': feat2.get('resize_size', (518, 518)), }