备份一些py文件
This commit is contained in:
309
backup/test_insightface_group.py
Normal file
309
backup/test_insightface_group.py
Normal file
@@ -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)}")
|
||||
Reference in New Issue
Block a user