103 lines
3.1 KiB
Python
103 lines
3.1 KiB
Python
"""
|
|
视频检查业务类 - RTSP专用
|
|
专门处理RTSP视频流中的人脸识别和检测
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from typing import Optional, List, Dict, Tuple
|
|
import time
|
|
from insightface.app import FaceAnalysis
|
|
from biz.base_face_biz import BaseFaceBiz
|
|
|
|
class VideoFacePrisonBiz(BaseFaceBiz):
|
|
"""
|
|
视频检查业务类 - RTSP专用
|
|
专门处理RTSP视频流中的人脸识别和检测
|
|
"""
|
|
|
|
def __init__(self, face_analysis: FaceAnalysis):
|
|
"""
|
|
初始化视频检查业务类
|
|
|
|
参数:
|
|
face_analysis: 已初始化好的FaceAnalysis实例
|
|
"""
|
|
super().__init__(face_analysis)
|
|
|
|
def draw_detections(self, frame: np.ndarray, results: List[Dict]) -> np.ndarray:
|
|
"""
|
|
重写绘制检测结果方法
|
|
只在检测到黑名单匹配时用红色绘制人脸框
|
|
|
|
参数:
|
|
frame: 原始帧图像
|
|
results: 检测结果列表
|
|
|
|
返回:
|
|
绘制后的帧图像
|
|
"""
|
|
for result in results:
|
|
# 只在黑名单匹配时绘制
|
|
if result['is_match']:
|
|
bbox = result['bbox']
|
|
|
|
# 使用红色绘制人脸框
|
|
x1, y1, x2, y2 = bbox
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
|
|
|
|
# 添加简单的匹配信息
|
|
best_match = result['best_match']
|
|
similarity = result['similarity']
|
|
|
|
# 绘制匹配信息
|
|
text = f"{best_match}: {similarity:.3f}"
|
|
text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
|
|
|
|
# 绘制文本背景
|
|
cv2.rectangle(frame, (x1, y1 - text_size[1] - 5),
|
|
(x1 + text_size[0], y1), (0, 0, 0), -1)
|
|
|
|
# 绘制文本
|
|
cv2.putText(frame, text, (x1, y1 - 5),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
|
|
|
|
return frame
|
|
|
|
def process_frame(self, frame: np.ndarray) -> Tuple[np.ndarray, List[Dict], float]:
|
|
"""
|
|
处理单帧图像
|
|
|
|
返回:
|
|
(原始帧, 识别结果列表, 处理时间ms)
|
|
"""
|
|
start_time = time.time()
|
|
|
|
# 人脸检测和识别
|
|
faces = self.app.get(frame)
|
|
|
|
results = []
|
|
for face in faces:
|
|
# 检查人脸质量是否可接受
|
|
is_acceptable, quality_metrics = self.is_face_quality_acceptable(face, frame)
|
|
|
|
# 查找最佳匹配
|
|
best_name, similarity = self.find_best_match(face.embedding)
|
|
|
|
is_match = best_name is not None and similarity >= self.similarity_threshold
|
|
|
|
result = {
|
|
'bbox': face.bbox.astype(int).tolist(),
|
|
'similarity': similarity,
|
|
'best_match': best_name,
|
|
'is_match': is_match,
|
|
'det_score': float(face.det_score),
|
|
'quality_metrics': quality_metrics,
|
|
'is_acceptable': is_acceptable
|
|
}
|
|
results.append(result)
|
|
|
|
processing_time = (time.time() - start_time) * 1000
|
|
return frame, results, processing_time
|
|
|