You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
124 lines
4.5 KiB
124 lines
4.5 KiB
"""Тесты 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()
|
|
|