digital_reception_ra/apps/face-service/main.py

"""face-service v2 для проекта «Цифровая рецепция».

Расширение time-tracker face-service: добавлены эндпоинты cross-camera re-id,
сохранения эмбеддингов с метаданными трека/камеры, узнавания пациента
(только среди эмбеддингов с согласием), удаления при отзыве согласия.

Эндпоинты:
  GET    /health                            — статус + флаг loaded
  POST   /embed                             — только эмбеддинг кадра (без БД)
  POST   /track-embeddings                  — сохранить эмбеддинг с привязкой к треку
  POST   /reid/search                       — cross-camera re-id (top-K в окне T мин)
  POST   /recognize                         — узнать пациента (patient_id) по кадру
  POST   /enroll                            — привязать эмбеддинги трека к пациенту
  DELETE /patient/{patient_id}/embeddings   — удалить все эмбеддинги пациента
  GET    /patient/{patient_id}/count        — кол-во эмбеддингов у пациента
"""
import os
import logging
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Optional

from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from database import (
    save_embedding_with_meta,
    attach_track_to_patient,
    find_topk_in_window,
    find_nearest_patient,
    delete_patient_embeddings,
    count_patient_embeddings,
)
from face_engine import detect_best_face, load_model, is_model_loaded

load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Порог узнавания пациента по лицу (ТЗ §4.3 — после первой партии данных тюним).
RECOGNITION_THRESHOLD = float(os.getenv("RECOGNITION_THRESHOLD", "0.5"))
# Порог склейки треков cross-camera (строже — иначе ложные склейки).
REID_THRESHOLD = float(os.getenv("REID_THRESHOLD", "0.35"))
# Окно для cross-camera re-id (минуты).
DEFAULT_REID_WINDOW_MIN = int(os.getenv("REID_WINDOW_MINUTES", "5"))


@asynccontextmanager
async def lifespan(app: FastAPI):
    if os.getenv("SKIP_MODEL_LOAD", "false").lower() == "true":
        logger.warning("SKIP_MODEL_LOAD=true — модель не загружается, frame-эндпоинты вернут 503")
        yield
        return
    try:
        load_model()
    except Exception as e:  # noqa: BLE001
        logger.error(f"Не удалось загрузить InsightFace: {e}. Frame-эндпоинты не будут работать.")
    yield


