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.
93 lines
3.1 KiB
93 lines
3.1 KiB
"""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
|
|
|