zzx/video
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ocr1.0

Browse Source
This commit is contained in:
ninghongbin
2025-09-03 16:22:21 +08:00
parent b7773f5f00
commit d83923d06b
5 changed files with 211 additions and 342 deletions
Download Patch File Download Diff File Expand all files Collapse all files

113
ocr/model_violation_detector.py
View File

@ -1,35 +1,30 @@
import cv2
from logger_config import logger
from ocr_violation_detector import OCRViolationDetector
from yolo_violation_detector import ViolationDetector as YoloViolationDetector
from face_recognizer import FaceRecognizer
import os
import cv2
import yaml
from pathlib import Path
from .ocr_violation_detector import OCRViolationDetector
from .yolo_violation_detector import ViolationDetector as YoloViolationDetector
from .face_recognizer import FaceRecognizer
class MultiModelViolationDetector:
"""
多模型违规检测封装类，串行调用OCR、人脸识别和YOLO模型（调整为YOLO最后检测），任一模型检测到违规即返回结果
多模型违规检测封装类，串行调用OCR、人脸识别和YOLO模型，任一模型检测到违规即返回结果
"""
def __init__(self,
forbidden_words_path: str,
ocr_config_path: str, # 新增OCR配置文件路径参数
ocr_config_path: str,
yolo_model_path: str,
known_faces_dir: str,
ocr_confidence_threshold: float = 0.5):
"""
初始化所有检测模型
Args:
forbidden_words_path: 违禁词文件路径
ocr_config_path: OCR配置文件（1.yaml）路径
yolo_model_path: YOLO模型文件路径
known_faces_dir: 已知人脸目录路径
ocr_confidence_threshold: OCR置信度阈值
"""
# 初始化OCR检测器（传入配置文件路径）
# 初始化OCR检测器
self.ocr_detector = OCRViolationDetector(
forbidden_words_path=forbidden_words_path,
ocr_config_path=ocr_config_path, # 传递配置文件路径
ocr_config_path=ocr_config_path,
ocr_confidence_threshold=ocr_confidence_threshold
)
@ -38,22 +33,16 @@ class MultiModelViolationDetector:
known_faces_dir=known_faces_dir
)
# 初始化YOLO检测器（调整为最后初始化）
# 初始化YOLO检测器
self.yolo_detector = YoloViolationDetector(
model_path=yolo_model_path
)
logger.info("多模型违规检测器初始化完成")
print("多模型违规检测器初始化完成")
def detect_violations(self, frame):
"""
串行调用三个检测模型（OCR → 人脸识别 → YOLO），任一检测到违规即返回结果
Args:
frame: 输入视频帧 (NumPy数组, BGR格式)
Returns:
tuple: (是否有违规, 违规类型, 违规详情)
违规类型: 'ocr' | 'yolo' | 'face' | None
违规详情: 对应模型的检测结果
"""
# 1. 首先进行OCR违禁词检测
try:
@ -63,10 +52,10 @@ class MultiModelViolationDetector:
"words": ocr_words,
"confidences": ocr_confs
}
logger.warning(f"OCR检测到违禁内容: {details}")
print(f"警告: OCR检测到违禁内容: {details}")
return (True, "ocr", details)
except Exception as e:
logger.error(f"OCR检测出错: {str(e)}", exc_info=True)
print(f"错误: OCR检测出错: {str(e)}")
# 2. 接着进行人脸识别检测
try:
@ -76,58 +65,72 @@ class MultiModelViolationDetector:
"name": face_name,
"similarity": face_similarity
}
logger.warning(f"人脸识别到违规人员: {details}")
print(f"警告: 人脸识别到违规人员: {details}")
return (True, "face", details)
except Exception as e:
logger.error(f"人脸识别出错: {str(e)}", exc_info=True)
print(f"错误: 人脸识别出错: {str(e)}")
# 3. 最后进行YOLO目标检测（调整为最后检测）
# 3. 最后进行YOLO目标检测
try:
yolo_results = self.yolo_detector.detect(frame)
# 检查是否有检测结果（根据实际业务定义何为违规目标）
if len(yolo_results.boxes) > 0:
# 提取检测到的目标信息
details = {
"classes": yolo_results.names,
"boxes": yolo_results.boxes.xyxy.tolist(), # 边界框坐标
"confidences": yolo_results.boxes.conf.tolist(), # 置信度
"class_ids": yolo_results.boxes.cls.tolist() # 类别ID
"boxes": yolo_results.boxes.xyxy.tolist(),
"confidences": yolo_results.boxes.conf.tolist(),
"class_ids": yolo_results.boxes.cls.tolist()
}
logger.warning(f"YOLO检测到违规目标: {details}")
print(f"警告: YOLO检测到违规目标: {details}")
return (True, "yolo", details)
except Exception as e:
logger.error(f"YOLO检测出错: {str(e)}", exc_info=True)
print(f"错误: YOLO检测出错: {str(e)}")
# 所有检测均未发现违规
return (False, None, None)
# # 使用示例
def load_config(config_path: str) -> dict:
"""加载YAML配置文件"""
try:
with open(config_path, 'r', encoding='utf-8') as f:
return yaml.safe_load(f)
except FileNotFoundError:
print(f"错误: 配置文件未找到: {config_path}")
raise
except yaml.YAMLError as e:
print(f"错误: 配置文件格式错误: {config_path}, 错误: {str(e)}")
raise
except Exception as e:
print(f"错误: 加载配置文件出错: {str(e)}")
raise
# 使用示例
# if __name__ == "__main__":
# # 配置文件路径（根据实际情况修改）
# FORBIDDEN_WORDS_PATH = r"D:\Git\bin\video\ocr\forbidden_words.txt"
# OCR_CONFIG_PATH = r"D:\Git\bin\video\ocr\config\1.yaml" # 新增OCR配置文件路径
# YOLO_MODEL_PATH = r"D:\Git\bin\video\ocr\models\best.pt"
# KNOWN_FACES_DIR = r"D:\Git\bin\video\ocr\known_faces"
# # 加载配置文件
# config = load_config("config.yaml") # 配置文件路径，可根据实际情况修改
#
# # 初始化多模型检测器
# detector = MultiModelViolationDetector(
# forbidden_words_path=FORBIDDEN_WORDS_PATH,
# ocr_config_path=OCR_CONFIG_PATH, # 传入OCR配置文件路径
# yolo_model_path=YOLO_MODEL_PATH,
# known_faces_dir=KNOWN_FACES_DIR,
# ocr_confidence_threshold=0.5
# forbidden_words_path=config["forbidden_words_path"],
# ocr_config_path=config["ocr_config_path"],
# yolo_model_path=config["yolo_model_path"],
# known_faces_dir=config["known_faces_dir"],
# ocr_confidence_threshold=config.get("ocr_confidence_threshold", 0.5)
# )
#
# # 读取测试图像（可替换为视频帧读取逻辑）
# test_image_path = r"D:\Git\bin\video\ocr\images\img.png"
# frame = cv2.imread(test_image_path)
# test_image_path = config.get("test_image_path") # 从配置文件获取测试图片路径
# if test_image_path:
# frame = cv2.imread(test_image_path)
#
# if frame is not None:
# has_violation, violation_type, details = detector.detect_violations(frame)
# if has_violation:
# print(f"检测到违规 - 类型: {violation_type}, 详情: {details}")
# if frame is not None:
# has_violation, violation_type, details = detector.detect_violations(frame)
# if has_violation:
# print(f"检测到违规 - 类型: {violation_type}, 详情: {details}")
# else:
# print("未检测到任何违规内容")
# else:
# print("未检测到任何违规内容")
# print(f"无法读取测试图像: {test_image_path}")
# else:
# print(f"无法读取测试图像: {test_image_path}")
# print("配置文件中未指定测试图像路径")