最新可用
This commit is contained in:
222
core/face.py
222
core/face.py
@ -1,27 +1,183 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image # 确保正确导入Image类
|
||||
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 load_model():
|
||||
"""加载人脸识别模型及已知人脸特征库"""
|
||||
global _face_app, _known_faces_embeddings, _known_faces_names
|
||||
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:
|
||||
_face_app = FaceAnalysis(name='buffalo_l', root=os.path.expanduser('~/.insightface'))
|
||||
_face_app.prepare(ctx_id=0, det_size=(640, 640))
|
||||
# 初始化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"Face model load failed: {e}")
|
||||
print(f"初始化失败: {e}")
|
||||
return False
|
||||
|
||||
# 从服务获取所有人脸姓名和特征值
|
||||
@ -62,19 +218,52 @@ def load_model():
|
||||
|
||||
def detect(frame, threshold=0.4):
|
||||
"""检测并识别人脸，返回结果元组(是否匹配到已知人脸, 结果字符串)"""
|
||||
global _face_app, _known_faces_embeddings, _known_faces_names
|
||||
global _face_app, _known_faces_embeddings, _known_faces_names, _using_gpu, _used_gpu_id
|
||||
global _ref_count, _last_used_time
|
||||
|
||||
if not _face_app or not _known_faces_names or frame is None:
|
||||
return (False, "未初始化或无效帧")
|
||||
# 验证前置条件
|
||||
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 # 新增标记：是否有匹配到的已知人脸
|
||||
has_matched = False # 标记是否有匹配到的已知人脸
|
||||
|
||||
for face in faces:
|
||||
# 特征归一化
|
||||
@ -109,5 +298,12 @@ def detect(frame, threshold=0.4):
|
||||
else:
|
||||
result_str = "; ".join(result_parts)
|
||||
|
||||
# 第一个返回值改为：是否匹配到已知人脸
|
||||
return (has_matched, result_str)
|
||||
# 减少引用计数，确保线程安全
|
||||
with _lock:
|
||||
_ref_count = max(0, _ref_count - 1)
|
||||
# 持续使用时更新最后使用时间
|
||||
if _ref_count > 0:
|
||||
_last_used_time = time.time()
|
||||
|
||||
# 第一个返回值为：是否匹配到已知人脸
|
||||
return (has_matched, result_str)
|
||||
Reference in New Issue
Block a user