diff --git a/rtsp_processor_1.py b/rtsp_processor_1.py index a26c3a9..42009a1 100644 --- a/rtsp_processor_1.py +++ b/rtsp_processor_1.py @@ -8,11 +8,12 @@ from pathlib import Path import subprocess import logging from typing import List, Dict, Tuple, Optional -from dataclasses import dataclass from collections import deque import yaml +from util.yolo_util import YOLODetector +from util.entity_utl import DetectionResult, AlarmConfig from util.log_util import TimeBasedDuplicateFilter @@ -35,23 +36,7 @@ logger = logging.getLogger(__name__) # 5为时间间隔,单位s. 5s内不会输出相同日志 logger.addFilter(TimeBasedDuplicateFilter(5)) -@dataclass -class DetectionResult: - """检测结果数据类""" - boxes: np.ndarray - confidences: np.ndarray - class_ids: np.ndarray - class_names: List[str] - timestamp: datetime -@dataclass -class AlarmConfig: - """告警配置""" - target_classes: List[str] # 目标类别 - confidence_threshold: float = 0.5 # 置信度阈值 - alarm_duration: int = 10 # 告警录制时长(秒) - cooldown_duration: int = 30 # 告警冷却时间(秒) - save_path: str = "./alarm_videos" # 保存路径 class FrameBuffer: """帧缓冲区,用于告警录制""" @@ -74,34 +59,7 @@ class FrameBuffer: return frames -class YOLODetector: - """YOLO检测器""" - def __init__(self, model_path: str = "yolov8n.pt"): - self.model = YOLO(model_path) - self.class_names = self.model.names - - def detect(self, frame: np.ndarray, confidence_threshold: float = 0.5) -> DetectionResult: - """执行目标检测""" - results = self.model(frame, conf=confidence_threshold, verbose=False) - - if len(results) > 0 and results[0].boxes is not None: - boxes = results[0].boxes.xyxy.cpu().numpy() - confidences = results[0].boxes.conf.cpu().numpy() - class_ids = results[0].boxes.cls.cpu().numpy().astype(int) - class_names = [self.class_names[id] for id in class_ids] - else: - boxes = np.array([]) - confidences = np.array([]) - class_ids = np.array([]) - class_names = [] - - return DetectionResult( - boxes=boxes, - confidences=confidences, - class_ids=class_ids, - class_names=class_names, - timestamp=datetime.now() - ) + class AlarmManager: """告警管理器""" @@ -425,6 +383,8 @@ class RTSPProcessor: def connect_to_rtsp_stream(self, url): """ 尝试连接到 RTSP 流 """ cap = cv2.VideoCapture(url) + # 设置缓冲区大小 + cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) # cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'HEVC')) if not cap.isOpened(): logger.error(f"Failed to connect to {url}") @@ -435,9 +395,6 @@ class RTSPProcessor: """捕获RTSP流帧""" cap = self.connect_to_rtsp_stream(self.rtsp_url) - # 设置缓冲区大小 - cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) - logger.info(f"开始捕获RTSP流: {self.rtsp_url}") try: diff --git a/util/entity_utl.py b/util/entity_utl.py new file mode 100644 index 0000000..d3ca645 --- /dev/null +++ b/util/entity_utl.py @@ -0,0 +1,22 @@ +import numpy as np +from datetime import datetime +from typing import List +from dataclasses import dataclass + +@dataclass +class DetectionResult: + """检测结果数据类""" + boxes: np.ndarray + confidences: np.ndarray + class_ids: np.ndarray + class_names: List[str] + timestamp: datetime + +@dataclass +class AlarmConfig: + """告警配置""" + target_classes: List[str] # 目标类别 + confidence_threshold: float = 0.5 # 置信度阈值 + alarm_duration: int = 10 # 告警录制时长(秒) + cooldown_duration: int = 30 # 告警冷却时间(秒) + save_path: str = "./alarm_videos" # 保存路径 \ No newline at end of file diff --git a/util/yolo_util.py b/util/yolo_util.py new file mode 100644 index 0000000..3f8bb03 --- /dev/null +++ b/util/yolo_util.py @@ -0,0 +1,33 @@ +from ultralytics import YOLO +import numpy as np +from util.entity_utl import DetectionResult +from datetime import datetime + +class YOLODetector: + """YOLO检测器""" + def __init__(self, model_path: str = "yolov8n.pt"): + self.model = YOLO(model_path) + self.class_names = self.model.names + + def detect(self, frame: np.ndarray, confidence_threshold: float = 0.5) -> DetectionResult: + """执行目标检测""" + results = self.model(frame, conf=confidence_threshold, verbose=False) + + if len(results) > 0 and results[0].boxes is not None: + boxes = results[0].boxes.xyxy.cpu().numpy() + confidences = results[0].boxes.conf.cpu().numpy() + class_ids = results[0].boxes.cls.cpu().numpy().astype(int) + class_names = [self.class_names[id] for id in class_ids] + else: + boxes = np.array([]) + confidences = np.array([]) + class_ids = np.array([]) + class_names = [] + + return DetectionResult( + boxes=boxes, + confidences=confidences, + class_ids=class_ids, + class_names=class_names, + timestamp=datetime.now() + ) \ No newline at end of file