113 lines
4.0 KiB
Python
113 lines
4.0 KiB
Python
import cv2
|
|
import numpy as np
|
|
import os
|
|
from tqdm import tqdm
|
|
import json
|
|
from deepface import DeepFace
|
|
import redis
|
|
from config import REDIS_IDENTITY, PATH_CONFIG
|
|
|
|
class FaceFeatureExtractor:
|
|
def __init__(self):
|
|
"""
|
|
初始化特征提取器和Redis连接
|
|
"""
|
|
self.redis_client = REDIS_IDENTITY
|
|
|
|
def get_feature(self, img_path):
|
|
"""
|
|
获取人脸特征向量
|
|
Args:
|
|
img_path: 图片路径
|
|
Returns:
|
|
人脸特征向量或None（如果检测失败）
|
|
"""
|
|
try:
|
|
# 使用 DeepFace 提取特征
|
|
embedding_objs = DeepFace.represent(
|
|
img_path=img_path,
|
|
detector_backend="retinaface",
|
|
align=True,
|
|
model_name="Facenet512"
|
|
)
|
|
|
|
if embedding_objs and len(embedding_objs) > 0:
|
|
return embedding_objs[0]['embedding'] # 返回第一个人脸的特征向量
|
|
return None
|
|
|
|
except Exception as e:
|
|
print(f"处理图片 {img_path} 时出错: {str(e)}")
|
|
return None
|
|
|
|
def process_dataset(self, dataset_dir):
|
|
"""
|
|
处理数据集并保存特征到Redis
|
|
Args:
|
|
dataset_dir: 数据集根目录
|
|
"""
|
|
total_processed = 0
|
|
total_failed = 0
|
|
|
|
# 计算所有文件夹中的图片总数
|
|
total_images = sum(len([f for f in os.listdir(os.path.join(dataset_dir, d))
|
|
if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
|
|
for d in os.listdir(dataset_dir)
|
|
if os.path.isdir(os.path.join(dataset_dir, d)))
|
|
|
|
with tqdm(total=total_images, desc="处理图片") as pbar:
|
|
# 处理所有文件夹
|
|
for person_name in sorted(os.listdir(dataset_dir)):
|
|
person_dir = os.path.join(dataset_dir, person_name)
|
|
if not os.path.isdir(person_dir):
|
|
continue
|
|
|
|
# 获取该人已有的特征向量列表
|
|
existing_features = []
|
|
if self.redis_client.exists(person_name):
|
|
existing_features = json.loads(self.redis_client.get(person_name))
|
|
|
|
person_features = existing_features
|
|
|
|
for img_name in os.listdir(person_dir):
|
|
if not img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
|
|
continue
|
|
|
|
# 检查图片是否已处理
|
|
if any(item['image'] == img_name for item in existing_features):
|
|
pbar.update(1)
|
|
continue
|
|
|
|
img_path = os.path.join(person_dir, img_name)
|
|
feature = self.get_feature(img_path)
|
|
|
|
if feature is not None:
|
|
# 特征向量已经是列表格式，无需调用 tolist()
|
|
person_features.append({
|
|
'image': img_name,
|
|
'embedding': feature
|
|
})
|
|
total_processed += 1
|
|
else:
|
|
total_failed += 1
|
|
|
|
pbar.set_postfix({'成功': total_processed, '失败': total_failed})
|
|
pbar.update(1)
|
|
|
|
# 将该人的所有特征向量保存到Redis
|
|
if person_features:
|
|
self.redis_client.set(person_name, json.dumps(person_features))
|
|
|
|
print(f"\n处理完成！成功: {total_processed}, 失败: {total_failed}")
|
|
|
|
def main():
|
|
# 设置路径
|
|
DATASET_DIR = PATH_CONFIG['data']
|
|
|
|
# 创建特征提取器
|
|
extractor = FaceFeatureExtractor()
|
|
|
|
# 处理数据集
|
|
extractor.process_dataset(DATASET_DIR)
|
|
|
|
if __name__ == "__main__":
|
|
main() |