136 lines
4.2 KiB
Python
136 lines
4.2 KiB
Python
# Copyright (C) 2018-2021 coneypo
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import os
|
|
import dlib
|
|
import csv
|
|
import numpy as np
|
|
import logging
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
# 要读取人脸图像文件的路径
|
|
path_images_from_camera = "data/data_faces"
|
|
|
|
# Dlib 检测器
|
|
detector = dlib.get_frontal_face_detector()
|
|
predictor = dlib.shape_predictor('models/dlib/shape_predictor_68_face_landmarks.dat')
|
|
face_reco_model = dlib.face_recognition_model_v1("models/dlib/dlib_face_recognition_resnet_model_v1.dat")
|
|
|
|
|
|
def return_128d_features(path_img):
|
|
"""返回单张图像的 128D 特征"""
|
|
try:
|
|
img_pil = Image.open(path_img)
|
|
img_np = np.array(img_pil)
|
|
img_rd = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
|
faces = detector(img_rd, 1)
|
|
|
|
logging.info("%-40s %-20s", "检测到人脸的图像:", path_img)
|
|
|
|
if len(faces) != 0:
|
|
shape = predictor(img_rd, faces[0])
|
|
face_descriptor = face_reco_model.compute_face_descriptor(img_rd, shape)
|
|
return face_descriptor
|
|
else:
|
|
logging.warning("未检测到人脸: %s", path_img)
|
|
return None
|
|
except Exception as e:
|
|
logging.error("处理图像时出错 %s: %s", path_img, e)
|
|
return None
|
|
|
|
|
|
def return_features_mean_personX(path_face_personX):
|
|
"""返回 personX 的 128D 特征均值"""
|
|
features_list_personX = []
|
|
photos_list = os.listdir(path_face_personX)
|
|
|
|
if photos_list:
|
|
for photo in photos_list:
|
|
photo_path = os.path.join(path_face_personX, photo)
|
|
logging.info("正在读取图像: %s", photo_path)
|
|
|
|
features_128d = return_128d_features(photo_path)
|
|
if features_128d is not None:
|
|
features_list_personX.append(features_128d)
|
|
else:
|
|
logging.warning("文件夹为空: %s", path_face_personX)
|
|
|
|
# 计算 128D 特征的均值
|
|
if features_list_personX:
|
|
features_mean_personX = np.array(features_list_personX).mean(axis=0)
|
|
else:
|
|
features_mean_personX = np.zeros(128, dtype=np.float64)
|
|
|
|
return features_mean_personX
|
|
|
|
|
|
def get_person_name_from_folder(folder_name):
|
|
"""从文件夹名称获取有意义的姓名"""
|
|
# 常见的文件夹前缀
|
|
prefixes = ['person_', 'face_', 'user_']
|
|
|
|
for prefix in prefixes:
|
|
if folder_name.startswith(prefix):
|
|
name_part = folder_name[len(prefix):]
|
|
# 如果剩下的部分是纯数字，使用完整文件夹名
|
|
if name_part.isdigit():
|
|
return folder_name
|
|
else:
|
|
return name_part
|
|
|
|
return folder_name
|
|
|
|
|
|
def main():
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
# 检查源文件夹是否存在
|
|
if not os.path.exists(path_images_from_camera):
|
|
logging.error("人脸图像文件夹不存在: %s", path_images_from_camera)
|
|
return
|
|
|
|
# 获取人脸文件夹列表
|
|
person_list = os.listdir(path_images_from_camera)
|
|
person_list.sort()
|
|
|
|
if not person_list:
|
|
logging.error("没有人脸文件夹可处理")
|
|
return
|
|
|
|
logging.info("找到 %d 个人脸文件夹: %s", len(person_list), person_list)
|
|
|
|
# 创建CSV文件
|
|
with open("data/features_all.csv", "w", newline="", encoding="utf-8") as csvfile:
|
|
writer = csv.writer(csvfile)
|
|
|
|
successful_count = 0
|
|
for person_folder in person_list:
|
|
folder_path = os.path.join(path_images_from_camera, person_folder)
|
|
|
|
if not os.path.isdir(folder_path):
|
|
continue
|
|
|
|
logging.info("处理文件夹: %s", person_folder)
|
|
|
|
# 提取特征
|
|
features_mean = return_features_mean_personX(folder_path)
|
|
|
|
# 获取有意义的姓名
|
|
person_name = get_person_name_from_folder(person_folder)
|
|
logging.info("使用姓名: %s", person_name)
|
|
|
|
# 构建行数据：姓名 + 128维特征
|
|
row_data = [person_name] + features_mean.tolist()
|
|
writer.writerow(row_data)
|
|
|
|
successful_count += 1
|
|
logging.info("完成: %s", person_name)
|
|
logging.info("-" * 50)
|
|
|
|
logging.info("成功处理 %d/%d 个人脸文件夹", successful_count, len(person_list))
|
|
logging.info("特征数据已保存到: data/features_all.csv")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main() |