app = FastAPI(title="reception/face-service", version="0.2.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ---------- Схемы ----------

class EmbedRequest(BaseModel):
    frame: str = Field(..., description="base64-encoded JPEG/PNG")


class FaceBbox(BaseModel):
    box: list[int]  # [x1, y1, x2, y2] в пикселях
    imgW: int
    imgH: int


class EmbedResponse(BaseModel):
    embedding: list[float]
    quality: float
    bbox: FaceBbox | None = None


class TrackEmbeddingRequest(BaseModel):
    frame: str
    track_id: str
    camera_id: str
    captured_at: datetime | None = None
    patient_id: str | None = None


class TrackEmbeddingResponse(BaseModel):
    id: str
    quality: float
    bbox: FaceBbox | None = None


class ReidSearchRequest(BaseModel):
    frame: str | None = None
    embedding: list[float] | None = None
    camera_id: str
    window_minutes: int | None = None
    k: int = 5
    exclude_same_camera: bool = True


class ReidMatch(BaseModel):
    track_id: str
    camera_id: str
    captured_at: str
    distance: float


class ReidSearchResponse(BaseModel):
    matches: list[ReidMatch]
    threshold: float


class RecognizeRequest(BaseModel):
    frame: str


class RecognizeResponse(BaseModel):
    patient_id: str
    confidence: float
    distance: float


class EnrollRequest(BaseModel):
    track_id: str
    patient_id: str


class EnrollResponse(BaseModel):
    ok: bool
    embeddings_attached: int


class TrackEmbeddingRawRequest(BaseModel):
    """Сохранить готовый эмбеддинг без детекции лица — для fixtures-runner."""
    embedding: list[float]
    track_id: str
    camera_id: str
    captured_at: datetime | None = None
    patient_id: str | None = None
    quality: float = 0.9


class RecognizeEmbeddingRequest(BaseModel):
    """Узнать пациента по готовому эмбеддингу — для fixtures-runner."""
    embedding: list[float]


# ---------- Эндпоинты ----------

@app.get("/health")
def health():
    return {
        "status": "ok",
        "model_loaded": is_model_loaded(),
        "recognition_threshold": RECOGNITION_THRESHOLD,
        "reid_threshold": REID_THRESHOLD,
    }


@app.post("/embed", response_model=Optional[EmbedResponse])
def embed(req: EmbedRequest):
    """Возвращает 512-d эмбеддинг лучшего лица на кадре, без записи в БД."""
    embedding, quality, bbox = detect_best_face(req.frame)
    if embedding is None:
        return None
    return EmbedResponse(embedding=embedding.tolist(), quality=quality, bbox=bbox)


@app.post("/track-embeddings", response_model=Optional[TrackEmbeddingResponse])
def store_track_embedding(req: TrackEmbeddingRequest):
    """Сохраняет эмбеддинг с привязкой к треку (используется video-ingest/fixtures)."""
    embedding, quality, bbox = detect_best_face(req.frame)
    if embedding is None:
        return None

    record_id = save_embedding_with_meta(
        embedding=embedding,
        track_id=req.track_id,
        camera_id=req.camera_id,
        quality=quality,
        captured_at=req.captured_at,
        patient_id=req.patient_id,
    )
    return TrackEmbeddingResponse(id=record_id, quality=quality, bbox=bbox)


@app.post("/reid/search", response_model=ReidSearchResponse)
def reid_search(req: ReidSearchRequest):
    """Cross-camera re-id: ищет top-K ближайших эмбеддингов с других камер в окне T мин."""
    if req.embedding is None and req.frame is None:
        raise HTTPException(400, "Нужен либо frame, либо embedding")

    if req.embedding is not None:
        embedding = req.embedding
    else:
        embedding, _quality, _bbox = detect_best_face(req.frame)
        if embedding is None:
            return ReidSearchResponse(matches=[], threshold=REID_THRESHOLD)
        embedding = embedding.tolist()

    window = req.window_minutes or DEFAULT_REID_WINDOW_MIN
    matches = find_topk_in_window(
        embedding=embedding,
        camera_id=req.camera_id,
        window_minutes=window,
        k=req.k,
        exclude_same_camera=req.exclude_same_camera,
    )

    return ReidSearchResponse(
        matches=[ReidMatch(**m) for m in matches if m["distance"] <= REID_THRESHOLD],
        threshold=REID_THRESHOLD,
    )


@app.post("/recognize", response_model=Optional[RecognizeResponse])
def recognize(req: RecognizeRequest):
    """Узнавание пациента: ищет ближайший эмбеддинг среди записей с patient_id."""
    embedding, _quality, _bbox = detect_best_face(req.frame)
    if embedding is None:
        return None

    result = find_nearest_patient(embedding, threshold=RECOGNITION_THRESHOLD)
    if result is None:
        return None

    return RecognizeResponse(**result)


@app.post("/track-embeddings/raw", response_model=TrackEmbeddingResponse)
def store_raw_track_embedding(req: TrackEmbeddingRawRequest):
    """Сохранить эмбеддинг без детекции лица. Использует fixtures-runner с
    синтетическими векторами; в продовом потоке не используется."""
    record_id = save_embedding_with_meta(
        embedding=req.embedding,
        track_id=req.track_id,
        camera_id=req.camera_id,
        quality=req.quality,
        captured_at=req.captured_at,
        patient_id=req.patient_id,
    )
    return TrackEmbeddingResponse(id=record_id, quality=req.quality)


@app.post("/recognize/embedding", response_model=Optional[RecognizeResponse])
def recognize_by_embedding(req: RecognizeEmbeddingRequest):
    """Узнать пациента по готовому эмбеддингу (для fixtures-runner)."""
    result = find_nearest_patient(req.embedding, threshold=RECOGNITION_THRESHOLD)
    if result is None:
        return None
    return RecognizeResponse(**result)


@app.post("/enroll", response_model=EnrollResponse)
def enroll(req: EnrollRequest):
    """Привязывает все эмбеддинги трека к пациенту (после согласия)."""
    affected = attach_track_to_patient(req.track_id, req.patient_id)
    if affected == 0:
        raise HTTPException(404, f"Не найдено эмбеддингов для трека {req.track_id}")
    logger.info(f"Enroll: трек {req.track_id} → пациент {req.patient_id} ({affected} эмбеддингов)")
    return EnrollResponse(ok=True, embeddings_attached=affected)


@app.delete("/patient/{patient_id}/embeddings")
def delete_embeddings(patient_id: str):
    deleted = delete_patient_embeddings(patient_id)
    logger.info(f"Отозвано согласие пациента {patient_id}: удалено {deleted} эмбеддингов")
    return {"ok": True, "deleted": deleted}


@app.get("/patient/{patient_id}/count")
def patient_count(patient_id: str):
    count = count_patient_embeddings(patient_id)
    return {"patient_id": patient_id, "count": count}