49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
from ultralytics import YOLO
|
||
import cv2
|
||
from logger_config import logger
|
||
|
||
class ViolationDetector:
|
||
"""
|
||
用于加载YOLOv8 .pt模型并进行违规内容检测的类。
|
||
"""
|
||
def __init__(self, model_path):
|
||
"""
|
||
初始化检测器。
|
||
|
||
Args:
|
||
model_path (str): YOLO .pt模型的路径。
|
||
"""
|
||
logger.info(f"正在从 '{model_path}' 加载YOLO模型...")
|
||
self.model = YOLO(model_path)
|
||
logger.info("YOLO模型加载成功。")
|
||
|
||
def detect(self, frame):
|
||
"""
|
||
对单帧图像进行目标检测。
|
||
|
||
Args:
|
||
frame: 输入的图像帧 (NumPy数组, BGR格式)。
|
||
|
||
Returns:
|
||
ultralytics.engine.results.Results: YOLO的检测结果对象。
|
||
"""
|
||
# conf可以根据您的模型效果进行调整
|
||
# --- 为了测试，我们暂时将置信度调低，例如 0.2 ---
|
||
results = self.model(frame, conf=0.2)
|
||
return results[0]
|
||
|
||
def draw_boxes(self, frame, result):
|
||
"""
|
||
在图像帧上绘制检测框。
|
||
|
||
Args:
|
||
frame: 原始图像帧。
|
||
result: YOLO的检测结果对象。
|
||
|
||
Returns:
|
||
numpy.ndarray: 绘制了检测框的图像帧。
|
||
"""
|
||
# 使用YOLO自带的plot功能，方便快捷
|
||
annotated_frame = result.plot()
|
||
return annotated_frame
|