robot-project/deepface/deepface_demo.py

#!/usr/bin/env python3
"""
DeepFace人脸识别Demo
功能：
1. 摄像头实时人脸识别（输出名字、年龄、性别、情绪）
2. 照片人脸识别
3. 人脸注册到数据库
"""

import cv2
import numpy as np
import os
import time
import argparse
from deepface import DeepFace
import pickle
from pathlib import Path


class DeepFaceDemo:
    """DeepFace人脸识别Demo类"""
    
    def __init__(self, db_path="face_database.pkl", detector_backend="opencv"):
        """
        初始化DeepFace Demo
        
        Args:
            db_path: 人脸数据库保存路径
            detector_backend: 人脸检测后端（opencv, ssd, mtcnn, retinaface等）
        """
        self.db_path = db_path
        self.detector_backend = detector_backend
        self.face_database = {}
        self.load_database()
        
        print(f"DeepFace Demo初始化完成")
        print(f"检测后端: {detector_backend}")
        print(f"数据库路径: {db_path}")
        print(f"已加载 {len(self.face_database)} 张人脸")
    
    def load_database(self):
        """从文件加载人脸数据库"""
        if os.path.exists(self.db_path):
            try:
                with open(self.db_path, 'rb') as f:
                    self.face_database = pickle.load(f)
                print(f"从 {self.db_path} 加载了 {len(self.face_database)} 张人脸")
            except Exception as e:
                print(f"加载数据库失败: {e}")
                self.face_database = {}
        else:
            print(f"未找到数据库文件 {self.db_path}，将创建新数据库")
            self.face_database = {}
    
    def save_database(self):
        """保存人脸数据库到文件"""
        try:
            with open(self.db_path, 'wb') as f:
                pickle.dump(self.face_database, f)
            print(f"人脸数据库已保存到 {self.db_path}")
            return True
        except Exception as e:
            print(f"保存数据库失败: {e}")
            return False
    
    def analyze_face(self, image):
        """
        分析人脸属性（年龄、性别、情绪、种族）
        
        Args:
            image: 图像（numpy数组或文件路径）
        
        Returns:
            dict: 包含年龄、性别、情绪等信息的字典
        """
        try:
            result = DeepFace.analyze(
                image,
                actions=['age', 'gender', 'emotion', 'race'],
                detector_backend=self.detector_backend,
                enforce_detection=False
            )
            return result
        except Exception as e:
            print(f"分析人脸时出错: {e}")
            return None
    
    def recognize_face(self, image, threshold=0.4):
        """
        识别人脸，返回最匹配的名字
        
        Args:
            image: 图像（numpy数组或文件路径）
            threshold: 识别阈值
        
        Returns:
            tuple: (name, confidence) - 识别出的名字和置信度
        """
        if not self.face_database:
            return "Unknown", 0.0
        
        try:
            # 在数据库中查找匹配的人脸
            dfs = DeepFace.find(
                image,
                db_path="database",
                detector_backend=self.detector_backend,
                enforce_detection=False
            )
            
            if dfs is not None and len(dfs) > 0 and not dfs[0].empty:
                # 获取最佳匹配
                best_match = dfs[0].iloc[0]
                confidence = best_match['distance']
                
                # 从文件名中提取名字
                identity = best_match['identity']
                name = os.path.basename(os.path.dirname(identity))
                
                # 如果距离小于阈值，则认为是匹配
                if confidence < threshold:
                    return name, confidence
            
            return "Unknown", 0.0
            
        except Exception as e:
            print(f"识别人脸时出错: {e}")
            return "Unknown", 0.0
    
    def register_face(self, image_path, name):
        """
        注册人脸到数据库
        
        Args:
            image_path: 人脸图像路径
            name: 人脸名字
        
        Returns:
            bool: 注册成功返回True
        """
        try:
            # 创建数据库目录
            db_dir = Path("database") / name
            db_dir.mkdir(parents=True, exist_ok=True)
            
            # 复制图像到数据库
            img = cv2.imread(image_path)
            if img is None:
                print(f"无法读取图像: {image_path}")
                return False
            
            # 保存图像到数据库
            save_path = db_dir / f"{name}_{int(time.time())}.jpg"
            cv2.imwrite(str(save_path), img)
            
            # 更新内存数据库
            self.face_database[name] = str(save_path)
            
            print(f"成功注册人脸: {name} -> {save_path}")
            return True
            
        except Exception as e:
            print(f"注册人脸时出错: {e}")
            return False
    
    def register_face_from_frame(self, frame, name):
        """
        从摄像头帧注册人脸
        
        Args:
            frame: 图像帧
            name: 人脸名字
        
        Returns:
            bool: 注册成功返回True
        """
        try:
            # 创建数据库目录
            db_dir = Path("database") / name
            db_dir.mkdir(parents=True, exist_ok=True)
            
            # 保存图像到数据库
            save_path = db_dir / f"{name}_{int(time.time())}.jpg"
            cv2.imwrite(str(save_path), frame)
            
            # 更新内存数据库
            self.face_database[name] = str(save_path)
            
            print(f"成功注册人脸: {name} -> {save_path}")
            return True
            
        except Exception as e:
            print(f"从帧注册人脸时出错: {e}")
            return False
    
    def process_image(self, image_path):
        """
        处理单张图像，显示人脸分析结果
        
        Args:
            image_path: 图像路径
        """
        print(f"\n处理图像: {image_path}")
        
        # 读取图像
        img = cv2.imread(image_path)
        if img is None:
            print(f"无法读取图像: {image_path}")
            return
        
        # 分析人脸
        result = self.analyze_face(img)
        
        if result is None or len(result) == 0:
            print("未检测到人脸")
            return
        
        # 处理每个检测到的人脸
        for i, face_info in enumerate(result):
            print(f"\n人脸 {i + 1}:")
            print(f"  年龄: {face_info.get('age', 'N/A')}")
            print(f"  性别: {face_info.get('dominant_gender', 'N/A')}")
            print(f"  情绪: {face_info.get('dominant_emotion', 'N/A')}")
            print(f"  种族: {face_info.get('dominant_race', 'N/A')}")
            
            # 识别人脸
            name, confidence = self.recognize_face(img)
            print(f"  识别: {name} (置信度: {confidence:.3f})")
            
            # 绘制结果
            if 'region' in face_info:
                region = face_info['region']
                x, y, w, h = region['x'], region['y'], region['w'], region['h']
                
                # 绘制人脸框
                cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
                
                # 准备显示文本
                text_lines = [
                    f"{name}",
                    f"Age: {face_info.get('age', 'N/A')}",
                    f"Gender: {face_info.get('dominant_gender', 'N/A')}",
                    f"Emotion: {face_info.get('dominant_emotion', 'N/A')}"
                ]
                
                # 绘制文本
                for j, line in enumerate(text_lines):
                    cv2.putText(img, line, (x, y - 10 - j * 20), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
        
        # 显示结果
        cv2.imshow('Face Recognition Result', img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    
    def run_camera(self, camera_id=0, enable_recognition=True):
        """
        运行摄像头实时人脸识别
        
        Args:
            camera_id: 摄像头ID
            enable_recognition: 是否启用人脸识别（否则只显示分析结果）
        """
        print(f"\n打开摄像头 {camera_id}...")
        
        # 打开摄像头
        cap = cv2.VideoCapture(camera_id)
        if not cap.isOpened():
            print(f"无法打开摄像头 {camera_id}")
            return
        
        # 设置摄像头参数
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        cap.set(cv2.CAP_PROP_FPS, 30)
        
        # 等待摄像头稳定
        time.sleep(1)
        
        print("开始实时人脸识别")
        print("按 'q' 键退出")
        print("按 's' 键保存当前帧")
        print("按 'r' 键注册当前人脸")
        
        frame_count = 0
        fps_start_time = time.time()
        
        while True:
            # 读取帧
            ret, frame = cap.read()
            if not ret:
                print("无法读取帧")
                break
            
            frame_count += 1
            
            # 计算FPS
            if frame_count % 30 == 0:
                fps = 30 / (time.time() - fps_start_time)
                fps_start_time = time.time()
                print(f"FPS: {fps:.1f}")
            
            # 分析人脸
            result = self.analyze_face(frame)
            
            if result is not None and len(result) > 0:
                for face_info in result:
                    # 获取人脸区域
                    if 'region' in face_info:
                        region = face_info['region']
                        x, y, w, h = region['x'], region['y'], region['w'], region['h']
                        
                        # 绘制人脸框
                        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                        
                        # 获取属性
                        age = face_info.get('age', 'N/A')
                        gender = face_info.get('dominant_gender', 'N/A')
                        emotion = face_info.get('dominant_emotion', 'N/A')
                        
                        # 识别人脸（如果启用）
                        if enable_recognition:
                            name, confidence = self.recognize_face(frame)
                            display_name = name
                        else:
                            display_name = "Unknown"
                        
                        # 准备显示文本
                        text_lines = [
                            f"{display_name}",
                            f"Age: {age}",
                            f"Gender: {gender}",
                            f"Emotion: {emotion}"
                        ]
                        
                        # 绘制文本
                        for j, line in enumerate(text_lines):
                            cv2.putText(frame, line, (x, y - 10 - j * 20), 
                                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
            
            # 显示帧
            cv2.imshow('DeepFace Real-time Recognition', frame)
            
            # 键盘控制
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                print("退出...")
                break
            elif key == ord('s'):
                # 保存当前帧
                save_path = f"capture_{int(time.time())}.jpg"
                cv2.imwrite(save_path, frame)
                print(f"已保存帧到: {save_path}")
            elif key == ord('r'):
                # 注册当前人脸
                name = input("请输入要注册的名字: ")
                if name:
                    if self.register_face_from_frame(frame, name):
                        self.save_database()
                        print(f"成功注册: {name}")
        
        # 释放资源
        cap.release()
        cv2.destroyAllWindows()
        
        # 保存数据库
        self.save_database()
    
    def register_from_camera(self, camera_id=0):
        """
        从摄像头注册人脸
        
        Args:
            camera_id: 摄像头ID
        """
        print(f"\n打开摄像头 {camera_id} 进行人脸注册...")
        
        # 打开摄像头
        cap = cv2.VideoCapture(camera_id)
        if not cap.isOpened():
            print(f"无法打开摄像头 {camera_id}")
            return
        
        # 设置摄像头参数
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        
        # 等待摄像头稳定
        time.sleep(1)
        
        print("按 'c' 键捕获人脸，按 'q' 键退出")
        
        while True:
            # 读取帧
            ret, frame = cap.read()
            if not ret:
                print("无法读取帧")
                break
            
            # 显示帧
            cv2.imshow('Register Face', frame)
            
            # 键盘控制
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                print("退出...")
                break
            elif key == ord('c'):
                # 捕获并注册人脸
                name = input("请输入要注册的名字: ")
                if name:
                    if self.register_face_from_frame(frame, name):
                        self.save_database()
                        print(f"成功注册: {name}")
                    else:
                        print("注册失败")
        
        # 释放资源
        cap.release()
        cv2.destroyAllWindows()


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="DeepFace人脸识别Demo")
    parser.add_argument("--camera", type=int, default=None, 
                       help="使用摄像头进行实时识别（指定摄像头ID）")
    parser.add_argument("--image", type=str, default=None, 
                       help="处理单张图像")
    parser.add_argument("--register", type=str, default=None, 
                       help="注册人脸（图像路径）")
    parser.add_argument("--register-camera", type=int, default=None, 
                       help="从摄像头注册人脸（指定摄像头ID）")
    parser.add_argument("--name", type=str, default=None, 
                       help="注册时使用的名字")
    parser.add_argument("--db-path", type=str, default="face_database.pkl", 
                       help="人脸数据库路径")
    parser.add_argument("--detector", type=str, default="opencv", 
                       choices=["opencv", "ssd", "mtcnn", "retinaface", "mediapipe"],
                       help="人脸检测后端")
    parser.add_argument("--no-recognition", action="store_true", 
                       help="禁用人脸识别（只显示分析结果）")
    
    args = parser.parse_args()
    
    # 创建DeepFace Demo实例
    demo = DeepFaceDemo(db_path=args.db_path, detector_backend=args.detector)
    
    # 根据参数执行不同操作
    if args.register:
        # 注册人脸
        if args.name is None:
            args.name = input("请输入要注册的名字: ")
        if demo.register_face(args.register, args.name):
            demo.save_database()
            print(f"成功注册: {args.name}")
        else:
            print("注册失败")
    
    elif args.register_camera is not None:
        # 从摄像头注册人脸
        if args.name is None:
            args.name = input("请输入要注册的名字: ")
        demo.register_from_camera(args.register_camera)
    
    elif args.camera is not None:
        # 运行摄像头实时识别
        enable_recognition = not args.no_recognition
        demo.run_camera(camera_id=args.camera, enable_recognition=enable_recognition)
    
    elif args.image:
        # 处理单张图像
        demo.process_image(args.image)
    
    else:
        # 默认运行摄像头
        print("未指定操作，运行摄像头实时识别...")
        demo.run_camera(camera_id=0)


if __name__ == "__main__":
    main()