482 lines
17 KiB
Python
482 lines
17 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
from insightface.app import FaceAnalysis
|
|
import pickle
|
|
from typing import List, Tuple, Dict, Optional
|
|
from datetime import datetime
|
|
import json
|
|
|
|
|
|
class SingleFaceComparisonSystem:
|
|
"""
|
|
单人脸比对系统
|
|
在图像中检测多个人脸，并与单个目标人脸进行比对
|
|
"""
|
|
|
|
def __init__(self, model_name: str = 'buffalo_l'):
|
|
"""
|
|
初始化人脸比对系统
|
|
|
|
Args:
|
|
model_name: 模型名称，可选 'buffalo_l', 'buffalo_s' 等
|
|
"""
|
|
# 初始化InsightFace
|
|
self.app = FaceAnalysis(name=model_name, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
|
|
self.app.prepare(ctx_id=0,
|
|
det_thresh=0.35, # 降低阈值，检测更多模糊/小脸
|
|
det_size=(640, 640))
|
|
|
|
# 目标人脸特征
|
|
self.target_embedding = None
|
|
self.target_person_id = None
|
|
self.target_image_path = None
|
|
|
|
# 相似度阈值
|
|
self.similarity_threshold = 0.3
|
|
|
|
# 颜色配置
|
|
self.colors = {
|
|
'match_high': (0, 255, 0), # 绿色 - 高相似度
|
|
'match_medium': (0, 255, 255), # 黄色 - 中等相似度
|
|
'match_low': (0, 165, 255), # 橙色 - 低相似度
|
|
'no_match': (0, 0, 255), # 红色 - 不匹配
|
|
'text': (255, 255, 255), # 白色 - 文字
|
|
'landmark': (255, 0, 0) # 蓝色 - 关键点
|
|
}
|
|
|
|
def set_target_face(self, image_path: str, person_id: str = "target_person") -> bool:
|
|
"""
|
|
设置目标人脸
|
|
|
|
Args:
|
|
image_path: 目标人脸图像路径
|
|
person_id: 目标人员ID
|
|
|
|
Returns:
|
|
bool: 设置是否成功
|
|
"""
|
|
if not os.path.exists(image_path):
|
|
print(f"目标图像文件不存在: {image_path}")
|
|
return False
|
|
|
|
# 提取人脸特征
|
|
faces = self.extract_face_features(image_path)
|
|
|
|
if not faces:
|
|
print(f"在目标图像 {image_path} 中未检测到人脸")
|
|
return False
|
|
|
|
if len(faces) > 1:
|
|
print(f"在目标图像 {image_path} 中检测到多个人脸，将使用第一个人脸")
|
|
|
|
# 存储目标人脸特征
|
|
self.target_embedding = np.array(faces[0]['embedding'])
|
|
self.target_person_id = person_id
|
|
self.target_image_path = image_path
|
|
|
|
print(f"成功设置目标人脸: {person_id}")
|
|
return True
|
|
|
|
def extract_face_features(self, image_path: str) -> List[Dict]:
|
|
"""
|
|
从图像中提取所有人脸特征
|
|
|
|
Args:
|
|
image_path: 图像路径
|
|
|
|
Returns:
|
|
List[Dict]: 包含多个人脸信息的列表
|
|
"""
|
|
img = cv2.imread(image_path)
|
|
if img is None:
|
|
raise ValueError(f"无法读取图像: {image_path}")
|
|
|
|
faces = self.app.get(img)
|
|
results = []
|
|
|
|
for i, face in enumerate(faces):
|
|
face_info = {
|
|
'bbox': face.bbox.astype(int).tolist(), # 人脸框 [x1, y1, x2, y2]
|
|
'kps': face.kps.astype(int).tolist(), # 关键点 [[x1,y1], [x2,y2], ...]
|
|
'embedding': face.embedding.tolist(), # 特征向量
|
|
'gender': 'Male' if face.gender == 1 else 'Female',
|
|
'age': int(face.age),
|
|
'det_score': float(face.det_score) # 检测置信度
|
|
}
|
|
results.append(face_info)
|
|
|
|
return results
|
|
|
|
def compare_with_target(self, query_embedding: np.ndarray) -> float:
|
|
"""
|
|
与目标人脸进行比对
|
|
|
|
Args:
|
|
query_embedding: 查询特征向量
|
|
|
|
Returns:
|
|
float: 相似度 (0-1)
|
|
"""
|
|
if self.target_embedding is None:
|
|
raise ValueError("未设置目标人脸")
|
|
|
|
# 归一化向量
|
|
query_embedding = query_embedding / np.linalg.norm(query_embedding)
|
|
target_embedding = self.target_embedding / np.linalg.norm(self.target_embedding)
|
|
|
|
# 计算余弦相似度
|
|
similarity = np.dot(query_embedding, target_embedding)
|
|
return float(similarity)
|
|
|
|
def get_similarity_color(self, similarity: float) -> tuple:
|
|
"""
|
|
根据相似度获取对应的颜色
|
|
|
|
Args:
|
|
similarity: 相似度 (0-1)
|
|
|
|
Returns:
|
|
tuple: BGR颜色值
|
|
"""
|
|
if similarity >= self.similarity_threshold:
|
|
return self.colors['match_high'] # 高相似度 - 绿色
|
|
# elif similarity >= self.similarity_threshold/2:
|
|
# return self.colors['match_low'] # 低相似度 - 橙色
|
|
else:
|
|
return self.colors['no_match'] # 不匹配 - 红色
|
|
|
|
def process_image_with_target_comparison(self, image_path: str, output_path: str = None,
|
|
draw_landmarks: bool = True, draw_info: bool = True) -> Dict:
|
|
"""
|
|
处理图像并与目标人脸进行比对
|
|
|
|
Args:
|
|
image_path: 输入图像路径
|
|
output_path: 输出图像路径
|
|
draw_landmarks: 是否绘制关键点
|
|
draw_info: 是否绘制详细信息
|
|
|
|
Returns:
|
|
Dict: 处理结果
|
|
"""
|
|
if self.target_embedding is None:
|
|
raise ValueError("请先设置目标人脸")
|
|
|
|
# 读取图像
|
|
img = cv2.imread(image_path)
|
|
if img is None:
|
|
raise ValueError(f"无法读取图像: {image_path}")
|
|
|
|
# 提取人脸特征
|
|
faces = self.extract_face_features(image_path)
|
|
comparison_results = {
|
|
'image_path': image_path,
|
|
'target_person_id': self.target_person_id,
|
|
'faces_detected': len(faces),
|
|
'comparisons': []
|
|
}
|
|
|
|
# 对每个检测到的人脸进行比对
|
|
for i, face in enumerate(faces):
|
|
# 与目标人脸比对
|
|
query_embedding = np.array(face['embedding'])
|
|
similarity = self.compare_with_target(query_embedding)
|
|
|
|
# 存储比对结果
|
|
face_comparison = {
|
|
'face_index': i,
|
|
'bbox': face['bbox'],
|
|
'similarity': similarity,
|
|
'is_match': similarity >= self.similarity_threshold,
|
|
'gender': face['gender'],
|
|
'age': face['age'],
|
|
'det_score': face['det_score']
|
|
}
|
|
comparison_results['comparisons'].append(face_comparison)
|
|
|
|
# 在图像上绘制标注
|
|
self._draw_face_comparison_annotation(img, face_comparison, draw_landmarks, draw_info)
|
|
|
|
# 保存输出图像
|
|
if output_path:
|
|
os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else '.', exist_ok=True)
|
|
cv2.imwrite(output_path, img)
|
|
print(f"比对结果图像已保存: {output_path}")
|
|
|
|
return comparison_results
|
|
|
|
def _draw_face_comparison_annotation(self, img: np.ndarray, face_info: Dict,
|
|
draw_landmarks: bool, draw_info: bool):
|
|
"""
|
|
在图像上绘制人脸比对标注
|
|
|
|
Args:
|
|
img: 图像数组
|
|
face_info: 人脸比对信息
|
|
draw_landmarks: 是否绘制关键点
|
|
draw_info: 是否绘制详细信息
|
|
"""
|
|
bbox = face_info['bbox'] # [x1, y1, x2, y2]
|
|
similarity = face_info['similarity']
|
|
is_match = face_info['is_match']
|
|
|
|
# 根据相似度确定颜色
|
|
color = self.get_similarity_color(similarity)
|
|
|
|
# 绘制人脸框
|
|
thickness = 3 if is_match else 2
|
|
cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, thickness)
|
|
|
|
# 绘制关键点
|
|
if draw_landmarks and 'kps' in face_info:
|
|
for kp in face_info['kps']:
|
|
cv2.circle(img, (kp[0], kp[1]), 2, self.colors['landmark'], -1)
|
|
|
|
# 绘制信息文本
|
|
if draw_info:
|
|
# 相似度信息
|
|
text_y = bbox[1] - 10
|
|
if text_y < 20:
|
|
text_y = bbox[3] + 20
|
|
|
|
similarity_text = f"Similarity: {similarity:.3f}"
|
|
match_status = "MATCH" if is_match else "NO MATCH"
|
|
|
|
cv2.putText(img, similarity_text, (bbox[0], text_y),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, self.colors['text'], 2)
|
|
|
|
# 匹配状态
|
|
cv2.putText(img, match_status, (bbox[0], text_y + 25),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
|
|
|
|
# 人脸属性信息
|
|
detail_text = f"{face_info['gender']}/{face_info['age']}"
|
|
cv2.putText(img, detail_text, (bbox[0], text_y + 50),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.colors['text'], 1)
|
|
|
|
def batch_process_with_target(self, input_dir: str, output_dir: str,
|
|
image_extensions: tuple = ('.jpg', '.jpeg', '.png')) -> Dict:
|
|
"""
|
|
批量处理目录中的所有图像，与目标人脸进行比对
|
|
|
|
Args:
|
|
input_dir: 输入图像目录
|
|
output_dir: 输出图像目录
|
|
image_extensions: 图像文件扩展名
|
|
|
|
Returns:
|
|
Dict: 批量处理结果统计
|
|
"""
|
|
if self.target_embedding is None:
|
|
raise ValueError("请先设置目标人脸")
|
|
|
|
if not os.path.exists(input_dir):
|
|
raise ValueError(f"输入目录不存在: {input_dir}")
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# 获取所有图像文件
|
|
image_files = []
|
|
for ext in image_extensions:
|
|
image_files.extend([f for f in os.listdir(input_dir) if f.lower().endswith(ext)])
|
|
|
|
print(f"找到 {len(image_files)} 个图像文件")
|
|
|
|
# 处理统计
|
|
stats = {
|
|
'target_person_id': self.target_person_id,
|
|
'total_images': len(image_files),
|
|
'processed_images': 0,
|
|
'total_faces': 0,
|
|
'matched_faces': 0,
|
|
'max_similarity': 0.0,
|
|
'min_similarity': 1.0,
|
|
'avg_similarity': 0.0,
|
|
'results': []
|
|
}
|
|
|
|
all_similarities = []
|
|
|
|
# 批量处理
|
|
for image_file in image_files:
|
|
input_path = os.path.join(input_dir, image_file)
|
|
output_path = os.path.join(output_dir, f"compared_{image_file}")
|
|
|
|
try:
|
|
# 处理单张图像
|
|
result = self.process_image_with_target_comparison(input_path, output_path)
|
|
|
|
# 更新统计
|
|
stats['processed_images'] += 1
|
|
stats['total_faces'] += result['faces_detected']
|
|
|
|
image_similarities = [comp['similarity'] for comp in result['comparisons']]
|
|
all_similarities.extend(image_similarities)
|
|
|
|
if image_similarities:
|
|
stats['max_similarity'] = max(stats['max_similarity'], max(image_similarities))
|
|
stats['min_similarity'] = min(stats['min_similarity'], min(image_similarities))
|
|
|
|
matched_count = sum(1 for comp in result['comparisons'] if comp['is_match'])
|
|
stats['matched_faces'] += matched_count
|
|
|
|
stats['results'].append({
|
|
'image_file': image_file,
|
|
'faces_detected': result['faces_detected'],
|
|
'matched_faces': matched_count,
|
|
'max_similarity': max(image_similarities) if image_similarities else 0,
|
|
'min_similarity': min(image_similarities) if image_similarities else 0
|
|
})
|
|
|
|
print(f"处理完成: {image_file} -> 检测到 {result['faces_detected']} 张脸, "
|
|
f"匹配 {matched_count} 张, 最高相似度: {max(image_similarities) if image_similarities else 0:.3f}")
|
|
|
|
except Exception as e:
|
|
print(f"处理图像 {image_file} 时出错: {e}")
|
|
|
|
# 计算平均相似度
|
|
if all_similarities:
|
|
stats['avg_similarity'] = sum(all_similarities) / len(all_similarities)
|
|
|
|
# 保存处理统计
|
|
stats_path = os.path.join(output_dir, "comparison_stats.json")
|
|
with open(stats_path, 'w', encoding='utf-8') as f:
|
|
json.dump(stats, f, indent=2, ensure_ascii=False)
|
|
|
|
print(f"批量比对完成! 统计信息已保存: {stats_path}")
|
|
return stats
|
|
|
|
def set_similarity_threshold(self, threshold: float):
|
|
"""设置相似度阈值"""
|
|
if 0 <= threshold <= 1:
|
|
self.similarity_threshold = threshold
|
|
print(f"相似度阈值已设置为: {threshold}")
|
|
else:
|
|
print("阈值必须在 0 到 1 之间")
|
|
|
|
def get_target_info(self) -> Dict:
|
|
"""获取目标人脸信息"""
|
|
if self.target_embedding is None:
|
|
return {"status": "未设置目标人脸"}
|
|
|
|
return {
|
|
"status": "已设置",
|
|
"person_id": self.target_person_id,
|
|
"image_path": self.target_image_path,
|
|
"similarity_threshold": self.similarity_threshold
|
|
}
|
|
|
|
|
|
# 使用示例和演示
|
|
def demo_usage():
|
|
"""演示使用方法"""
|
|
|
|
# 创建人脸比对系统
|
|
face_system = SingleFaceComparisonSystem()
|
|
|
|
# 1. 设置目标人脸
|
|
print("=== 设置目标人脸 ===")
|
|
|
|
# 创建测试数据目录
|
|
os.makedirs("test_data/target", exist_ok=True)
|
|
os.makedirs("test_data/query", exist_ok=True)
|
|
os.makedirs("test_data/output", exist_ok=True)
|
|
|
|
# target_image = "test_data/register/person2.jpg"
|
|
# target_image = "test_data/register/person1.png"
|
|
target_image = "test_data/register/ztk.jpg"
|
|
|
|
|
|
# 设置目标人脸
|
|
if os.path.exists(target_image):
|
|
face_system.set_target_face(target_image, "target_person")
|
|
else:
|
|
print(f"目标图像不存在: {target_image}")
|
|
print("请将目标人脸图像放在 test_data/target/ 目录下")
|
|
return
|
|
|
|
# 2. 查看目标信息
|
|
print("\n=== 目标人脸信息 ===")
|
|
target_info = face_system.get_target_info()
|
|
print(f"目标人员ID: {target_info['person_id']}")
|
|
print(f"目标图像路径: {target_info['image_path']}")
|
|
print(f"相似度阈值: {target_info['similarity_threshold']}")
|
|
|
|
# 3. 单张图像比对示例
|
|
print("\n=== 单张图像比对 ===")
|
|
query_image = "test_data/query/multi_face.jpg" # 包含多个人的图像
|
|
|
|
if os.path.exists(query_image):
|
|
result = face_system.process_image_with_target_comparison(
|
|
image_path=query_image,
|
|
output_path="test_data/output/compared_multi_face.jpg",
|
|
draw_landmarks=True,
|
|
draw_info=True
|
|
)
|
|
|
|
print(f"检测到 {result['faces_detected']} 张人脸:")
|
|
for face in result['comparisons']:
|
|
status = "匹配" if face['is_match'] else "不匹配"
|
|
color_name = "绿色" if face['similarity'] >= 0.8 else \
|
|
"黄色" if face['similarity'] >= 0.6 else \
|
|
"橙色" if face['similarity'] >= 0.4 else "红色"
|
|
|
|
print(f" 人脸 {face['face_index'] + 1}: 相似度 {face['similarity']:.3f} "
|
|
f"({color_name}, {status}), {face['gender']}/{face['age']}岁")
|
|
else:
|
|
print(f"查询图像不存在: {query_image}")
|
|
print("请将包含多人脸的图像放在 test_data/query/ 目录下")
|
|
|
|
# 4. 批量处理示例
|
|
print("\n=== 批量比对图像 ===")
|
|
if os.path.exists("test_data/query") and len(os.listdir("test_data/query")) > 0:
|
|
stats = face_system.batch_process_with_target(
|
|
input_dir="test_data/query/2",
|
|
output_dir="test_data/output/batch_comparison/2"
|
|
)
|
|
|
|
print(f"批量比对统计:")
|
|
print(f" 目标人员: {stats['target_person_id']}")
|
|
print(f" 总图像数: {stats['total_images']}")
|
|
print(f" 成功处理: {stats['processed_images']}")
|
|
print(f" 总人脸数: {stats['total_faces']}")
|
|
print(f" 匹配人脸: {stats['matched_faces']}")
|
|
print(f" 最高相似度: {stats['max_similarity']:.3f}")
|
|
print(f" 最低相似度: {stats['min_similarity']:.3f}")
|
|
print(f" 平均相似度: {stats['avg_similarity']:.3f}")
|
|
|
|
# 5. 调整相似度阈值
|
|
print("\n=== 调整阈值 ===")
|
|
face_system.set_similarity_threshold(0.3) # 提高阈值，更严格
|
|
|
|
|
|
# 高级使用示例
|
|
def advanced_usage():
|
|
"""高级使用示例"""
|
|
face_system = SingleFaceComparisonSystem()
|
|
|
|
# 设置目标人脸
|
|
face_system.set_target_face("path/to/target_face.jpg", "suspect_A")
|
|
|
|
# 设置更严格的阈值
|
|
face_system.set_similarity_threshold(0.75)
|
|
|
|
# 处理单个图像并保存结果
|
|
result = face_system.process_image_with_target_comparison(
|
|
image_path="path/to/surveillance_image.jpg",
|
|
output_path="path/to/analyzed_image.jpg"
|
|
)
|
|
|
|
# 输出详细信息
|
|
print(f"分析完成: {result['image_path']}")
|
|
print(f"目标人员: {result['target_person_id']}")
|
|
print(f"检测到人脸: {result['faces_detected']}")
|
|
|
|
for comp in result['comparisons']:
|
|
print(f"人脸 {comp['face_index'] + 1}: 相似度={comp['similarity']:.3f}, "
|
|
f"匹配={comp['is_match']}, 位置={comp['bbox']}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
demo_usage() |