适配华为npu

src/video_face_recognition_cann.py

@@ -8,7 +8,7 @@ import os
 import glob
 
 
-#改进后的人脸质量显示
+# 改进后的人脸质量显示
 
 class VideoFaceRecognition:
     """
@@ -19,14 +19,37 @@ class VideoFaceRecognition:
 
     # buffalo_l， buffalo_sc
 
-    def __init__(self, model_name: str = 'buffalo_l', use_gpu: bool = True):
+    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_size = 640  # 320快速 640中等 1280慢
 
-        #white
+        # white
         self.list_mode = "whitelist"  # "blacklist" 或 "whitelist"
-        self.det_threshold = 0.7    # 人脸置信度
+        self.det_threshold = 0.7  # 人脸置信度
         self.clarity_threshold = 1000.0  # 清晰度阈值，低于此值认为人脸模糊
         self.min_face_size = 30  # 最小人脸像素尺寸
         self.pitch_threshold = 40  #
@@ -44,12 +67,25 @@ class VideoFaceRecognition:
         # self.quality_threshold = 0.6  # 质量得分阈值
         # self.similarity_threshold = 0.3
 
+        # 根据设备类型选择配置
+        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)
+        self.app = FaceAnalysis(name=model_name, providers=providers)
         self.app.prepare(
-            ctx_id=0 if use_gpu else -1,
+            ctx_id=ctx_id,
             det_thresh=self.det_threshold,
-            det_size=(self.det_size,self. det_size)
+            det_size=(self.det_size, self.det_size)
         )
 
         # 名单相关变量
@@ -59,7 +95,7 @@ class VideoFaceRecognition:
         self.frame_count = 0
         self.processing_times = []
 
-        print(f"✅ 视频人脸识别系统初始化完成 - GPU: {use_gpu}")
+        print(f"✅ 视频人脸识别系统初始化完成 - 设备: {device_type.upper()}")
 
     def set_list_mode(self, mode: str):
         """设置名单模式"""
@@ -141,7 +177,6 @@ class VideoFaceRecognition:
 
         return best_name, best_similarity
 
-
     def calculate_clarity(self, face_region: np.ndarray) -> float:
         """
         计算人脸区域的清晰度/模糊度
@@ -301,8 +336,6 @@ class VideoFaceRecognition:
         quality_metrics = result['quality_metrics']
         best_match = result['best_match']
 
-
-
         # 选择颜色
         if not is_acceptable:
             color = (128, 128, 128)  # 灰色 - 质量不可接受
@@ -346,7 +379,6 @@ class VideoFaceRecognition:
         text_lines.append(f"Width: {quality_metrics['bbox_width']:.1f}")
         text_lines.append(f"Height: {quality_metrics['bbox_height']:.1f}")
 
-
         # 计算文本区域大小
         max_text_width = 0
         total_text_height = 0
@@ -620,7 +652,12 @@ class VideoFaceRecognition:
 # 使用示例
 def main():
     # 创建视频识别系统
-    video_system = VideoFaceRecognition(use_gpu=True)
+    # 使用NPU：
+    # video_system = VideoFaceRecognition(use_gpu=False, use_npu=True, npu_device_id=0)
+    # 使用GPU：
+    video_system = VideoFaceRecognition(use_gpu=True, use_npu=False)
+    # 使用CPU：video_system = VideoFaceRecognition(use_gpu=False, use_npu=False)
+    # video_system = VideoFaceRecognition(use_gpu=True, use_npu=False)  # 默认使用GPU
 
     # 设置名单模式
     # video_system.set_list_mode("blacklist")  # 黑名单模式
@@ -640,8 +677,6 @@ def main():
     #     min_face_size=30
     # )
 
-
-
     # # 选择处理模式
     # print("请选择处理模式:")
     # print("1. 处理视频文件")
@@ -654,7 +689,7 @@ def main():
     if choice == "1":
         # 处理视频文件
         video_path = "test_data/video/video_2.mp4"
-        output_path = "test_data/output_video/video_2_white_7_gpu.mp4"
+        output_path = "test_data/output_video/video_2_white_8_gpu.mp4"
         # output_path = "test_data/output_video/video_2_black_2.mp4"
 
         # 性能优化：跳帧处理
@@ -681,4 +716,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()