309 lines
10 KiB
Python
309 lines
10 KiB
Python
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) |