PokemonCardSearch/code/modules/feature_extractor.py

"""
DINOv2 特征提取模块
复用 dinov2_base_retrieval_392_PokemonCN04 模型，输出 768 维 L2 归一化特征向量。
预处理逻辑与训练时保持一致（Letterbox 392×392 + ImageNet 归一化）。
"""
import os
import numpy as np


class CardFeatureExtractor:
    """DINOv2 卡牌特征提取器"""

    def __init__(self, model_dir, device=None):
        import torch
        import torch.nn as nn
        from transformers import Dinov2Model
        from safetensors.torch import load_file

        if not os.path.isdir(model_dir):
            raise NotADirectoryError(f"模型目录不存在: {model_dir}")

        self.torch = torch
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))

        # 加载骨干网络
        self.backbone = Dinov2Model.from_pretrained(model_dir, local_files_only=True)
        self.feature_dim = self.backbone.config.hidden_size
        self.image_size = self.backbone.config.image_size  # 392

        # BN Neck（训练时用的，推理需加载其权重）
        self.bn_neck = nn.BatchNorm1d(self.feature_dim)
        st_path = os.path.join(model_dir, "model.safetensors")
        if os.path.exists(st_path):
            self._load_bn_weights(st_path)

        self.backbone.to(self.device).eval()
        self.bn_neck.to(self.device).eval()

        # ImageNet 归一化参数
        self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32)

    def _load_bn_weights(self, safetensors_path):
        from safetensors.torch import load_file
        state = load_file(safetensors_path, device="cpu")
        bn_dict = {k.replace("bn_neck.", ""): v for k, v in state.items() if k.startswith("bn_neck.")}
        if bn_dict:
            self.bn_neck.load_state_dict(bn_dict)

    def _letterbox(self, img_rgb):
        """Letterbox 等比缩放 + 黑边填充到 image_size×image_size"""
        import cv2
        shape = img_rgb.shape[:2]  # h, w
        target = self.image_size
        r = min(target / shape[0], target / shape[1])
        new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))  # w, h
        dw = (target - new_unpad[0]) / 2
        dh = (target - new_unpad[1]) / 2

        if (shape[1], shape[0]) != new_unpad:
            img_rgb = cv2.resize(img_rgb, new_unpad, interpolation=cv2.INTER_CUBIC)

        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        img_rgb = cv2.copyMakeBorder(img_rgb, top, bottom, left, right,
                                     cv2.BORDER_CONSTANT, value=(0, 0, 0))
        if img_rgb.shape[:2] != (target, target):
            img_rgb = cv2.resize(img_rgb, (target, target), interpolation=cv2.INTER_CUBIC)

        img = img_rgb.astype(np.float32) / 255.0
        img = (img - self.mean) / self.std
        return img.transpose(2, 0, 1)  # CHW

    @property
    def _to_tensor(self):
        import torch.nn.functional as F
        return self.torch, F

    def extract(self, imgs, normalize=True):
        """
        批量提取特征。
        Args:
            imgs: RGB numpy 数组列表 (每个 H,W,3 uint8)
            normalize: 是否 L2 归一化
        Returns:
            np.ndarray (N, 768)，失败的行填充 NaN
        """
        import torch.nn.functional as F
        torch = self.torch

        tensors, valid_idx = [], []
        for i, img in enumerate(imgs):
            if img is None or img.ndim != 3:
                continue
            try:
                tensors.append(torch.from_numpy(self._letterbox(img)).float())
                valid_idx.append(i)
            except Exception:
                continue

        out = np.full((len(imgs), self.feature_dim), np.nan, dtype=np.float32)
        if not tensors:
            return out

        batch = torch.stack(tensors).to(self.device)
        with torch.no_grad():
            outputs = self.backbone(batch)
            cls = outputs.last_hidden_state[:, 0, :]
            feat = self.bn_neck(cls)
            if normalize:
                feat = F.normalize(feat, p=2, dim=1)
        feats_np = feat.cpu().numpy()
        for idx, vec in zip(valid_idx, feats_np):
            out[idx] = vec
        return out