75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
# face_recognition_algorithm.py
|
||
from insightface.app import FaceAnalysis
|
||
|
||
|
||
class FaceRecognitionAlgorithm:
|
||
"""
|
||
人脸识别核心算法类
|
||
负责模型加载和设备配置，与业务逻辑分离
|
||
"""
|
||
|
||
def __init__(self, model_name: str = 'buffalo_l', use_gpu: bool = True, use_npu: bool = False,
|
||
npu_device_id: int = 0):
|
||
# 设备配置映射（NPU采用华为指定的完整参数）
|
||
self.DEVICE_CONFIG = {
|
||
"cpu": (['CPUExecutionProvider'], -1),
|
||
"gpu": (['CUDAExecutionProvider'], 0),
|
||
"npu": (
|
||
[
|
||
(
|
||
"CANNExecutionProvider",
|
||
{
|
||
"device_id": npu_device_id,
|
||
"arena_extend_strategy": "kNextPowerOfTwo",
|
||
"npu_mem_limit": 16 * 1024 * 1024 * 1024,
|
||
"op_select_impl_mode": "high_precision",
|
||
"precision_mode": "allow_fp32_to_fp16",
|
||
"enable_cann_graph": True,
|
||
},
|
||
),
|
||
"CPUExecutionProvider"
|
||
],
|
||
npu_device_id
|
||
)
|
||
}
|
||
|
||
# 算法参数设置
|
||
self.det_size = 640 # 320快速 640中等 1280慢
|
||
self.det_threshold = 0.5 # 人脸置信度
|
||
|
||
# 根据设备类型选择配置
|
||
if use_npu:
|
||
device_type = "npu"
|
||
print(f"✅ 使用NPU设备，设备ID: {npu_device_id}")
|
||
elif use_gpu:
|
||
device_type = "gpu"
|
||
print("✅ 使用GPU设备")
|
||
else:
|
||
device_type = "cpu"
|
||
print("✅ 使用CPU设备")
|
||
|
||
providers, ctx_id = self.DEVICE_CONFIG[device_type]
|
||
|
||
# 初始化人脸识别模型
|
||
self.app = FaceAnalysis(name=model_name, providers=providers)
|
||
self.app.prepare(
|
||
ctx_id=ctx_id,
|
||
det_thresh=self.det_threshold,
|
||
det_size=(self.det_size, self.det_size)
|
||
)
|
||
|
||
print(f"✅ 人脸识别算法初始化完成 - 设备: {device_type.upper()}")
|
||
|
||
def set_det_threshold(self, threshold: float):
|
||
"""设置检测阈值"""
|
||
self.det_threshold = threshold
|
||
print(f"✅ 检测阈值设置为: {threshold}")
|
||
|
||
def set_det_size(self, size: int):
|
||
"""设置检测尺寸"""
|
||
self.det_size = size
|
||
print(f"✅ 检测尺寸设置为: {size}")
|
||
|
||
def get_app(self) -> FaceAnalysis:
|
||
"""获取底层的FaceAnalysis实例"""
|
||
return self.app |