module/scene_matcher.py

"""场景匹配器 — DINOv2-only 模式（无 MASt3R）

复杂场景判定改为纯 DINOv2 余弦相似度，不调用 MASt3R 3D 几何匹配。
阈值由 config.DINOV2_ONLY_SAME_THRESHOLD 控制。
支持 DINOv2 patch 级别匹配可视化。
"""
import torch

from module.config import DEVICE, DINOV2_ONLY_SAME_THRESHOLD


class SceneMatcher:
    """DINOv2-only 场景匹配器"""

    def __init__(self, model=None, dinov2_extractor=None, device=None):
        self.model = model  # None（MASt3R 已禁用）
        self.dinov2_extractor = dinov2_extractor
        self.device = device or DEVICE

    def compare(self, img_path1, img_path2, dinov2_sim_override=None):
        """比对两张图片是否为同一场景（DINOv2-only）

        Returns:
            dict: 与原始 SceneMatcher.compare() 兼容的结果格式
        """
        # 获取 DINOv2 相似度
        if dinov2_sim_override is not None:
            dinov2_sim = float(dinov2_sim_override)
        elif self.dinov2_extractor is not None:
            try:
                dinov2_sim = self.dinov2_extractor.compute_similarity(img_path1, img_path2)
                if dinov2_sim is None:
                    dinov2_sim = 0.0
            except Exception:
                dinov2_sim = 0.0
        else:
            dinov2_sim = 0.0

        # DINOv2-only 判定
        is_same = dinov2_sim >= DINOV2_ONLY_SAME_THRESHOLD

        # 如果判定为同一场景，计算 patch 匹配信息（用于可视化）
        patch_match_info = None
        if is_same and self.dinov2_extractor is not None:
            try:
                patch_match_info = self.dinov2_extractor.compute_patch_matches(img_path1, img_path2, top_k=50)
            except Exception as e:
                print(f"  [可视化] patch 匹配计算失败: {e}")

        return {
            'is_same_scene': bool(is_same),
            'mast3r_is_same': None,
            'dinov2_is_same': bool(is_same),
            'similarity_score': round(float(dinov2_sim), 4),
            'match_count': len(patch_match_info['matches']) if patch_match_info else 0,
            'raw_match_count': len(patch_match_info['matches']) if patch_match_info else 0,
            'inlier_ratio': 0.0,
            'avg_confidence': round(float(dinov2_sim), 4),
            'dinov2_similarity': round(float(dinov2_sim), 4),
            'gamma_info': [],
            'patch_match_info': patch_match_info,
        }