474 lines
16 KiB
Python
474 lines
16 KiB
Python
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() |