华为npu insightface demo

src/test_insightface_group.py

@@ -0,0 +1,309 @@
+import cv2
+import insightface
+import numpy as np
+import os
+import datetime
+import argparse
+from insightface.app import FaceAnalysis
+
+# 设备配置映射（NPU采用用户指定的完整参数）
+DEVICE_CONFIG = {
+    "cpu": (['CPUExecutionProvider'], -1),
+	"gpu": (['CUDAExecutionProvider'], 0),
+	"npu": (
+	    [
+		    (
+			    "CANNExecutionProvider",
+				{
+				    "device_id": 1,
+					"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"
+		],
+		0
+	)
+}
+#allow_fp32_to_fp16
+#核心配置参数
+THRESHOLD = 0.65
+IMAGE_EXTENSIONS = ('.jpg','.jepg','.png','.bmp','.gif')
+NPU_REQUIREMENTS = {
+    "依赖包": "onnxruntime-cann（华为官方）+ onnxruntime（基础）",
+	"驱动要求": "Ascend CANN Toolkit ≥ 5.0.3",
+	"硬件要求": "华为昇腾芯片（如Ascend 310/910）",
+	"文档链接": "https://onnxruntime.ai/docs/execution-providers/community-maintained/CANN-ExecutionProvider.html"
+}
+
+# Initialize face analysis model
+#app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])  # Use 'CUDAExecutionProvider' for GPU
+#app.prepare(ctx_id=-1)  # ctx_id=-1 for CPU, 0 for GPU
+
+def parse_args():
+    """解析命令行参数"""
+    parser = argparse.ArgumentParser(description="人脸特征提取与同一人照片分组工具（支持CPU/GPU/NPU）")
+    parser.add_argument("-d","--device", type=str, choices=DEVICE_CONFIG.keys(),
+        default="cpu", help=f"指定运行设备")
+    parser.add_argument("-t","--threshold", type=float, default=THRESHOLD, help=f"相似度阈值，值越大匹配越严格")
+    parser.add_argument("--npu-device-id", type=int, default="0", help=f"覆盖NPU设备ID")
+    parser.add_argument("--npu-mem-limit", type=int, default="16", help=f"NPU内存限制")
+    args = parser.parse_args()
+    return args
+
+def get_face_embedding(image_path, app):
+    """Extract face embedding from an image"""
+    img = cv2.imread(image_path)
+    if img is None:
+        raise ValueError(f"Could not read image: {image_path}")
+    
+    faces = app.get(img)
+    
+    if len(faces) < 1:
+        raise ValueError("No faces detected in the image")
+    if len(faces) > 1:
+        print("Warning: Multiple faces detected. Using first detected face")
+    
+    return faces[0].embedding
+
+def compare_faces(emb1, emb2, threshold): # Adjust this threshold according to your usecase.
+    """Compare two embeddings using cosine similarity"""
+    similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
+    return similarity, similarity > threshold
+	
+def get_all_images_files(directory="."):
+    """获取目录下所有支持的图片文件"""
+    image_files	= []
+    for filename in os.listdir(directory):
+        file_path = os.path.abspath(filename)
+        if filename.lower().endswith(IMAGE_EXTENSIONS) and os.path.isfile(file_path):
+            image_files.append(file_path)
+	
+    return sorted(image_files)
+
+def find_same_person_groups(embedding_dict, threshold):
+    """基于特征向量分组同一人照片"""
+    ungrouped_files = list(embedding_dict.keys())
+    same_person_groups = []
+    single_groups = []
+
+    while ungrouped_files:
+        current_file =ungrouped_files.pop(0)
+        current_emb = embedding_dict[current_file]
+        current_group = [current_file]
+        to_remove = []
+		
+        for candidate_file in ungrouped_files:
+            candidate_emb = embedding_dict[candidate_file]
+            similarity, is_same = compare_faces(current_emb, candidate_emb, threshold)
+            if is_same:
+                current_group.append(candidate_file)
+                to_remove.append(candidate_file)
+                print(f"匹配成功：{os.path.basename(current_file)}和{os.path.basename(candidate_file)}（相似度：{similarity:.4f}）")
+        for file in to_remove:
+            ungrouped_files.remove(file)
+		
+        if len(current_group) >= 2:
+            same_person_groups.append(current_group)
+        else:
+            single_groups.append(current_group)
+	
+    return same_person_groups, single_groups
+
+def generate_log_filename(device):
+    """生成带设备类型和时间戳的日志文件名"""
+    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
+    return f"face_embeddings_{device}_{timestamp}.log"
+def write_embedding_log(log_entries, args, npu_config, log_filename):
+    """写入详细日志"""
+    with open(log_filename, 'w', encoding='utf-8') as f:
+        f.write("="*80 + "\n")
+        f.write(f"人脸特征向量分析日志\n")
+        f.write(f"生成时间： {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
+        f.write(f"运行设备：{args.device}\n")
+        if args.device == "npu":
+            f.write(f"NPU配置：{npu_config}\n")
+        f.write(f"相似度阈值：{args.threshold}\n")
+        f.write(f"支持图片格式：{','.join(IMAGE_EXTENSIONS)}\n")
+        f.write(f"处理图片总数：{len(log_entries)}\n")
+        f.write("="*80 + "\n\n")
+		
+        for entry in log_entries:
+            f.write(entry + "\n" + "-"*60 + "\n\n")
+			
+    print(f"\n 日志文件已保存至：{os.path.abspath(log_filename)}")
+	
+def check_npu_environment():
+    """检查NPU环境是否满足要求（仅检查核心依赖，不验证环境变量）"""
+    try:
+        import onnxruntime as ort
+        available_providers = ort.get_available_providers()
+        if 'CANNExecutionProvider' not in available_providers:
+            return False, "未检测到CANNExecutionProvider（请安装onnxruntime-cann包）"
+
+        return True,"NPU核心依赖检查通过"
+    except ImportError:
+        return False, "未安装onnxruntime（基础依赖）"
+    except Exception as e:
+        return False, f"NPU环境检查失败： {str(e)}"
+
+def main():
+    args = parse_args()
+    device = args.device
+    threshold = args.threshold
+    npu_device_id = args.npu_device_id	
+    npu_mem_limit = args.npu_mem_limit*1024*1024*1024
+    #生成带设备和时间戳的日志文件名
+    log_filename = generate_log_filename(device)
+
+    #打印启动信息
+    print("人脸特征提取与同一人照片分组工具（支持CPU/GPU/NPU）")
+    print("="*80)
+    print(f"核心配置：")
+    print(f"运行设备：{device}")
+    print(f"相似度阈值：{threshold}")
+    print(f"搜索目录：{os.getcwd()}")
+    print(f"支持格式：{','.join(IMAGE_EXTENSIONS)}")
+    print(f"日志文件：{log_filename}")
+    if device == "npu":
+        print(f"NPU设备ID：{npu_device_id}")
+        print(f"NPU内存限制：{args.npu_mem_limit}GB")
+    print("="*80 + "\n")
+	
+    #设备环境预检查
+    if device == "gpu":
+        available_providers = insightface.utils.get_available_providers()
+        if 'CUDAExecutionProvider' not in available_providers:
+            print(" 警告：未检测到CUDA环境，GPU模式可能运行失败！")
+            print(" 解决方案：1.安装CUDA≥11.0 + cuDNN 2.安装onnxruntime-gpu 3. 切换至cpu模式")
+        elif device == "npu":
+            npu_ok,npu_msg = check_npu_environment()
+            print(f" NPU环境检查：{npu_msg}")
+            if not npu_ok:
+                print(" NPU环境不满足要求，建议按以下步骤配置：")
+                for key, value in NPU_REQUIREMENTS.items():
+                    print(f"  - {key}:{value}")
+                return
+				
+				
+    #初始化模型
+    try:
+        providers, ctx_id = DEVICE_CONFIG[device]
+        npu_config = None
+		
+        if device == "npu":
+            npu_provider = list(providers[0])
+            npu_provider[1]["device_id"] = npu_device_id
+            npu_provider[1]["npu_mem_limit"] = npu_mem_limit
+            providers[0] = tuple(npu_provider)
+            npu_config = providers[0][1]
+            ctx_id = npu_device_id
+            app = FaceAnalysis(name='buffalo_l', providers=providers)
+        else:
+            app = FaceAnalysis(name='buffalo_l', providers=providers)
+			
+        app.prepare(ctx_id=ctx_id)
+        print(f" 模型初始化成功（设备：{device}, ctx_id: {ctx_id}）")
+        if device == "npu":
+            print(f" NPU最终配置：")
+            for key, value in npu_config.items():
+                if key == "npu_mem_limit":
+                    print(f"  - {key}: {value/(1024*1024*1024)}GB")
+                else:
+                    print(f"  - {key}: {value}")
+    except Exception as e:
+        print(f"模型初始化失败！")
+        print(f"Error:{str(e)}")
+        print(f" 解决方案:")
+        if device == "gpu":
+            print("   1.确认CUDA驱动已安装 2.确认onnxruntime-gpu版本与CUDA匹配 3.尝试切换CPU模式")
+        elif device == "npu":
+            print("   1.确认CANN Toolkit已正确安装 2.确认onnxruntime-cann版本兼容 3.检查设备ID和内存限制是否合理")
+        return
+
+    #1. 获取所有图片文件
+    image_files = get_all_images_files()
+    if not image_files:
+        print("未找到任何支持的图片文件（检查目录下是否有jpg/png等格式图片）")
+        return
+    print(f"  找到{len(image_files)}个图片文件，开始提取特征向量...\n")
+
+    #2. 提取特征向量并记录日志
+    embedding_dict = {}
+    log_entries = []
+    for img_path in image_files:
+        img_name = os.path.basename(img_path)
+        try:
+            emb = get_face_embedding(img_path, app)
+            embedding_dict[img_name] = emb
+            log_entry = f"【文件】：{img_name}\n" \
+                        f"【路径】：{img_path}\n" \
+                        f"【状态】：成功\n" \
+                        f"【特征向量维度】：{len(emb)}\n" \
+                        f"【特征向量】：{emb.tolist()}"
+            log_entries.append(log_entry)
+            print(f"  处理成功： {img_name}（特征向量维度：{len(emb)})")
+        except Exception as e:
+            log_entry = f"【文件】：{img_name}\n" \
+                        f"【路径】：{img_path}\n" \
+                        f"【状态】：失败\n" \
+                        f":【错误信息】:{str(e)}"
+            log_entries.append(log_entry)
+            print(f"  处理失败： {img_name} - 原因：{str(e)}")
+
+    #3. 写入日志文件（使用动态生成的文件名）
+    write_embedding_log(log_entries, args, npu_config, log_filename)
+
+    #4. 同一人分组分析    
+    if not embedding_dict:
+        print("\n 没有成功提取到人脸特征向量，无法进行分组分析")
+        return
+    
+    print(f"\n 开始分组分析（有效人脸数：{len(embedding_dict)}）...")
+    same_groups, single_groups = find_same_person_groups(embedding_dict, threshold)
+    
+    #5. 输出分组结果
+    print("\n" + "="*80)
+    print("同一人照片分组结果（每组≥2张）:")
+    print("="*80)
+    if same_groups:
+        for i, group in enumerate(same_groups, 1):
+            group_names = [os.path.basename(file) for file in group]
+            print(f"组{i}:{', '.join(group_names)}")
+    else:
+        print(f"  未发现同一人的多张照片")
+        
+    print("\n" + "="*80)
+    print("无匹配的单独照片：")
+    #print("\n"*80)
+    if single_groups:
+        for group in single_groups:
+            print(f"  - {os.path.basename(group[0])}")
+    else:
+        print(f"所有照片均已分组（无单独照片）")
+    print("\n" + "="*80)
+    print(f" 处理完成！详细日志请查看：{log_filename}")
+    
+if __name__ == "__main__":
+    main()
+# Paths to your Indian face images
+#image1_path = "path/to/face1.jpg"
+#image2_path = "path/to/face2.jpg"
+
+#try:
+    # Get embeddings
+#    emb1 = get_face_embedding(image1_path)
+#    emb2 = get_face_embedding(image2_path)
+    
+    # Compare faces
+#    similarity_score, is_same_person = compare_faces(emb1, emb2)
+    
+#    print(f"Similarity Score: {similarity_score:.4f}")
+#    print(f"Same person? {'YES' if is_same_person else 'NO'}")
+    
+#except Exception as e:
+#    print(f"Error: {str(e)}")