video/core/face.py

import os
import numpy as np
import cv2
import gc
import time
import threading
from PIL import Image
from insightface.app import FaceAnalysis
# 导入获取人脸信息的服务
from service.face_service import get_all_face_name_with_eigenvalue

# 用于检查GPU状态
try:
    import pynvml

    pynvml.nvmlInit()
    _nvml_available = True
except ImportError:
    print("警告: pynvml库未安装、无法检测GPU状态、将默认使用0号GPU")
    _nvml_available = False

# 全局变量
_face_app = None
_known_faces_embeddings = {}  # 存储姓名到特征值的映射
_known_faces_names = []  # 存储所有已知姓名
_using_gpu = False  # 标记是否使用GPU
_used_gpu_id = -1  # 记录当前使用的GPU ID

# 资源管理变量
_ref_count = 0
_last_used_time = 0
_lock = threading.Lock()
_release_timeout = 8  # 5秒无使用则释放
_is_releasing = False  # 标记是否正在释放

# 调试用计数器
_debug_counter = {
    "created": 0,
    "released": 0,
    "detected": 0
}


def check_gpu_availability(gpu_id, threshold=0.7):
    """检查指定GPU是否可用（内存使用率低于阈值）"""
    if not _nvml_available:
        return True  # 无法检测时默认认为可用

    try:
        handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
        mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
        usage = mem_info.used / mem_info.total
        # 内存使用率低于阈值则认为可用
        return usage < threshold
    except Exception as e:
        print(f"检查GPU {gpu_id} 状态时出错: {e}")
        return False


def select_best_gpu(preferred_gpus=[0, 1]):
    """选择最佳可用GPU、严格按照首选列表顺序检查、优先使用0号GPU"""
    # 首先检查首选GPU列表
    for gpu_id in preferred_gpus:
        try:
            # 检查GPU是否存在
            if _nvml_available:
                pynvml.nvmlDeviceGetHandleByIndex(gpu_id)

            # 检查GPU是否可用
            if check_gpu_availability(gpu_id):
                print(f"GPU {gpu_id} 可用、将使用该GPU")
                return gpu_id
            else:
                if gpu_id == 0:
                    print(f"GPU 0 内存使用率过高（繁忙）、尝试切换到其他GPU")
        except Exception as e:
            print(f"GPU {gpu_id} 不存在或无法访问: {e}")
            continue

    # 如果所有首选GPU都不可用、返回-1表示使用CPU
    print("所有指定的GPU都不可用、将使用CPU进行计算")
    return -1


def _release_engine():
    """释放人脸识别引擎资源"""
    global _face_app, _is_releasing, _known_faces_embeddings, _known_faces_names
    if not _face_app or _is_releasing:
        return

    try:
        _is_releasing = True
        # 释放InsightFace资源
        if hasattr(_face_app, 'model'):
            # 清除模型资源
            _face_app.model = None
        _face_app = None

        # 清空人脸数据
        _known_faces_embeddings.clear()
        _known_faces_names.clear()

        _debug_counter["released"] += 1
        print(f"Face recognition engine released. Stats: {_debug_counter}")

        # 清理GPU缓存
        gc.collect()
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.ipc_collect()
        except ImportError:
            pass
        try:
            import tensorflow as tf
            tf.keras.backend.clear_session()
        except ImportError:
            pass
    finally:
        _is_releasing = False


def _monitor_thread():
    """监控线程、检查并释放超时未使用的资源"""
    global _ref_count, _last_used_time, _face_app
    while True:
        time.sleep(5)  # 每5秒检查一次
        with _lock:
            # 只有当引擎存在、没有引用且超时、才释放
            if _face_app and _ref_count == 0 and not _is_releasing:
                elapsed = time.time() - _last_used_time
                if elapsed > _release_timeout:
                    print(f"Idle timeout ({elapsed:.1f}s > {_release_timeout}s), releasing face engine")
                    _release_engine()


