diff --git a/src/video_check_biz.py b/src/video_check_biz.py new file mode 100644 index 0000000..0520cb7 --- /dev/null +++ b/src/video_check_biz.py @@ -0,0 +1,217 @@ +""" +视频检查业务类 +继承自BaseFaceBiz,专门处理视频中的人脸特征提取 +""" + +import cv2 +import numpy as np +from typing import Optional, List, Dict, Tuple +import os +from insightface.app import FaceAnalysis + +from src.base_face_biz import BaseFaceBiz + + +class VideoCheckBiz(BaseFaceBiz): + """ + 视频检查业务类 + 专门处理视频中的人脸特征提取和质量评估 + """ + + def __init__(self, face_analysis: FaceAnalysis): + """ + 初始化视频检查业务类 + + 参数: + face_analysis: 已初始化好的FaceAnalysis实例 + """ + super().__init__(face_analysis) + + def extract_best_face_from_video(self, video_path: str, frame_skip: int = 10) -> Optional[np.ndarray]: + """ + 从视频中提取最佳人脸特征 + + 参数: + video_path: 视频文件路径 + frame_skip: 跳帧数,每隔多少帧取一帧 + + 返回: + 最佳人脸的特征向量,如果未找到合适的人脸则返回None + """ + if not os.path.exists(video_path): + print(f"❌ 视频文件不存在: {video_path}") + return None + + # 打开视频文件 + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + print(f"❌ 无法打开视频文件: {video_path}") + return None + + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + fps = cap.get(cv2.CAP_PROP_FPS) + + print(f"📹 视频信息: {total_frames}帧, {fps:.1f}FPS") + print(f"🔍 跳帧设置: 每{frame_skip}帧取一帧") + + best_quality_score = -1.0 + best_feature = None + best_frame_info = None + processed_frames = 0 + valid_frames = 0 + + frame_index = 0 + while True: + ret, frame = cap.read() + if not ret: + break + + # 跳过指定帧数 + if frame_index % frame_skip != 0: + frame_index += 1 + continue + + processed_frames += 1 + + # 人脸检测 + faces = self.app.get(frame) + + # 检查人脸数量 + if len(faces) == 0: + # 无人脸,跳过 + frame_index += 1 + continue + elif len(faces) > 1: + # 超过1个人脸,舍弃该帧 + frame_index += 1 + continue + + # 只有1个人脸 + face = faces[0] + + # 检查人脸质量 + is_acceptable, quality_metrics = self.is_face_quality_acceptable(face, frame) + + if not is_acceptable: + # 质量不可接受,跳过 + frame_index += 1 + continue + + valid_frames += 1 + + # 计算综合质量得分 + quality_score = self._calculate_face_quality_score(face, quality_metrics) + + # 更新最佳人脸 + if quality_score > best_quality_score: + best_quality_score = quality_score + best_feature = face.embedding + best_frame_info = { + 'frame_index': frame_index, + 'quality_score': quality_score, + 'quality_metrics': quality_metrics + } + + frame_index += 1 + + cap.release() + + # 输出结果统计 + print(f"📊 处理统计: 共处理{processed_frames}帧, 有效帧{valid_frames}帧") + + if best_feature is not None: + print(f"✅ 找到最佳人脸: 帧{best_frame_info['frame_index']}, 质量得分: {best_quality_score:.3f}") + print(f" 检测得分: {best_frame_info['quality_metrics']['det_score']:.3f}") + print(f" 清晰度: {best_frame_info['quality_metrics']['clarity_score']:.1f}") + print(f" 姿态角度: pitch={best_frame_info['quality_metrics']['pitch']:.1f}°, yaw={best_frame_info['quality_metrics']['yaw']:.1f}°") + print(f" 人脸尺寸: {best_frame_info['quality_metrics']['bbox_width']}x{best_frame_info['quality_metrics']['bbox_height']}") + else: + print("❌ 未找到合适的人脸") + + return best_feature + + def _calculate_face_quality_score(self, face, quality_metrics: Dict) -> float: + """ + 计算人脸的综合质量得分 + + 参数: + face: 人脸检测结果 + quality_metrics: 质量指标字典 + + 返回: + 综合质量得分 (0-1) + """ + # 基础得分:检测置信度 + base_score = quality_metrics['det_score'] + + # 清晰度得分 (归一化到0-1) + clarity_score = min(quality_metrics['clarity_score'] / self.clarity_threshold, 1.0) + + # 姿态得分 (基于角度偏差) + pitch_score = 1.0 - min(abs(quality_metrics['pitch']) / self.pitch_threshold, 1.0) + yaw_score = 1.0 - min(abs(quality_metrics['yaw']) / self.yaw_threshold, 1.0) + + # 尺寸得分 (基于最小尺寸要求) + min_dim = min(quality_metrics['bbox_width'], quality_metrics['bbox_height']) + size_score = min(min_dim / self.min_face_size, 1.0) + + # 综合得分权重 + weights = { + 'base': 0.3, # 检测置信度 + 'clarity': 0.3, # 清晰度 + 'pose': 0.2, # 姿态 + 'size': 0.2 # 尺寸 + } + + # 计算综合得分 + pose_score = (pitch_score + yaw_score) / 2 + + total_score = ( + base_score * weights['base'] + + clarity_score * weights['clarity'] + + pose_score * weights['pose'] + + size_score * weights['size'] + ) + + return total_score + + def extract_best_face_from_video_with_details(self, video_path: str, frame_skip: int = 10) -> Optional[Dict]: + """ + 从视频中提取最佳人脸特征(带详细信息) + + 参数: + video_path: 视频文件路径 + frame_skip: 跳帧数 + + 返回: + 包含特征向量和详细信息的字典,如果未找到合适的人脸则返回None + """ + best_feature = self.extract_best_face_from_video(video_path, frame_skip) + + if best_feature is not None: + # 重新处理视频获取详细信息 + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + return { + 'feature': best_feature, + 'frame_info': None, + 'video_info': { + 'path': video_path, + 'frame_skip': frame_skip + } + } + + # 这里可以添加更多详细信息的提取逻辑 + cap.release() + + return { + 'feature': best_feature, + 'frame_info': { + 'frame_skip': frame_skip + }, + 'video_info': { + 'path': video_path + } + } + + return None \ No newline at end of file