94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
"""InsightFace wrapper: load model, decode images, extract 512-d embeddings.
|
|
|
|
Скопировано и расширено из work-pcs-adm-time-tracker. Импорты InsightFace и PIL
|
|
сделаны ленивыми — face-service может запускаться без них (SKIP_MODEL_LOAD=true)
|
|
для интеграционных тестов raw-embedding эндпоинтов.
|
|
"""
|
|
import os
|
|
import base64
|
|
import io
|
|
import logging
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MODEL_NAME = os.getenv("MODEL_NAME", "buffalo_l")
|
|
DET_SCORE_THRESHOLD = float(os.getenv("DET_SCORE_THRESHOLD", "0.7"))
|
|
|
|
_app = None # FaceAnalysis | None — ленивый импорт.
|
|
|
|
|
|
def load_model():
|
|
global _app
|
|
if _app is not None:
|
|
return _app
|
|
from insightface.app import FaceAnalysis # ленивый импорт
|
|
logger.info(f"Загружаю модель InsightFace '{MODEL_NAME}'...")
|
|
app = FaceAnalysis(name=MODEL_NAME, providers=["CPUExecutionProvider"])
|
|
app.prepare(ctx_id=0, det_thresh=DET_SCORE_THRESHOLD, det_size=(640, 640))
|
|
_app = app
|
|
logger.info("Модель загружена.")
|
|
return _app
|
|
|
|
|
|
def decode_image(base64_str: str) -> np.ndarray:
|
|
"""Декодирует base64-строку в numpy-массив (BGR, формат OpenCV)."""
|
|
from PIL import Image # ленивый импорт
|
|
if "," in base64_str:
|
|
base64_str = base64_str.split(",", 1)[1]
|
|
|
|
image_bytes = base64.b64decode(base64_str)
|
|
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
|
img_array = np.array(image)
|
|
return img_array[:, :, ::-1].copy()
|
|
|
|
|
|
def detect_best_face(base64_str: str):
|
|
"""Возвращает (embedding, quality, bbox_norm) лучшего лица или (None, None, None).
|
|
|
|
bbox_norm — [x1, y1, x2, y2] в нормализованных 0..1 координатах относительно
|
|
размера изображения. UI рисует overlay поверх displayed image.
|
|
"""
|
|
app = load_model()
|
|
|
|
try:
|
|
img = decode_image(base64_str)
|
|
except Exception as e:
|
|
logger.warning(f"Ошибка декодирования изображения: {e}")
|
|
return None, None, None
|
|
|
|
faces = app.get(img)
|
|
if not faces:
|
|
return None, None, None
|
|
|
|
best_face = max(faces, key=lambda f: f.det_score)
|
|
if best_face.det_score < DET_SCORE_THRESHOLD:
|
|
return None, None, None
|
|
|
|
embedding = best_face.normed_embedding.astype(np.float32)
|
|
quality = float(best_face.det_score)
|
|
|
|
h, w = img.shape[:2]
|
|
box = best_face.bbox.tolist() # [x1, y1, x2, y2] в пикселях
|
|
bbox = {
|
|
"box": [
|
|
max(0, int(box[0])),
|
|
max(0, int(box[1])),
|
|
min(w, int(box[2])),
|
|
min(h, int(box[3])),
|
|
],
|
|
"imgW": w,
|
|
"imgH": h,
|
|
}
|
|
return embedding, quality, bbox
|
|
|
|
|
|
def get_embedding(base64_str: str) -> np.ndarray | None:
|
|
"""Обратная совместимость: только эмбеддинг лучшего лица."""
|
|
embedding, _, _ = detect_best_face(base64_str)
|
|
return embedding
|
|
|
|
|
|
def is_model_loaded() -> bool:
|
|
return _app is not None
|