def load_model(prefer_gpu=True, preferred_gpus=[0, 1]):
    """加载人脸识别模型及已知人脸特征库、默认优先使用0号GPU"""
    global _face_app, _known_faces_embeddings, _known_faces_names, _using_gpu, _used_gpu_id

    # 确保监控线程只启动一次
    if not any(t.name == "FaceMonitor" for t in threading.enumerate()):
        threading.Thread(target=_monitor_thread, daemon=True, name="FaceMonitor").start()
        print("Face monitor thread started")

    # 如果正在释放中、等待释放完成
    while _is_releasing:
        time.sleep(0.1)

    # 如果已经初始化、直接返回
    if _face_app:
        return True

    # 初始化InsightFace模型
    try:
        # 初始化InsightFace
        print("正在初始化InsightFace人脸识别引擎...")
        _face_app = FaceAnalysis(name='buffalo_l', root='~/.insightface')

        # 选择合适的GPU、默认优先使用0号
        ctx_id = 0
        if prefer_gpu:
            ctx_id = select_best_gpu(preferred_gpus)
            _using_gpu = ctx_id != -1
            _used_gpu_id = ctx_id if _using_gpu else -1

            if _using_gpu:
                print(f"成功初始化、使用GPU {ctx_id} 进行计算")
            else:
                print("成功初始化、使用CPU进行计算")

        # 准备模型
        _face_app.prepare(ctx_id=ctx_id, det_size=(640, 640))
        print("InsightFace人脸识别引擎初始化成功。")
        _debug_counter["created"] += 1
        print(f"Face engine initialized. Stats: {_debug_counter}")

    except Exception as e:
        print(f"初始化失败: {e}")
        return False

    # 从服务获取所有人脸姓名和特征值
    try:
        face_data = get_all_face_name_with_eigenvalue()

        # 处理获取到的人脸数据
        for person_name, eigenvalue_data in face_data.items():
            # 处理特征值数据 - 兼容数组和字符串两种格式
            if isinstance(eigenvalue_data, np.ndarray):
                # 如果已经是numpy数组、直接使用
                eigenvalue = eigenvalue_data.astype(np.float32)
            elif isinstance(eigenvalue_data, str):
                # 清理字符串: 移除方括号、换行符和多余空格
                cleaned = eigenvalue_data.replace('[', '').replace(']', '').replace('\n', '').strip()
                # 按空格或逗号分割（处理可能的不同分隔符）
                values = [v for v in cleaned.split() if v]
                # 转换为数组
                eigenvalue = np.array(list(map(float, values)), dtype=np.float32)
            else:
                # 不支持的类型
                print(f"Unsupported eigenvalue type for {person_name}")
                continue

            # 归一化处理
            norm = np.linalg.norm(eigenvalue)
            if norm != 0:
                eigenvalue = eigenvalue / norm

            _known_faces_embeddings[person_name] = eigenvalue
            _known_faces_names.append(person_name)

    except Exception as e:
        print(f"Error loading face data from service: {e}")

    return True if _face_app else False


def detect(frame, threshold=0.4):
    """检测并识别人脸、返回结果元组(是否匹配到已知人脸, 结果字符串)"""
    global _face_app, _known_faces_embeddings, _known_faces_names, _using_gpu, _used_gpu_id
    global _ref_count, _last_used_time

    # 验证前置条件
    if frame is None or frame.size == 0:
        return (False, "无效帧数据")

    # 增加引用计数并获取引擎实例
    engine = None
    with _lock:
        _ref_count += 1
        _last_used_time = time.time()
        _debug_counter["detected"] += 1

        # 初始化引擎（如果未初始化且不在释放中）
        if not _face_app and not _is_releasing:
            if not load_model(prefer_gpu=True):
                _ref_count -= 1  # 恢复引用计数
                return (False, "引擎初始化失败")

        # 获取当前引擎引用
        engine = _face_app

    # 检查引擎是否可用
    if not engine or not _known_faces_names:
        with _lock:
            _ref_count = max(0, _ref_count - 1)
        return (False, "人脸识别引擎不可用或未初始化")

    try:
        # 如果使用GPU、确保输入帧在处理前是连续的数组
        if _using_gpu and not frame.flags.contiguous:
            frame = np.ascontiguousarray(frame)

        faces = _face_app.get(frame)
    except Exception as e:
        print(f"Face detect error: {e}")
        # 检测到错误时尝试重新选择GPU并重新初始化
        print("尝试重新选择GPU并重新初始化...")
        with _lock:
            _ref_count = max(0, _ref_count - 1)
        load_model(prefer_gpu=True)  # 重新初始化时保持默认GPU优先级
        return (False, f"检测错误: {str(e)}")

    result_parts = []
    has_matched = False  # 标记是否有匹配到的已知人脸

    for face in faces:
        # 特征归一化
        embedding = face.embedding.astype(np.float32)
        norm = np.linalg.norm(embedding)
        if norm == 0:
            continue
        embedding = embedding / norm

        # 对比已知人脸
        max_sim, best_name = -1.0, "Unknown"
        for name in _known_faces_names:
            known_emb = _known_faces_embeddings[name]
            sim = np.dot(embedding, known_emb)
            if sim > max_sim:
                max_sim = sim
                best_name = name

        # 判断匹配结果
        is_match = max_sim >= threshold
        if is_match:
            has_matched = True  # 只要有一个匹配成功、就标记为True

        bbox = face.bbox
        result_parts.append(
            f"{'匹配' if is_match else '不匹配'}: {best_name} (相似度: {max_sim:.2f}, 边界框: {bbox})"
        )

    # 构建结果字符串
    if not result_parts:
        result_str = "未检测到人脸"
    else:
        result_str = "; ".join(result_parts)

    # 减少引用计数、确保线程安全
    with _lock:
        _ref_count = max(0, _ref_count - 1)
        # 持续使用时更新最后使用时间
        if _ref_count > 0:
            _last_used_time = time.time()

    # 第一个返回值为: 是否匹配到已知人脸
    return (has_matched, result_str)