添加配置文件

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()

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"

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地址
+    RTSP_URL = config["RTSP_URL"]  # 替换为实际RTSP地址
-    MODEL_PATH = "yolov8n.pt"  # YOLO模型路径
+    MODEL_PATH = config["MODEL_PATH"]  # YOLO模型路径
-    DETECTION_INTERVAL = 5  # 每5帧检测一次
+    DETECTION_INTERVAL = config["DETECTION_INTERVAL"]  # 每5帧检测一次
-    OUTPUT_STREAM_URL = "rtmp://localhost:1935/live/processed"  # 输出流地址
+    OUTPUT_STREAM_URL = config["OUTPUT_STREAM_URL"]  # 输出流地址
-    
+
     # 告警配置
     alarm_config = AlarmConfig(
-        target_classes=["person", "car", "truck"],  # 目标类别
+        target_classes=config["TARGET_CLASSES"],  # 目标类别
-        confidence_threshold=0.6,
+        confidence_threshold=config["confidence_threshold"],
-        alarm_duration=15,  # 告警录制15秒
+        alarm_duration=config["alarm_duration"],  # 告警录制15秒
-        cooldown_duration=60*60,  # 冷却60秒
+        cooldown_duration=config["cooldown_duration"],  # 冷却60秒
-        save_path="./alarm_videos"
+        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()
+    main()