205 lines
8.6 KiB
Python
205 lines
8.6 KiB
Python
"""Метрики для оценки качества перекрёстной геолокализации (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)
|