125 lines
4.5 KiB
Python
125 lines
4.5 KiB
Python
"""Тесты cross-camera re-id логики.
|
|
|
|
Используем синтетические 512-мерные эмбеддинги (без InsightFace).
|
|
Проверяем, что find_topk_in_window:
|
|
1. Возвращает соседей в правильном порядке по cos-дистанции.
|
|
2. Фильтрует по camera_id (исключает ту же камеру).
|
|
3. Фильтрует по временному окну.
|
|
"""
|
|
from datetime import datetime, timedelta
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from database import (
|
|
save_embedding_with_meta,
|
|
find_topk_in_window,
|
|
find_nearest_patient,
|
|
attach_track_to_patient,
|
|
delete_patient_embeddings,
|
|
)
|
|
|
|
|
|
def normed(vec: np.ndarray) -> np.ndarray:
|
|
return (vec / np.linalg.norm(vec)).astype(np.float32)
|
|
|
|
|
|
def make_embedding(seed: int) -> np.ndarray:
|
|
rng = np.random.default_rng(seed)
|
|
return normed(rng.standard_normal(512))
|
|
|
|
|
|
def test_topk_in_window_basic(seed_camera_and_track):
|
|
"""Из 3 эмбеддингов на 3 разных камерах находим 2 ближайших к query (исключая саму камеру query)."""
|
|
cam_a, track_a = seed_camera_and_track("A")
|
|
cam_b, track_b = seed_camera_and_track("B")
|
|
cam_c, track_c = seed_camera_and_track("C")
|
|
|
|
base = make_embedding(seed=42)
|
|
|
|
# Соседи: tweak base слегка для cam_b, сильнее для cam_c.
|
|
near = normed(base + 0.05 * make_embedding(seed=43))
|
|
far = normed(base + 0.5 * make_embedding(seed=44))
|
|
|
|
save_embedding_with_meta(base, track_a, cam_a, quality=0.9, captured_at=datetime.utcnow())
|
|
save_embedding_with_meta(near, track_b, cam_b, quality=0.9, captured_at=datetime.utcnow())
|
|
save_embedding_with_meta(far, track_c, cam_c, quality=0.9, captured_at=datetime.utcnow())
|
|
|
|
# Запрос с cam_a — должен вернуть cam_b раньше cam_c, cam_a исключаем.
|
|
results = find_topk_in_window(
|
|
embedding=base.tolist(),
|
|
camera_id=cam_a,
|
|
window_minutes=5,
|
|
k=5,
|
|
exclude_same_camera=True,
|
|
)
|
|
|
|
cam_ids_in_results = [r["camera_id"] for r in results]
|
|
assert cam_a not in cam_ids_in_results
|
|
assert cam_b in cam_ids_in_results
|
|
assert cam_c in cam_ids_in_results
|
|
# Порядок: ближе → дальше
|
|
assert results[0]["camera_id"] == cam_b
|
|
assert results[0]["distance"] < results[-1]["distance"]
|
|
|
|
|
|
def test_topk_filters_by_window(seed_camera_and_track):
|
|
"""Старый эмбеддинг (вне окна) не должен попадать в результат."""
|
|
cam_a, track_a = seed_camera_and_track("A")
|
|
cam_b, track_b = seed_camera_and_track("B")
|
|
|
|
base = make_embedding(seed=7)
|
|
|
|
save_embedding_with_meta(
|
|
base, track_b, cam_b, quality=0.9,
|
|
captured_at=datetime.utcnow() - timedelta(hours=1), # вне окна 5 мин
|
|
)
|
|
|
|
results = find_topk_in_window(
|
|
embedding=base.tolist(),
|
|
camera_id=cam_a,
|
|
window_minutes=5,
|
|
k=5,
|
|
)
|
|
|
|
cam_ids = [r["camera_id"] for r in results]
|
|
assert cam_b not in cam_ids
|
|
|
|
|
|
def test_find_nearest_patient_only_consented(db_conn, seed_camera_and_track):
|
|
"""find_nearest_patient ищет только среди эмбеддингов с patient_id IS NOT NULL."""
|
|
cam_a, track_a = seed_camera_and_track("A")
|
|
base = make_embedding(seed=100)
|
|
|
|
# Сохраняем эмбеддинг без patient_id.
|
|
save_embedding_with_meta(base, track_a, cam_a, quality=0.9)
|
|
|
|
# Ищем — никого не должно найти.
|
|
assert find_nearest_patient(base.tolist(), threshold=0.5) is None
|
|
|
|
# Создаём пациента и привязываем трек.
|
|
import uuid
|
|
patient_id = str(uuid.uuid4())
|
|
with db_conn.cursor() as cur:
|
|
cur.execute(
|
|
"INSERT INTO patients (id, full_name, updated_at) VALUES (%s, %s, NOW())",
|
|
(patient_id, "Тестовый Пациент"),
|
|
)
|
|
db_conn.commit()
|
|
|
|
affected = attach_track_to_patient(track_a, patient_id)
|
|
assert affected == 1
|
|
|
|
# Теперь должны найти.
|
|
result = find_nearest_patient(base.tolist(), threshold=0.5)
|
|
assert result is not None
|
|
assert result["patient_id"] == patient_id
|
|
assert result["distance"] < 0.01 # тот же эмбеддинг
|
|
|
|
# Очистка.
|
|
deleted = delete_patient_embeddings(patient_id)
|
|
assert deleted == 1
|
|
with db_conn.cursor() as cur:
|
|
cur.execute("DELETE FROM patients WHERE id = %s", (patient_id,))
|
|
db_conn.commit()
|