289 lines
10 KiB
Python
289 lines
10 KiB
Python
"""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}
|