robot-project/face_recognition.py

#!/usr/bin/env python3
"""
实时人脸识别系统

功能：
1. 添加人脸到数据库并设定名字
2. 保存和加载人脸数据库
3. 摄像头实时人脸检测和识别
4. 在检测到的人脸上显示对应名字
"""

import os
import time
import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image
import pickle
# 导入FreeType库
import freetype
import sys
import traceback


class FaceRecognitionSystem:
    """实时人脸识别系统类"""
    
    # 关键：创建 ONNX Runtime 配置，禁用融合卷积
    session_options = ort.SessionOptions()
    # 禁用融合卷积（解决 FusedConv 算子报错）
    session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
    # 或更精细控制：只禁用融合卷积，保留其他优化
    session_options.add_session_config_entry("ep.cuda.disable_fused_conv", "1")

    def __init__(self, db_path="face_database.pkl",
                model_name="buffalo_l",
                force_cuda=True):
        """
        初始化人脸识别系统
        
        Args:
            db_path: 人脸数据库保存路径
            model_name: 使用的模型名称
        """
        self.db_path = db_path
        self.model_name = model_name
        
        # 初始化人脸分析应用
        print("初始化人脸识别模型...")
        print("尝试使用TensorRT加速...")
        
        # TensorRT执行提供程序的配置选项
        trt_options = {
            'trt_engine_cache_enable': True,  # 启用引擎缓存
            'trt_engine_cache_path': './trt_cache',  # 引擎缓存路径
            'trt_fp16_enable': True,  # 启用FP16精度以提升性能
            'trt_max_workspace_size': 2 * 1024 * 1024 * 1024,  # 最大工作空间2GB
        }
        
        try:
            # 优先使用TensorRT，然后是CUDA，最后是CPU
            self.app = FaceAnalysis(
                name=model_name,
                modules=['detection', 'recognition'],
                providers=[
                    #('TensorrtExecutionProvider', trt_options),
                    'CUDAExecutionProvider',
                    'CPUExecutionProvider'
                ] 
            )
            self.app.prepare(ctx_id=0, det_size=(640, 640))
            print("成功使用TensorRT加速初始化")

        except BaseException as e:
            print(f"TensorRT启动失败: {e}", file=sys.stderr)
            print("尝试使用CUDA执行...")
            try:
                # 回退到CUDA
                self.app = FaceAnalysis(
                    name=model_name,
                    modules=['detection', 'recognition'],
                    providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] 
                )
                self.app.prepare(ctx_id=0, det_size=(640, 640))
                print("成功使用CUDA执行")
            except BaseException as e2:
                print(f"CUDA启动失败: {e2}", file=sys.stderr)
                # 确保self.app被正确初始化，使用CPU作为最后的备选方案
                print("尝试使用纯CPU执行...")
                try:
                    self.app = FaceAnalysis(
                        name=model_name,
                        modules=['detection', 'recognition'],
                        providers=['CPUExecutionProvider']  # 只使用CPU执行提供程序
                    )
                    self.app.prepare(ctx_id=-1, det_size=(640, 640))
                    print("成功使用纯CPU执行")
                except BaseException as e3:
                    print(f"纯CPU执行也失败: {e3}", file=sys.stderr)
                    raise RuntimeError(f"无法初始化FaceAnalysis: {e3}") from e2
        # 初始化人脸数据库
        self.face_database = {}
        self.load_database()
        
        # 初始化FreeType字体
        self.font = None
        self._init_font()
        
        print(f"系统初始化完成。已加载 {len(self.face_database)} 张人脸。")
    
    def _reinit_app_with_fallback(self, use_tensorrt=True):
        """
        重新初始化FaceAnalysis应用，支持TensorRT/CUDA/CPU回退
        
        Args:
            use_tensorrt: 是否尝试使用TensorRT
        
        Returns:
            bool: 初始化成功返回True，失败返回False
        """
        # TensorRT执行提供程序的配置选项
        trt_options = {
            'trt_engine_cache_enable': True,
            'trt_engine_cache_path': './trt_cache',
            'trt_fp16_enable': True,
            'trt_max_workspace_size': 2 * 1024 * 1024 * 1024,
        }
        
        if use_tensorrt:
            try:
                self.app = FaceAnalysis(
                    name=self.model_name,
                    modules=['detection', 'recognition'],
                    providers=[
                        ('TensorRTExecutionProvider', trt_options),
                        'CUDAExecutionProvider',
                        'CPUExecutionProvider'
                    ]
                )
                self.app.prepare(ctx_id=0, det_size=(640, 640))
                print("成功使用TensorRT重新初始化")
                return True
            except Exception as e:
                print(f"TensorRT重新初始化失败: {e}")
                
                traceback.print_exc()
        
        # 回退到CUDA
        try:
            self.app = FaceAnalysis(
                name=self.model_name,
                modules=['detection', 'recognition'],
                providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
            )
            self.app.prepare(ctx_id=0, det_size=(640, 640))
            print("成功使用CUDA重新初始化")
            return True
        except Exception as e:
            print(f"CUDA重新初始化失败: {e}")
            
            traceback.print_exc()
        
        # 最后回退到CPU
        try:
            self.app = FaceAnalysis(
                name=self.model_name,
                modules=['detection', 'recognition'],
                providers=['CPUExecutionProvider']
            )
            self.app.prepare(ctx_id=-1, det_size=(640, 640))
            print("成功使用CPU重新初始化")
            return True
        except Exception as e:
            print(f"CPU重新初始化也失败: {e}")
            return False
    
    def _init_font(self):
        """
        初始化FreeType字体
        """
        # Linux系统的中文字体路径
        font_paths = [
            "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",  # 文泉驿微米黑
            "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",     # 文泉驿正黑
            "/usr/share/fonts/truetype/simhei.ttf",             # 黑体
            "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",  # 备用字体
        ]
        
        # 尝试加载字体
        for font_path in font_paths:
            if os.path.exists(font_path):
                try:
                    self.font = freetype.Face(font_path)
                    print(f"成功加载字体: {font_path}")
                    return
                except Exception as e:
                    print(f"加载字体 {font_path} 失败: {e}")
        
        print("警告：无法加载任何字体，将使用OpenCV默认字体")
        self.font = None
    
    def _put_text_freetype(self, image, text, pos, font_size=18, color=(0, 255, 0)):
        """
        使用FreeType绘制文字（支持中文）
        
        Args:
            image: 要绘制文字的图像
            text: 要绘制的文字
            pos: 文字位置 (x, y)
            font_size: 字体大小
            color: 文字颜色 (B, G, R)
        
        Returns:
            image: 绘制了文字的图像
        """
        
        try:
            # 设置字体大小
            self.font.set_char_size(font_size * 64)
            
            # 分离颜色通道
            b, g, r = color
            
            # 获取文字位置
            x, y = pos
            
            # 绘制文字
            for char in text:
                # 加载字符
                self.font.load_char(char)
                
                # 获取字符的位图
                bitmap = self.font.glyph.bitmap
                width = bitmap.width
                height = bitmap.rows
                
                # 获取字符的偏移量
                left = self.font.glyph.bitmap_left
                top = self.font.glyph.bitmap_top
                
                # 计算字符在图像中的位置
                char_x = x + left
                char_y = y - top
                
                # 确保字符在图像范围内
                if char_x < 0 or char_y < 0 or char_x + width > image.shape[1] or char_y + height > image.shape[0]:
                    # 跳过超出图像范围的字符
                    x += (self.font.glyph.advance.x >> 6)
                    continue
                
                # 绘制字符的每个像素
                for i in range(height):
                    for j in range(width):
                        if i < bitmap.rows and j < bitmap.width:
                            pixel = bitmap.buffer[i * bitmap.pitch + j]
                            if pixel > 0:
                                # 设置像素颜色
                                image[char_y + i, char_x + j] = (b, g, r)
                
                # 更新x坐标，准备绘制下一个字符
                x += (self.font.glyph.advance.x >> 6)
        except Exception as e:
            print(f"绘制文字时出错: {e}")
            # 回退到OpenCV默认字体
            cv2.putText(image, text, pos, cv2.FONT_HERSHEY_SIMPLEX, font_size/25, color, 2)
        
        return image
    
    def add_face(self, image_path, name):
        """
        添加人脸到数据库
        
        Args:
            image_path: 包含人脸的图像路径
            name: 人脸对应的名字
        
        Returns:
            bool: 添加成功返回True，失败返回False
        """
        print(f"正在添加人脸: {name}...")
        
        # 读取图像
        img = cv2.imread(image_path)
        if img is None:
            print(f"错误：无法读取图像 {image_path}")
            return False
        
        try:
            # 检测人脸
            faces = self.app.get(img)
            print(f"DEBUG: 检测到 {len(faces)} 张人脸")
            
            if len(faces) == 0:
                # 尝试调整图像大小后再次检测
                print("尝试调整图像大小后再次检测...")
                h, w = img.shape[:2]
                max_size = 1280
                if h > max_size or w > max_size:
                    scale = max_size / max(h, w)
                    new_h, new_w = int(h * scale), int(w * scale)
                    img_resized = cv2.resize(img, (new_w, new_h))
                    faces = self.app.get(img_resized)
                    print(f"DEBUG: 调整大小后检测到 {len(faces)} 张人脸")
            
            if len(faces) == 0:
                print("错误：图像中未检测到人脸")
                return False
            elif len(faces) > 1:
                print("警告：图像中检测到多张人脸，仅使用第一张")
            
            # 获取人脸特征
            face = faces[0]
            print(f"DEBUG: 人脸信息: {face.keys()}")
            if 'embedding' not in face:
                print("错误：无法提取人脸特征")
                print("DEBUG: 可用的人脸信息: {face.keys()}")
                return False
            
            # 添加到数据库
            self.face_database[name] = face['embedding']
            print(f"成功添加人脸: {name}")
            return True
        except Exception as e:
            error_msg = str(e)
            print(f"添加人脸时出错: {error_msg}")
            
            # 检查是否是执行提供程序相关错误
            is_provider_error = any(keyword in error_msg for keyword in 
                                  ['CUDA', 'cuda', 'CUDNN', 'cudnn', 'GPU', 'gpu', 
                                   'FusedConv', 'TensorRT', 'tensorrt', 'TRT', 'trt'])
            
            if is_provider_error:
                print("检测到执行提供程序相关错误，尝试重新初始化...")
                # 重新初始化应用（会尝试TensorRT -> CUDA -> CPU的回退）
                if self._reinit_app_with_fallback(use_tensorrt=True):
                    # 重新尝试添加人脸
                    return self.add_face(image_path, name)
                else:
                    print("重新初始化失败")
                    return False
            return False
    
    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 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)} 张人脸")
                return True
            except Exception as e:
                print(f"加载数据库失败: {e}")
                self.face_database = {}
        else:
            print(f"未找到数据库文件 {self.db_path}，将创建新数据库")
            self.face_database = {}
        return False
    
    def recognize_face(self, face, threshold=0.70):
        """
        识别人脸，返回最匹配的名字和相似度分数
        
        Args:
            face: 检测到的人脸对象
            threshold: 识别阈值
        
        Returns:
            tuple: (str, float) - 识别出的名字和相似度分数，未识别出返回("Unknown", 0.0)
        """
        if 'embedding' not in face:
            return "Unknown", 0.0
        
        embedding = face['embedding']
        best_name = "Unknown"
        best_score = threshold
        
        # 检查数据库是否为空
        if not self.face_database:
            return "Unknown", 0.0
        
        # 遍历数据库，寻找最匹配的人脸
        for name, db_embedding in self.face_database.items():
            # 计算余弦相似度
            similarity = np.dot(embedding, db_embedding)
            
            if similarity > best_score:
                best_score = similarity
                best_name = name
        
        # 确保返回的名字是有效的字符串
        if not isinstance(best_name, str):
            best_name = str(best_name)
        # 移除可能的特殊字符
        best_name = ''.join(c for c in best_name if c.isprintable())
        if not best_name:
            best_name = "Unknown"
        return best_name, best_score
    
    def run_camera(self, camera_id, save_db_on_exit=True):
        """
        运行摄像头实时人脸识别
        
        Args:
            camera_id: 摄像头ID，默认10
            save_db_on_exit: 退出时是否保存数据库
        """
        print(f"打开摄像头 {camera_id}...")
        
        # 尝试不同的摄像头打开方式
        cap = None
        
        # 方式1：直接使用摄像头ID
        try:
            cap = cv2.VideoCapture(camera_id)
            if cap.isOpened():
                print(f"成功打开摄像头 {camera_id}（直接方式）")
            else:
                print(f"直接方式打开摄像头 {camera_id}失败，尝试其他方式...")
                cap.release()
                cap = None
        except Exception as e:
            print(f"直接方式打开摄像头失败：{e}")
            cap = None
        
        # 方式3：使用cv2.CAP_V4L2（Linux特有）
        if cap is None and os.name == 'posix':
            try:
                cap = cv2.VideoCapture(camera_id, cv2.CAP_V4L2)
                if cap.isOpened():
                    print(f"成功打开摄像头 {camera_id}（CAP_V4L2方式）")
                else:
                    print(f"CAP_V4L2方式打开摄像头 {camera_id}失败，尝试其他方式...")
                    cap.release()
                    cap = None
            except Exception as e:
                print(f"CAP_V4L2方式打开摄像头失败：{e}")
                cap = None
        
        if cap is None:
            print(f"错误：无法打开摄像头 {camera_id}")
            # 列出可用摄像头
            print("尝试检测可用摄像头...")
            for i in range(3):
                test_cap = cv2.VideoCapture(i)
                if test_cap.isOpened():
                    print(f"摄像头 {i} 可用")
                    test_cap.release()
            return False
        
        # 设置摄像头参数
        print("设置摄像头参数...")
        
        # 设置分辨率和帧率
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        cap.set(cv2.CAP_PROP_FPS, 30)
        
        # 尝试设置MJPG编码，提高兼容性（忽略可能的错误）
        try:
            cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
            print("  使用MJPG编码")
        except Exception as e:
            print(f"  使用默认编码（MJPG设置失败：{e}")
        
        # 摄像头初始化后添加延迟，让摄像头稳定
        print("摄像头初始化完成，等待1秒让摄像头稳定...")
        time.sleep(1)
        
        # 预读取几帧，丢弃初始的不稳定帧
        for _ in range(5):
            ret, _ = cap.read()
            if not ret:
                # 短暂休眠后重试
                time.sleep(0.1)
        
        print("开始实时人脸识别。按 'q' 键退出，按 's' 键保存数据库。")
        
        frame_count = 0
        consecutive_failures = 0  # 连续失败计数器
        
        while True:
            # 读取帧
            ret, frame = cap.read()
            
            if not ret:
                consecutive_failures += 1
                print(f"警告：无法读取帧 {frame_count} (连续失败: {consecutive_failures})")
                
                # 仅等待后重试，不重新打开摄像头，确保摄像头持续开启
                # 连续失败多次后才考虑重新初始化
                if consecutive_failures >= 30:  # 连续失败30次后才重新初始化
                    print("连续读取失败30次，尝试重置摄像头...")
                    consecutive_failures = 0  # 重置失败计数器
                    
                    # 释放并重新打开摄像头
                    cap.release()
                    time.sleep(1)  # 等待1秒
                    cap = cv2.VideoCapture(camera_id)
                    if not cap.isOpened():
                        print(f"错误：无法重新打开摄像头 {camera_id}")
                        break
                    # 重新设置参数
                    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
                    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
                    cap.set(cv2.CAP_PROP_FPS, 30)
                else:
                    time.sleep(0.1)  # 短暂等待后重试
                continue
            
            # 读取成功，重置失败计数器
            consecutive_failures = 0
            frame_count += 1
            
            # 检查帧是否有效
            if frame is None or frame.size == 0:
                print(f"警告：无效帧 {frame_count}")
                continue
            
            # 创建空白结果帧，防止原帧有问题
            result_frame = np.zeros((480, 640, 3), dtype=np.uint8)
            
            try:
                # 检测人脸
                faces = self.app.get(frame)
                
                # 绘制结果
                result_frame = frame.copy()
                
                for face in faces:
                    # 获取置信度
                    confidence = face['det_score']
                    
                    # 识别人脸，获取名字和相似度
                    name, similarity = self.recognize_face(face)
                    
                    # 只有当置信度大于0.7且相似度大于200时才显示对应名字，否则显示Unknown
                    if confidence > 0.7 and similarity > 100:
                        display_name = name
                    else:
                        display_name = "Unknown"
                    
                    # 获取边界框
                    bbox = face['bbox']
                    bbox = list(map(int, bbox))
                    
                    # 绘制边界框
                    cv2.rectangle(result_frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
                    
                    # 绘制名字 - 确保文字在图像范围内
                    name_y = max(20, bbox[1] - 10)  # 确保文字不会超出图像顶部
                    name_y = min(name_y, result_frame.shape[0] - 20)  # 确保文字不会超出图像底部
                    self._put_text_freetype(result_frame, display_name, (bbox[0], name_y), 
                                          font_size=18, color=(0, 255, 0))
                    
                    # 绘制置信度和相似度
                    conf_y = bbox[3] + 20
                    conf_y = min(conf_y, result_frame.shape[0] - 10)  # 确保文字不会超出图像底部
                    self._put_text_freetype(result_frame, f"置信度: {confidence:.2f}, 相似度: {similarity:.2f}", 
                                          (bbox[0], conf_y), font_size=12, color=(0, 255, 0))
            except Exception as e:
                error_msg = str(e)
                print(f"处理帧时出错: {error_msg}")
                
                # 检查是否是执行提供程序相关错误
                is_provider_error = any(keyword in error_msg for keyword in 
                                      ['CUDA', 'cuda', 'CUDNN', 'cudnn', 'GPU', 'gpu', 
                                       'FusedConv', 'TensorRT', 'tensorrt', 'TRT', 'trt'])
                
                if is_provider_error:
                    print("检测到执行提供程序相关错误，尝试重新初始化...")
                    # 重新初始化应用（会尝试TensorRT -> CUDA -> CPU的回退）
                    if self._reinit_app_with_fallback(use_tensorrt=True):
                        try:
                            # 尝试重新检测
                            faces = self.app.get(frame)
                            result_frame = frame.copy()
                            for face in faces:
                                confidence = face['det_score']
                                name, similarity = self.recognize_face(face)
                                if confidence > 0.7 and similarity > 100:
                                    display_name = name
                                else:
                                    display_name = "Unknown"
                                bbox = list(map(int, face['bbox']))
                                cv2.rectangle(result_frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
                                name_y = max(20, bbox[1] - 10)
                                name_y = min(name_y, result_frame.shape[0] - 20)
                                self._put_text_freetype(result_frame, display_name, (bbox[0], name_y), 
                                                      font_size=18, color=(0, 255, 0))
                                conf_y = bbox[3] + 20
                                conf_y = min(conf_y, result_frame.shape[0] - 10)
                                self._put_text_freetype(result_frame, f"置信度: {confidence:.2f}, 相似度: {similarity:.2f}", 
                                                      (bbox[0], conf_y), font_size=12, color=(0, 255, 0))
                        except Exception as retry_e:
                            print(f"重新检测失败: {retry_e}")
                            result_frame = np.zeros((480, 640, 3), dtype=np.uint8)
                            self._put_text_freetype(result_frame, f"重新检测失败: {retry_e}", (10, 50), 
                                                  font_size=14, color=(0, 0, 255))
                    else:
                        result_frame = np.zeros((480, 640, 3), dtype=np.uint8)
                        self._put_text_freetype(result_frame, "执行提供程序重新初始化失败", (10, 50), 
                                              font_size=14, color=(0, 0, 255))
                else:
                    # 使用空白帧，避免程序崩溃
                    result_frame = np.zeros((480, 640, 3), dtype=np.uint8)
                    self._put_text_freetype(result_frame, f"处理错误: {error_msg}", (10, 50), 
                                          font_size=14, color=(0, 0, 255))
            
            # 显示 FPS
            fps = 1.0 / (frame_count + 1e-6)  # 简单FPS计算
            self._put_text_freetype(result_frame, f"FPS: {fps:.1f}", (10, 30), 
                                  font_size=22, color=(0, 0, 255))
            
            # 确保窗口持续显示的关键逻辑
            try:
                # 确保窗口已创建并保持打开
                if not hasattr(self, '_window_created') or not self._window_created:
                    cv2.namedWindow('实时人脸识别', cv2.WINDOW_NORMAL)
                    cv2.resizeWindow('实时人脸识别', 640, 480)
                    self._window_created = True
                    print("摄像头窗口已创建并持续显示")
                
                # 持续显示结果
                cv2.imshow('实时人脸识别', result_frame)
                
                # 键盘控制 - 使用waitKey(1)确保窗口持续刷新
                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    # 退出
                    print("收到退出命令，正在关闭...")
                    break
                elif key == ord('s'):
                    # 保存数据库
                    self.save_database()
            except Exception as e:
                print(f"窗口显示错误: {e}")
                # 窗口显示错误时重新创建窗口
                if hasattr(self, '_window_created'):
                    delattr(self, '_window_created')
                # 不退出，继续尝试显示
                continue
        
        # 释放资源
        print("释放摄像头资源...")
        cap.release()
        print("关闭所有窗口...")
        cv2.destroyAllWindows()
        
        # 退出时保存数据库
        if save_db_on_exit:
            self.save_database()
        
        print("系统已退出")
        return True
    
    def add_face_interactive(self, name, camera_id):
        """
        交互式添加人脸（从摄像头捕获）
        
        Args:
            name: 人脸对应的名字
            camera_id: 摄像头ID，默认10
        
        Returns:
            bool: 添加成功返回True，失败返回False
        """
        print(f"准备添加人脸: {name}")
        print("请面对摄像头，按 'c' 键捕获图像，按 'q' 键取消")
        
        # 尝试打开摄像头
        cap = cv2.VideoCapture(camera_id)
        if not cap.isOpened():
            print(f"❌ 无法打开摄像头，索引: {camera_id}")
            
            # 列出可用摄像头
            print("ℹ️  尝试检测可用摄像头...")
            available_cameras = []
            for i in range(5):  # 尝试检测前5个摄像头
                test_cap = cv2.VideoCapture(i)
                if test_cap.isOpened():
                    available_cameras.append(i)
                    test_cap.release()
            
            if available_cameras:
                print(f"✅ 检测到以下可用摄像头: {available_cameras}")
            else:
                print("❌ 未检测到任何可用摄像头")
            
            return False
        
        # 重新设置参数
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        cap.set(cv2.CAP_PROP_FPS, 30)
        
        # 摄像头初始化后添加延迟，让摄像头稳定
        print("摄像头初始化完成，等待1秒让摄像头稳定...")
        time.sleep(1)
        
        # 预读取几帧，丢弃初始的不稳定帧
        for _ in range(5):
            ret, _ = cap.read()
            if not ret:
                # 短暂休眠后重试
                time.sleep(0.1)
        
        # 创建空白帧，防止摄像头初始帧为黑
        blank_frame = np.zeros((480, 640, 3), dtype=np.uint8)
        self._put_text_freetype(blank_frame, "初始化摄像头...", (10, 30), 
                              font_size=18, color=(0, 0, 255))
        
        # 创建窗口
        cv2.namedWindow('捕获人脸', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('捕获人脸', 640, 480)
        
        frame_count = 0
        consecutive_failures = 0
        
        while True:
            ret, frame = cap.read()
            
            if not ret or frame is None or frame.size == 0:
                consecutive_failures += 1
                print(f"⚠️  无法读取帧 ({frame_count}/{consecutive_failures})")
                
                # 使用空白帧
                display_frame = blank_frame.copy()
                self._put_text_freetype(display_frame, f"无法读取帧，尝试次数: {consecutive_failures}", (10, 60), 
                                      font_size=14, color=(0, 0, 255))
                
                if consecutive_failures > 10:
                    print("❌ 连续10次无法读取帧，退出")
                    cap.release()
                    cv2.destroyAllWindows()
                    return False
            else:
                consecutive_failures = 0
                frame_count += 1
                
                # 确保帧尺寸正确
                if frame.shape[0] != 480 or frame.shape[1] != 640:
                    display_frame = cv2.resize(frame, (640, 480))
                else:
                    display_frame = frame.copy()
            
            # 在图像上显示提示信息
            self._put_text_freetype(display_frame, "按 'c' 键捕获，按 'q' 键取消", (10, 30), 
                                  font_size=18, color=(0, 0, 255))
            
            # 显示当前帧计数
            self._put_text_freetype(display_frame, f"帧计数: {frame_count}", (10, 450), 
                                  font_size=14, color=(0, 255, 0))
            
            # 显示图像
            cv2.imshow('捕获人脸', display_frame)
            
            # 等待键盘输入
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                # 取消
                cap.release()
                cv2.destroyAllWindows()
                print("已取消添加")
                return False
            elif key == ord('c'):
                # 捕获图像
                cap.release()
                cv2.destroyAllWindows()
                
                # 确保捕获的是有效图像
                if consecutive_failures == 0 and frame is not None and frame.size > 0:
                    # 检测并添加人脸
                    return self.add_face_from_frame(frame, name)
                else:
                    print("❌ 无法捕获有效图像")
                    return False
    
    def add_face_from_frame(self, frame, name):
        """
        从图像帧添加人脸
        
        Args:
            frame: 图像帧
            name: 人脸对应的名字
        
        Returns:
            bool: 添加成功返回True，失败返回False
        """
        try:
            # 检测人脸
            faces = self.app.get(frame)
            if len(faces) == 0:
                print("错误：图像中未检测到人脸")
                return False
            elif len(faces) > 1:
                print("警告：图像中检测到多张人脸，仅使用第一张")
            
            # 获取人脸特征
            face = faces[0]
            if 'embedding' not in face:
                print("错误：无法提取人脸特征")
                return False
            
            # 添加到数据库
            self.face_database[name] = face['embedding']
            print(f"成功添加人脸: {name}")
            return True
        except Exception as e:
            error_msg = str(e)
            print(f"从帧添加人脸时出错: {error_msg}")
            
            # 检查是否是执行提供程序相关错误
            is_provider_error = any(keyword in error_msg for keyword in 
                                  ['CUDA', 'cuda', 'CUDNN', 'cudnn', 'GPU', 'gpu', 
                                   'FusedConv', 'TensorRT', 'tensorrt', 'TRT', 'trt'])
            
            if is_provider_error:
                print("检测到执行提供程序相关错误，尝试重新初始化...")
                # 重新初始化应用（会尝试TensorRT -> CUDA -> CPU的回退）
                if self._reinit_app_with_fallback(use_tensorrt=True):
                    # 重新尝试添加人脸
                    return self.add_face_from_frame(frame, name)
                else:
                    print("重新初始化失败")
                    return False
            return False


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="实时人脸识别系统")
    parser.add_argument("--add-face", nargs=2, metavar=("IMAGE_PATH", "NAME"), 
                       help="添加人脸到数据库")
    parser.add_argument("--add-webcam", metavar="NAME", 
                       help="从摄像头添加人脸到数据库")
    parser.add_argument("--db-path", default="face_database.pkl", 
                       help="人脸数据库路径")
    parser.add_argument("--camera-id", type=int, default=10, 
                       help="摄像头ID")
    parser.add_argument("--model", default="buffalo_l", 
                       help="使用的模型名称")
    parser.add_argument("--force-cuda", action="store_true", default=True, 
                       help="是否强制使用CUDA，默认强制使用")
    parser.add_argument("--allow-cpu-fallback", action="store_false", dest="force_cuda", 
                       help="允许在CUDA不可用时回退到CPU")
    
    args = parser.parse_args()
    
    print(f"命令行参数: force_cuda={args.force_cuda}")
    
    # 初始化系统 - 强制使用CUDA
    system = FaceRecognitionSystem(db_path=args.db_path, 
                                  model_name=args.model, 
                                  force_cuda=args.force_cuda)
    
    # 添加人脸
    if args.add_face:
        image_path, name = args.add_face
        system.add_face(image_path, name)
        system.save_database()
    
    # 从摄像头添加人脸
    elif args.add_webcam:
        system.add_face_interactive(args.add_webcam, args.camera_id)
        system.save_database()
    
    # 运行摄像头
    else:
        system.run_camera(camera_id=args.camera_id)


if __name__ == "__main__":
    main()