优化代码风格
This commit is contained in:
ZZX9599
42
core/face.py
42
core/face.py
@ -16,7 +16,7 @@ try:
|
||||
pynvml.nvmlInit()
|
||||
_nvml_available = True
|
||||
except ImportError:
|
||||
print("警告: pynvml库未安装，无法检测GPU状态，将默认使用0号GPU")
|
||||
print("警告: pynvml库未安装、无法检测GPU状态、将默认使用0号GPU")
|
||||
_nvml_available = False
|
||||
|
||||
# 全局变量
|
||||
@ -58,7 +58,7 @@ def check_gpu_availability(gpu_id, threshold=0.7):
|
||||
|
||||
|
||||
def select_best_gpu(preferred_gpus=[0, 1]):
|
||||
"""选择最佳可用GPU，严格按照首选列表顺序检查，优先使用0号GPU"""
|
||||
"""选择最佳可用GPU、严格按照首选列表顺序检查、优先使用0号GPU"""
|
||||
# 首先检查首选GPU列表
|
||||
for gpu_id in preferred_gpus:
|
||||
try:
|
||||
@ -68,17 +68,17 @@ def select_best_gpu(preferred_gpus=[0, 1]):
|
||||
|
||||
# 检查GPU是否可用
|
||||
if check_gpu_availability(gpu_id):
|
||||
print(f"GPU {gpu_id} 可用，将使用该GPU")
|
||||
print(f"GPU {gpu_id} 可用、将使用该GPU")
|
||||
return gpu_id
|
||||
else:
|
||||
if gpu_id == 0:
|
||||
print(f"GPU 0 内存使用率过高（繁忙），尝试切换到其他GPU")
|
||||
print(f"GPU 0 内存使用率过高（繁忙）、尝试切换到其他GPU")
|
||||
except Exception as e:
|
||||
print(f"GPU {gpu_id} 不存在或无法访问: {e}")
|
||||
continue
|
||||
|
||||
# 如果所有首选GPU都不可用，返回-1表示使用CPU
|
||||
print("所有指定的GPU都不可用，将使用CPU进行计算")
|
||||
# 如果所有首选GPU都不可用、返回-1表示使用CPU
|
||||
print("所有指定的GPU都不可用、将使用CPU进行计算")
|
||||
return -1
|
||||
|
||||
|
||||
@ -122,12 +122,12 @@ def _release_engine():
|
||||
|
||||
|
||||
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:
|
||||
@ -136,7 +136,7 @@ def _monitor_thread():
|
||||
|
||||
|
||||
def load_model(prefer_gpu=True, preferred_gpus=[0, 1]):
|
||||
"""加载人脸识别模型及已知人脸特征库，默认优先使用0号GPU"""
|
||||
"""加载人脸识别模型及已知人脸特征库、默认优先使用0号GPU"""
|
||||
global _face_app, _known_faces_embeddings, _known_faces_names, _using_gpu, _used_gpu_id
|
||||
|
||||
# 确保监控线程只启动一次
|
||||
@ -144,11 +144,11 @@ def load_model(prefer_gpu=True, preferred_gpus=[0, 1]):
|
||||
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
|
||||
|
||||
@ -158,7 +158,7 @@ def load_model(prefer_gpu=True, preferred_gpus=[0, 1]):
|
||||
print("正在初始化InsightFace人脸识别引擎...")
|
||||
_face_app = FaceAnalysis(name='buffalo_l', root='~/.insightface')
|
||||
|
||||
# 选择合适的GPU，默认优先使用0号
|
||||
# 选择合适的GPU、默认优先使用0号
|
||||
ctx_id = 0
|
||||
if prefer_gpu:
|
||||
ctx_id = select_best_gpu(preferred_gpus)
|
||||
@ -166,9 +166,9 @@ def load_model(prefer_gpu=True, preferred_gpus=[0, 1]):
|
||||
_used_gpu_id = ctx_id if _using_gpu else -1
|
||||
|
||||
if _using_gpu:
|
||||
print(f"成功初始化，使用GPU {ctx_id} 进行计算")
|
||||
print(f"成功初始化、使用GPU {ctx_id} 进行计算")
|
||||
else:
|
||||
print("成功初始化，使用CPU进行计算")
|
||||
print("成功初始化、使用CPU进行计算")
|
||||
|
||||
# 准备模型
|
||||
_face_app.prepare(ctx_id=ctx_id, det_size=(640, 640))
|
||||
@ -188,10 +188,10 @@ def load_model(prefer_gpu=True, preferred_gpus=[0, 1]):
|
||||
for person_name, eigenvalue_data in face_data.items():
|
||||
# 处理特征值数据 - 兼容数组和字符串两种格式
|
||||
if isinstance(eigenvalue_data, np.ndarray):
|
||||
# 如果已经是numpy数组，直接使用
|
||||
# 如果已经是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]
|
||||
@ -217,7 +217,7 @@ def load_model(prefer_gpu=True, preferred_gpus=[0, 1]):
|
||||
|
||||
|
||||
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
|
||||
|
||||
@ -248,7 +248,7 @@ def detect(frame, threshold=0.4):
|
||||
return (False, "人脸识别引擎不可用或未初始化")
|
||||
|
||||
try:
|
||||
# 如果使用GPU，确保输入帧在处理前是连续的数组
|
||||
# 如果使用GPU、确保输入帧在处理前是连续的数组
|
||||
if _using_gpu and not frame.flags.contiguous:
|
||||
frame = np.ascontiguousarray(frame)
|
||||
|
||||
@ -285,7 +285,7 @@ def detect(frame, threshold=0.4):
|
||||
# 判断匹配结果
|
||||
is_match = max_sim >= threshold
|
||||
if is_match:
|
||||
has_matched = True # 只要有一个匹配成功，就标记为True
|
||||
has_matched = True # 只要有一个匹配成功、就标记为True
|
||||
|
||||
bbox = face.bbox
|
||||
result_parts.append(
|
||||
@ -298,12 +298,12 @@ def detect(frame, threshold=0.4):
|
||||
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)
|
||||
Reference in New Issue
Block a user