From e09432f0e4384c556713457b2cfb88279f4b8512 Mon Sep 17 00:00:00 2001 From: haotian <2421912570@qq.com> Date: Wed, 17 Sep 2025 11:13:12 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E9=85=8D=E7=BD=AE=E6=96=87?= =?UTF-8?q?=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backup/rtsp_processor_0.py | 474 +++++++++++++++++++++++++++++++++++++ config/config.yaml | 13 + rtsp_processor_1.py | 48 ++-- 3 files changed, 517 insertions(+), 18 deletions(-) create mode 100644 backup/rtsp_processor_0.py create mode 100644 config/config.yaml diff --git a/backup/rtsp_processor_0.py b/backup/rtsp_processor_0.py new file mode 100644 index 0000000..f078e34 --- /dev/null +++ b/backup/rtsp_processor_0.py @@ -0,0 +1,474 @@ +import cv2 +import numpy as np +import time +import threading +import queue +from datetime import datetime, timedelta +from pathlib import Path +import subprocess +import logging +from typing import List, Dict, Tuple, Optional +from dataclasses import dataclass +from collections import deque +import json + +try: + from ultralytics import YOLO +except ImportError: + print("请安装ultralytics: pip install ultralytics") + raise + +# 配置日志 +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger(__name__) + +@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: + """帧缓冲区，用于告警录制""" + def __init__(self, max_duration: int = 60, fps: int = 25): + self.max_frames = max_duration * fps + self.buffer = deque(maxlen=self.max_frames) + self.fps = fps + + def add_frame(self, frame: np.ndarray, timestamp: datetime): + self.buffer.append((frame.copy(), timestamp)) + + def get_frames_in_range(self, start_time: datetime, duration: int) -> List[Tuple[np.ndarray, datetime]]: + """获取指定时间范围内的帧""" + end_time = start_time + timedelta(seconds=duration) + frames = [] + + for frame, timestamp in self.buffer: + if start_time <= timestamp <= end_time: + frames.append((frame, timestamp)) + + 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: + """告警管理器""" + def __init__(self, config: AlarmConfig): + self.config = config + self.last_alarm_time = None + self.is_alarming = False + self.alarm_start_time = None + + # 创建保存目录 + Path(config.save_path).mkdir(parents=True, exist_ok=True) + + def check_alarm_trigger(self, detection_result: DetectionResult) -> bool: + """检查是否触发告警""" + current_time = datetime.now() + + # 检查冷却时间 + if (self.last_alarm_time and + current_time - self.last_alarm_time < timedelta(seconds=self.config.cooldown_duration)): + return False + + # 检查是否检测到目标类别 + for class_name, confidence in zip(detection_result.class_names, detection_result.confidences): + if (class_name in self.config.target_classes and + confidence >= self.config.confidence_threshold): + return True + + return False + + def start_alarm(self) -> bool: + """开始告警""" + if not self.is_alarming: + self.is_alarming = True + self.alarm_start_time = datetime.now() + logger.info(f"告警开始: {self.alarm_start_time}") + return True + return False + + def should_stop_alarm(self) -> bool: + """检查是否应该停止告警""" + if self.is_alarming and self.alarm_start_time: + duration = (datetime.now() - self.alarm_start_time).total_seconds() + return duration >= self.config.alarm_duration + return False + + def stop_alarm(self): + """停止告警""" + if self.is_alarming: + self.is_alarming = False + self.last_alarm_time = datetime.now() + logger.info(f"告警结束: {self.last_alarm_time}") + + def save_alarm_video(self, frames: List[Tuple[np.ndarray, datetime]], fps: int = 25): + """保存告警视频""" + if not frames: + logger.warning("没有帧数据可保存") + return None + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"alarm_{timestamp}.mp4" + filepath = Path(self.config.save_path) / filename + + # 获取帧尺寸 + height, width = frames[0][0].shape[:2] + + # 创建视频写入器 + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter(str(filepath), fourcc, fps, (width, height)) + + try: + for frame, _ in frames: + out.write(frame) + + logger.info(f"告警视频已保存: {filepath}") + return str(filepath) + + except Exception as e: + logger.error(f"保存视频失败: {e}") + return None + + finally: + out.release() + +class StreamServer: + """流媒体服务器接口""" + def __init__(self, output_url: str = "rtmp://localhost:1935/live/stream"): + self.output_url = output_url + self.ffmpeg_process = None + + def start_streaming(self, width: int, height: int, fps: int = 25): + """启动流媒体推送""" + ffmpeg_cmd = [ + 'ffmpeg', + '-y', # 覆盖输出文件 + '-f', 'rawvideo', + '-vcodec', 'rawvideo', + '-pix_fmt', 'bgr24', + '-s', f'{width}x{height}', + '-r', str(fps), + '-i', '-', # 从stdin读取 + '-c:v', 'libx264', + '-pix_fmt', 'yuv420p', + '-preset', 'ultrafast', + '-f', 'flv', + self.output_url + ] + + try: + self.ffmpeg_process = subprocess.Popen( + ffmpeg_cmd, + stdin=subprocess.PIPE, + stderr=subprocess.PIPE + ) + logger.info(f"流媒体服务启动: {self.output_url}") + return True + except Exception as e: + logger.error(f"启动流媒体失败: {e}") + return False + + def send_frame(self, frame: np.ndarray): + """发送帧到流媒体服务器""" + if self.ffmpeg_process and self.ffmpeg_process.stdin: + try: + self.ffmpeg_process.stdin.write(frame.tobytes()) + self.ffmpeg_process.stdin.flush() + except Exception as e: + logger.error(f"发送帧失败: {e}") + + def stop_streaming(self): + """停止流媒体推送""" + if self.ffmpeg_process: + self.ffmpeg_process.stdin.close() + self.ffmpeg_process.wait() + self.ffmpeg_process = None + logger.info("流媒体服务已停止") + +class RTSPProcessor: + """RTSP流处理器""" + def __init__(self, + rtsp_url: str, + model_path: str = "yolov8n.pt", + detection_interval: int = 5, + alarm_config: AlarmConfig = None, + output_stream_url: str = "rtmp://localhost:1935/live/stream"): + + self.rtsp_url = rtsp_url + self.detection_interval = detection_interval + self.frame_count = 0 + self.running = False + + # 初始化组件 + self.detector = YOLODetector(model_path) + self.alarm_manager = AlarmManager(alarm_config or AlarmConfig(target_classes=["person"])) + self.frame_buffer = FrameBuffer() + self.stream_server = StreamServer(output_stream_url) + + # 最后的检测结果 + self.last_detection = None + + # 线程安全队列 + self.frame_queue = queue.Queue(maxsize=100) + + def draw_detections(self, frame: np.ndarray, detection_result: DetectionResult) -> np.ndarray: + """在帧上绘制检测结果""" + if detection_result is None or len(detection_result.boxes) == 0: + return frame + + annotated_frame = frame.copy() + + for box, confidence, class_name in zip( + detection_result.boxes, + detection_result.confidences, + detection_result.class_names + ): + x1, y1, x2, y2 = box.astype(int) + + # 绘制边界框 + color = (0, 255, 0) # 绿色 + if class_name in self.alarm_manager.config.target_classes: + color = (0, 0, 255) # 目标类别用红色 + + cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2) + + # 绘制标签 + label = f"{class_name}: {confidence:.2f}" + label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0] + cv2.rectangle(annotated_frame, (x1, y1 - label_size[1] - 10), + (x1 + label_size[0], y1), color, -1) + cv2.putText(annotated_frame, label, (x1, y1 - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2) + + # 添加告警状态显示 + if self.alarm_manager.is_alarming: + cv2.putText(annotated_frame, "ALARM ACTIVE", (10, 30), + cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 3) + + # 添加时间戳 + timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + cv2.putText(annotated_frame, timestamp, (10, annotated_frame.shape[0] - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) + + return annotated_frame + + def process_frame(self, frame: np.ndarray) -> np.ndarray: + """处理单帧""" + current_time = datetime.now() + detection_result = None + + # 按间隔进行检测 + if self.frame_count % self.detection_interval == 0: + detection_result = self.detector.detect(frame, + self.alarm_manager.config.confidence_threshold) + self.last_detection = detection_result + + # 检查告警触发 + if self.alarm_manager.check_alarm_trigger(detection_result): + if self.alarm_manager.start_alarm(): + # 告警开始时的处理 + pass + + # 使用最后的检测结果绘制 + annotated_frame = self.draw_detections(frame, self.last_detection) + + # 添加帧到缓冲区 + self.frame_buffer.add_frame(annotated_frame, current_time) + + # 检查告警结束 + if self.alarm_manager.should_stop_alarm(): + # 获取告警期间的帧 + alarm_frames = self.frame_buffer.get_frames_in_range( + self.alarm_manager.alarm_start_time, + self.alarm_manager.config.alarm_duration + ) + + # 保存告警视频 + if alarm_frames: + threading.Thread( + target=self.alarm_manager.save_alarm_video, + args=(alarm_frames, 25), + daemon=True + ).start() + + self.alarm_manager.stop_alarm() + + return annotated_frame + + def capture_frames(self): + """捕获RTSP流帧""" + cap = cv2.VideoCapture(self.rtsp_url) + + if not cap.isOpened(): + logger.error(f"无法打开RTSP流: {self.rtsp_url}") + return + + # 设置缓冲区大小 + cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) + + logger.info(f"开始捕获RTSP流: {self.rtsp_url}") + + try: + while self.running: + ret, frame = cap.read() + if not ret: + logger.warning("读取帧失败，尝试重连...") + cap.release() + time.sleep(5) + cap = cv2.VideoCapture(self.rtsp_url) + continue + + if not self.frame_queue.full(): + self.frame_queue.put(frame) + else: + # 丢弃最旧的帧 + try: + self.frame_queue.get_nowait() + self.frame_queue.put(frame) + except queue.Empty: + pass + + finally: + cap.release() + logger.info("RTSP捕获已停止") + + def process_frames(self): + """处理帧线程""" + first_frame = True + + while self.running: + try: + frame = self.frame_queue.get(timeout=1) + + # 初始化流媒体服务器 + if first_frame: + height, width = frame.shape[:2] + if self.stream_server.start_streaming(width, height): + first_frame = False + else: + logger.error("流媒体服务器启动失败") + break + + # 处理帧 + processed_frame = self.process_frame(frame) + + # 发送到流媒体服务器 + self.stream_server.send_frame(processed_frame) + + self.frame_count += 1 + + except queue.Empty: + continue + except Exception as e: + logger.error(f"处理帧时出错: {e}") + + logger.info("帧处理已停止") + + def start(self): + """启动处理器""" + self.running = True + + # 启动捕获线程 + capture_thread = threading.Thread(target=self.capture_frames, daemon=True) + capture_thread.start() + + # 启动处理线程 + process_thread = threading.Thread(target=self.process_frames, daemon=True) + process_thread.start() + + logger.info("RTSP处理器已启动") + + return capture_thread, process_thread + + def stop(self): + """停止处理器""" + self.running = False + self.stream_server.stop_streaming() + logger.info("RTSP处理器已停止") + +def main(): + """主函数""" + # 配置参数 + RTSP_URL = "rtsp://10.0.0.17:8554/camera_test/2" # 替换为实际RTSP地址 + MODEL_PATH = "yolov8n.pt" # YOLO模型路径 + DETECTION_INTERVAL = 5 # 每5帧检测一次 + OUTPUT_STREAM_URL = "rtmp://localhost:1935/live/processed" # 输出流地址 + + # 告警配置 + alarm_config = AlarmConfig( + target_classes=["person"], # 目标类别 + confidence_threshold=0.6, + alarm_duration=15, # 告警录制15秒 + cooldown_duration=60*5, # 冷却60秒 + save_path="./alarm_videos" + ) + + # 创建处理器 + processor = RTSPProcessor( + rtsp_url=RTSP_URL, + model_path=MODEL_PATH, + detection_interval=DETECTION_INTERVAL, + alarm_config=alarm_config, + output_stream_url=OUTPUT_STREAM_URL + ) + + try: + # 启动处理 + threads = processor.start() + + # 保持运行 + logger.info("系统运行中，按Ctrl+C停止...") + while True: + time.sleep(1) + + except KeyboardInterrupt: + logger.info("收到停止信号") + + finally: + processor.stop() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/config/config.yaml b/config/config.yaml new file mode 100644 index 0000000..51fcc46 --- /dev/null +++ b/config/config.yaml @@ -0,0 +1,13 @@ +RTSP_URL : "rtsp://10.0.0.17:8554/camera_test/2" # 替换为实际RTSP地址 +MODEL_PATH : "yolov8n.pt" # YOLO模型路径 +DETECTION_INTERVAL : 5 # 每5帧检测一次 +OUTPUT_STREAM_URL : "rtmp://localhost:1935/live/processed" # 输出流地址 +TARGET_CLASSES: ["person", "car", "truck"] + + +alarm_save_path: "./alarm_videos" +confidence_threshold: 0.6 +alarm_duration: 15 # 告警录制15秒 +cooldown_duration: 600 # 冷却600秒 +output_stream_url: "rtmp://localhost:1935/live/stream" +output_url: "rtmp://127.0.0.1:1935/live/stream" \ No newline at end of file diff --git a/rtsp_processor_1.py b/rtsp_processor_1.py index cd0994f..af20d63 100644 --- a/rtsp_processor_1.py +++ b/rtsp_processor_1.py @@ -10,7 +10,7 @@ import logging from typing import List, Dict, Tuple, Optional from dataclasses import dataclass from collections import deque -import json +import yaml try: from ultralytics import YOLO @@ -18,6 +18,9 @@ except ImportError: print("请安装ultralytics: pip install ultralytics") raise +with open("./config/config.yaml", "r") as f: + config = yaml.safe_load(f) + # 配置日志 logging.basicConfig( level=logging.INFO, @@ -412,6 +415,15 @@ class RTSPProcessor: return annotated_frame + def connect_to_rtsp_stream(self, url): + """ 尝试连接到 RTSP 流 """ + cap = cv2.VideoCapture(url) + cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'HEVC')) + if not cap.isOpened(): + logger.error(f"Failed to connect to {url}") + return None + return cap + def capture_frames(self): """捕获RTSP流帧""" cap = cv2.VideoCapture(self.rtsp_url) @@ -506,43 +518,43 @@ class RTSPProcessor: def main(): """主函数""" # 配置参数 - RTSP_URL = "rtsp://10.0.0.17:8554/camera_test/2" # 替换为实际RTSP地址 - MODEL_PATH = "yolov8n.pt" # YOLO模型路径 - DETECTION_INTERVAL = 5 # 每5帧检测一次 - OUTPUT_STREAM_URL = "rtmp://localhost:1935/live/processed" # 输出流地址 - + RTSP_URL = config["RTSP_URL"] # 替换为实际RTSP地址 + MODEL_PATH = config["MODEL_PATH"] # YOLO模型路径 + DETECTION_INTERVAL = config["DETECTION_INTERVAL"] # 每5帧检测一次 + OUTPUT_STREAM_URL = config["OUTPUT_STREAM_URL"] # 输出流地址 + # 告警配置 alarm_config = AlarmConfig( - target_classes=["person", "car", "truck"], # 目标类别 - confidence_threshold=0.6, - alarm_duration=15, # 告警录制15秒 - cooldown_duration=60*60, # 冷却60秒 - save_path="./alarm_videos" + target_classes=config["TARGET_CLASSES"], # 目标类别 + confidence_threshold=config["confidence_threshold"], + alarm_duration=config["alarm_duration"], # 告警录制15秒 + cooldown_duration=config["cooldown_duration"], # 冷却60秒 + save_path=config["alarm_save_path"], ) - + # 创建处理器 processor = RTSPProcessor( rtsp_url=RTSP_URL, model_path=MODEL_PATH, detection_interval=DETECTION_INTERVAL, alarm_config=alarm_config, - output_stream_url=OUTPUT_STREAM_URL + output_stream_url=OUTPUT_STREAM_URL, ) - + try: # 启动处理 threads = processor.start() - + # 保持运行 logger.info("系统运行中，按Ctrl+C停止...") while True: time.sleep(1) - + except KeyboardInterrupt: logger.info("收到停止信号") - + finally: processor.stop() if __name__ == "__main__": - main() \ No newline at end of file + main()