233 lines
8.5 KiB
Python
Executable File
233 lines
8.5 KiB
Python
Executable File
# Copyright (C) 2018-2021 coneypo
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import dlib
|
|
import numpy as np
|
|
import cv2
|
|
import os
|
|
import shutil
|
|
import time
|
|
import logging
|
|
|
|
# Dlib 正向人脸检测器
|
|
detector = dlib.get_frontal_face_detector()
|
|
|
|
|
|
class Face_Register:
|
|
def __init__(self):
|
|
self.path_photos_from_camera = "data/data_faces/"
|
|
self.font = cv2.FONT_ITALIC
|
|
|
|
self.existing_faces_cnt = 0 # 已注册的人数统计
|
|
self.ss_cnt = 0 # 当前人员已保存图片数量
|
|
self.current_frame_faces_cnt = 0 # 用于存储摄像头中出现的人数
|
|
|
|
self.save_flag = 1
|
|
self.press_n_flag = 0
|
|
self.current_person_name = "" # 新增：当前录入人姓名
|
|
|
|
# FPS
|
|
self.frame_time = 0
|
|
self.frame_start_time = 0
|
|
self.fps = 0
|
|
self.fps_show = 0
|
|
self.start_time = time.time()
|
|
|
|
# 创建数据目录
|
|
def pre_work_mkdir(self):
|
|
if os.path.isdir(self.path_photos_from_camera):
|
|
pass
|
|
else:
|
|
os.makedirs(self.path_photos_from_camera, exist_ok=True)
|
|
|
|
# 删除所有数据
|
|
def pre_work_del_old_face_folders(self):
|
|
folders_rd = os.listdir(self.path_photos_from_camera)
|
|
for folder in folders_rd:
|
|
shutil.rmtree(os.path.join(self.path_photos_from_camera, folder))
|
|
if os.path.isfile("data/features_all.csv"):
|
|
os.remove("data/features_all.csv")
|
|
|
|
# 统计已经录入的人脸数据
|
|
def check_existing_faces_cnt(self):
|
|
if os.listdir(self.path_photos_from_camera):
|
|
person_list = os.listdir(self.path_photos_from_camera)
|
|
# 统计数量
|
|
self.existing_faces_cnt = len(person_list)
|
|
else:
|
|
self.existing_faces_cnt = 0
|
|
|
|
# 计算帧数
|
|
def update_fps(self):
|
|
now = time.time()
|
|
if str(self.start_time).split(".")[0] != str(now).split(".")[0]:
|
|
self.fps_show = self.fps
|
|
self.start_time = now
|
|
self.frame_time = now - self.frame_start_time
|
|
self.fps = 1.0 / self.frame_time
|
|
self.frame_start_time = now
|
|
|
|
# 添加图像说明
|
|
def draw_note(self, img_rd):
|
|
# 添加说明
|
|
cv2.putText(img_rd, "Face Register", (20, 40), self.font, 1, (255, 255, 255), 1, cv2.LINE_AA)
|
|
cv2.putText(img_rd, "FPS: " + str(self.fps_show.__round__(2)), (20, 100), self.font, 0.8, (0, 255, 0), 1,
|
|
cv2.LINE_AA)
|
|
cv2.putText(img_rd, "Faces: " + str(self.current_frame_faces_cnt), (20, 140), self.font, 0.8, (0, 255, 0), 1,
|
|
cv2.LINE_AA)
|
|
|
|
# 修改：显示当前录入人姓名
|
|
if self.current_person_name:
|
|
cv2.putText(img_rd, f"Name: {self.current_person_name}", (20, 180), self.font, 0.8, (255, 255, 0), 1,
|
|
cv2.LINE_AA)
|
|
|
|
cv2.putText(img_rd, "N: Input Name & Create folder", (20, 320), self.font, 0.8, (255, 255, 255), 1, cv2.LINE_AA)
|
|
cv2.putText(img_rd, "S: Save current face", (20, 350), self.font, 0.8, (255, 255, 255), 1, cv2.LINE_AA)
|
|
cv2.putText(img_rd, "Q: Quit", (20, 380), self.font, 0.8, (255, 255, 255), 1, cv2.LINE_AA)
|
|
cv2.putText(img_rd, "D: Delete all data", (20, 410), self.font, 0.8, (255, 255, 255), 1, cv2.LINE_AA)
|
|
|
|
# 用户输入姓名
|
|
def get_person_name_from_input(self):
|
|
print("\n请输入姓名（中文或英文）:")
|
|
name = input().strip()
|
|
return name if name else f"person_{self.existing_faces_cnt + 1}"
|
|
|
|
# 创建对应的文件夹保存数据
|
|
def create_person_folder(self, person_name):
|
|
# 清理文件名中的非法字符
|
|
safe_name = "".join(c for c in person_name if c.isalnum() or c in (' ', '-', '_')).rstrip()
|
|
if not safe_name:
|
|
safe_name = f"person_{self.existing_faces_cnt + 1}"
|
|
|
|
# 创建文件夹路径
|
|
folder_name = f"person_{safe_name}"
|
|
current_face_dir = os.path.join(self.path_photos_from_camera, folder_name)
|
|
|
|
os.makedirs(current_face_dir, exist_ok=True)
|
|
logging.info("新建人脸文件夹: %s", current_face_dir)
|
|
|
|
return current_face_dir, safe_name
|
|
|
|
# 人脸录入过程
|
|
def get_face_process(self, stream):
|
|
# 1. 新建储存人脸图像文件目录
|
|
self.pre_work_mkdir()
|
|
|
|
# 2. 检查已有人脸文件
|
|
self.check_existing_faces_cnt()
|
|
|
|
current_face_dir = ""
|
|
|
|
while stream.isOpened():
|
|
flag, img_rd = stream.read()
|
|
if not flag:
|
|
break
|
|
|
|
kk = cv2.waitKey(1)
|
|
faces = detector(img_rd, 0)
|
|
|
|
# 4. 按下 'n' 输入姓名并新建文件夹
|
|
if kk == ord('n'):
|
|
person_name = self.get_person_name_from_input()
|
|
current_face_dir, self.current_person_name = self.create_person_folder(person_name)
|
|
self.ss_cnt = 0
|
|
self.press_n_flag = 1
|
|
print(f"已创建文件夹: {current_face_dir}")
|
|
print("请调整位置并按 'S' 保存人脸")
|
|
|
|
# 5. 按下 'd' 删除所有数据
|
|
elif kk == ord('d'):
|
|
confirm = input("确定要删除所有数据吗？(y/n): ")
|
|
if confirm.lower() == 'y':
|
|
self.pre_work_del_old_face_folders()
|
|
self.existing_faces_cnt = 0
|
|
self.current_person_name = ""
|
|
self.press_n_flag = 0
|
|
print("所有数据已删除")
|
|
continue
|
|
|
|
# 6. 检测到人脸
|
|
if len(faces) != 0:
|
|
for k, d in enumerate(faces):
|
|
# 计算矩形框大小
|
|
height = d.bottom() - d.top()
|
|
width = d.right() - d.left()
|
|
hh = int(height / 2)
|
|
ww = int(width / 2)
|
|
|
|
# 判断人脸是否在范围内
|
|
if (d.right() + ww > 640 or d.bottom() + hh > 480 or
|
|
d.left() - ww < 0 or d.top() - hh < 0):
|
|
cv2.putText(img_rd, "OUT OF RANGE", (20, 300), self.font, 0.8, (0, 0, 255), 1, cv2.LINE_AA)
|
|
color_rectangle = (0, 0, 255)
|
|
save_flag = 0
|
|
else:
|
|
color_rectangle = (255, 255, 255)
|
|
save_flag = 1
|
|
|
|
# 绘制人脸框
|
|
cv2.rectangle(img_rd,
|
|
(d.left() - ww, d.top() - hh),
|
|
(d.right() + ww, d.bottom() + hh),
|
|
color_rectangle, 2)
|
|
|
|
# 创建空白图像用于保存人脸
|
|
img_blank = np.zeros((height * 2, width * 2, 3), np.uint8)
|
|
|
|
if save_flag and kk == ord('s'):
|
|
# 检查是否已创建文件夹
|
|
if self.press_n_flag:
|
|
self.ss_cnt += 1
|
|
# 提取人脸区域
|
|
for ii in range(height * 2):
|
|
for jj in range(width * 2):
|
|
img_blank[ii][jj] = img_rd[d.top() - hh + ii][d.left() - ww + jj]
|
|
|
|
# 保存人脸图像
|
|
filename = f"img_face_{self.ss_cnt}.jpg"
|
|
filepath = os.path.join(current_face_dir, filename)
|
|
cv2.imwrite(filepath, img_blank)
|
|
|
|
logging.info("保存人脸: %s", filepath)
|
|
print(f"已保存第 {self.ss_cnt} 张人脸图片")
|
|
else:
|
|
logging.warning("请先按 'N' 输入姓名创建文件夹")
|
|
|
|
self.current_frame_faces_cnt = len(faces)
|
|
|
|
# 绘制说明文字
|
|
self.draw_note(img_rd)
|
|
|
|
# 按下 'q' 退出
|
|
if kk == ord('q'):
|
|
break
|
|
|
|
# 更新 FPS
|
|
self.update_fps()
|
|
|
|
cv2.imshow("Face Register", img_rd)
|
|
|
|
def run(self):
|
|
cap = cv2.VideoCapture(0)
|
|
if not cap.isOpened():
|
|
print("错误: 无法打开摄像头")
|
|
return
|
|
|
|
# 设置摄像头参数
|
|
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
|
|
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
|
|
|
|
self.get_face_process(cap)
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
print("程序结束")
|
|
|
|
|
|
def main():
|
|
logging.basicConfig(level=logging.INFO)
|
|
Face_Register_con = Face_Register()
|
|
Face_Register_con.run()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main() |