"""Метрики для оценки качества перекрёстной геолокализации (CVGL). R@K (один positive на запрос): R@K = (1/N) Σ_i 𝟙[rank(i) ≤ K] где rank(i) — позиция (1-indexed) истинного соответствия. R@K (несколько positive, конвенция «hit-if-any», Sample4Geo): достаточно, чтобы ХОТЯ БЫ ОДИН positive попал в top-K. AP для одного запроса — доля позитивных предсказаний среди top-K; mAP — среднее AP по всем запросам. q2g (query-to-gallery): БПЛА → спутник (основная метрика); g2q (gallery-to-query): спутник → БПЛА (диагностика симметричности). """ from __future__ import annotations import numpy as np import torch @torch.no_grad() def compute_retrieval_metrics( query_embeddings: torch.Tensor, gallery_embeddings: torch.Tensor, ks: tuple[int, ...] = (1, 5, 10), positives: list[list[int]] | None = None, compute_map: bool = True, return_ranks: bool = False, ) -> dict[str, float]: """Вычислить метрики поиска q2g. Args: query_embeddings: [N, D] нормализованные слитые drone-эмбеддинги. gallery_embeddings: [M, D] нормализованные слитые satellite-эмбеддинги. ks: Значения K для R@K. positives: Опционально — список positive-индексов галереи для каждого запроса (для конвенции «hit-if-any» при нескольких правильных ответах). Если None — предполагается диагональное соответствие query[i] ↔ gallery[i] (один positive на запрос). compute_map: Считать ли mAP. return_ranks: Если True — добавить в результат ключ "_best_ranks": np.ndarray [N] с 0-indexed рангом первого попавшего positive по каждому запросу (M, если ни один не попал). Из него выводится hit@k = (ranks < k) для любого k — нужно для bootstrap CI и paired-теста (протокол §8.3 / §9.2). Returns: dict: recall@1, recall@5, recall@10, recall@1%, mean_rank, median_rank, mAP (+ "_best_ranks", если return_ranks=True). """ N = query_embeddings.shape[0] M = gallery_embeddings.shape[0] # Матрица косинусного сходства (эмбеддинги уже L2-нормализованы) sim = query_embeddings @ gallery_embeddings.T # [N, M] sorted_indices = sim.argsort(dim=1, descending=True) # [N, M] # positive-множество для каждого запроса if positives is None: # Диагональное соответствие: для запроса i правильный ответ — gallery i positives = [[i] for i in range(N)] # Ранг ПЕРВОГО найденного positive (hit-if-any) для каждого запроса best_ranks = torch.full((N,), M, dtype=torch.long) # AP по каждому запросу (для mAP) aps = torch.zeros(N, dtype=torch.float) sorted_np = sorted_indices.cpu().numpy() for i in range(N): pos_set = set(positives[i]) if not pos_set: continue order = sorted_np[i] # индексы галереи по убыванию сходства # Позиции всех positive в ранжированном списке (0-indexed) hit_positions = [r for r, idx in enumerate(order) if idx in pos_set] if not hit_positions: continue # Ранг первого попадания (1-indexed) — для R@K по «hit-if-any» best_ranks[i] = hit_positions[0] # AP запроса: среднее precision@(каждое попадание) # precision на k-м попадании = (номер попадания) / (позиция+1) if compute_map: ap = 0.0 for hit_num, pos in enumerate(hit_positions, start=1): precision_at_hit = hit_num / (pos + 1) ap += precision_at_hit aps[i] = ap / len(pos_set) results: dict[str, float] = {} # R@K: доля запросов, где первый positive попал в top-K for k in ks: results[f"recall@{k}"] = (best_ranks < k).float().mean().item() # R@1% k_1pct = max(1, M // 100) results["recall@1%"] = (best_ranks < k_1pct).float().mean().item() # Mean / Median rank (1-indexed для читаемости) results["mean_rank"] = (best_ranks.float() + 1).mean().item() results["median_rank"] = (best_ranks.float() + 1).median().item() # mAP if compute_map: results["mAP"] = aps.mean().item() if return_ranks: results["_best_ranks"] = best_ranks.cpu().numpy() return results @torch.no_grad() def compute_bidirectional_metrics( query_embeddings: torch.Tensor, gallery_embeddings: torch.Tensor, ks: tuple[int, ...] = (1, 5, 10), positives_q2g: list[list[int]] | None = None, positives_g2q: list[list[int]] | None = None, return_ranks: bool = False, ) -> dict[str, float]: """Метрики в обе стороны. q2g (query-to-gallery): БПЛА → спутник — основная метрика. g2q (gallery-to-query): спутник → БПЛА — диагностика симметричности признакового пространства. Args: return_ranks: Если True — добавить ключи "q2g_best_ranks" и "g2q_best_ranks" (np.ndarray, 0-indexed ранги первого попадания) для последующего bootstrap CI / paired-теста. Числовые метрики остаются float, per-query массивы лежат под отдельными ключами. """ q2g = compute_retrieval_metrics( query_embeddings, gallery_embeddings, ks=ks, positives=positives_q2g, return_ranks=return_ranks, ) g2q = compute_retrieval_metrics( gallery_embeddings, query_embeddings, ks=ks, positives=positives_g2q, return_ranks=return_ranks, ) q2g_ranks = q2g.pop("_best_ranks", None) g2q_ranks = g2q.pop("_best_ranks", None) results = {} for key, val in q2g.items(): results[f"q2g_{key}"] = val for key, val in g2q.items(): results[f"g2q_{key}"] = val # Средний R@K по обоим направлениям for k in ks: results[f"avg_recall@{k}"] = (q2g[f"recall@{k}"] + g2q[f"recall@{k}"]) / 2 if return_ranks: results["q2g_best_ranks"] = q2g_ranks results["g2q_best_ranks"] = g2q_ranks return results def compute_meter_distance( predicted_indices: np.ndarray, query_coords: np.ndarray, gallery_coords: np.ndarray, ) -> dict[str, float]: """Расстояние в метрах между предсказанной и истинной позицией. Берёт координаты top-1 предсказанного спутникового тайла и сравнивает с истинными координатами запроса. """ predicted_coords = gallery_coords[predicted_indices] diffs = predicted_coords - query_coords distances = np.linalg.norm(diffs, axis=1) results = { "mean_distance": float(distances.mean()), "median_distance": float(np.median(distances)), } for threshold in (5, 10, 25, 50, 100): results[f"recall@{threshold}m"] = float((distances < threshold).mean()) return results def format_metrics(metrics: dict[str, float], prefix: str = "") -> str: """Форматировать метрики для логирования.""" lines = [] for key, val in metrics.items(): if "recall" in key or "mAP" in key or "AP" in key or "acc" in key: lines.append(f"{prefix}{key}: {val:.4f} ({val*100:.1f}%)") elif "rank" in key: lines.append(f"{prefix}{key}: {val:.1f}") elif "distance" in key: lines.append(f"{prefix}{key}: {val:.1f}m") else: lines.append(f"{prefix}{key}: {val:.4f}") return "\n".join(lines)