import cv2 import time import threading import queue import yaml import os import json import base64 import asyncio import websockets from dataclasses import dataclass from typing import Optional, Dict, Any, Tuple # ========================= # 配置与数据结构 # ========================= @dataclass class CameraConfig: id: int name: str rtsp_url: str RTSP_TARGET_FPS = 10.0 # 固定 10 帧/秒 FRAMES_PER_SEGMENT = 600 # 每 600 帧一个 mp4 VIDEO_OUTPUT_DIR = "./videos" # 视频输出目录 WS_HOST = "0.0.0.0" # WebSocket 服务端监听地址 WS_PORT = 8765 # WebSocket 服务端端口 # 已连接的 WebSocket 客户端集合 ws_clients = set() # ========================= # WebSocket 服务线程 # ========================= class WebSocketSender(threading.Thread): """ WebSocket 服务端线程： - 在 WS_HOST:WS_PORT 上启动 websockets 服务器 - 从 send_queue 中读取消息，广播给所有已连接客户端 """ def __init__(self, send_queue: "queue.Queue[Dict[str, Any]]", stop_event: threading.Event): super().__init__(daemon=True) self.send_queue = send_queue self.stop_event = stop_event async def _ws_handler(self, websocket): # 新客户端连接 ws_clients.add(websocket) try: async for _ in websocket: # 当前忽略客户端发送的消息 pass finally: # 客户端断开 ws_clients.discard(websocket) async def _broadcaster(self): """从队列中取出消息并广播给所有连接的客户端""" while not self.stop_event.is_set(): try: # 在线程池中阻塞等待队列消息 msg = await asyncio.to_thread(self.send_queue.get, timeout=0.5) except queue.Empty: continue data = json.dumps(msg) dead = [] for ws in list(ws_clients): try: await ws.send(data) except Exception: dead.append(ws) for ws in dead: ws_clients.discard(ws) self.send_queue.task_done() async def _run_async(self): async with websockets.serve(self._ws_handler, WS_HOST, WS_PORT): print(f"[INFO] WebSocket server started at ws://{WS_HOST}:{WS_PORT}") await self._broadcaster() def run(self): asyncio.run(self._run_async()) # ========================= # RTSP 抓流线程 # ========================= class RTSPCaptureWorker(threading.Thread): """ 只负责从 RTSP 读取原始帧，放入 raw_frame_queue。 不负责抽帧、不负责写视频。 """ def __init__( self, camera_cfg: CameraConfig, raw_frame_queue: "queue.Queue[Dict[str, Any]]", stop_event: threading.Event, ): super().__init__(daemon=True) self.camera_cfg = camera_cfg self.raw_frame_queue = raw_frame_queue self.stop_event = stop_event def run(self): cap = cv2.VideoCapture(self.camera_cfg.rtsp_url, cv2.CAP_FFMPEG) if not cap.isOpened(): print(f"[ERROR] Cannot open RTSP stream: {self.camera_cfg.rtsp_url}") return print(f"[INFO] Start capturing: id={self.camera_cfg.id}, name={self.camera_cfg.name}") while not self.stop_event.is_set(): ok, frame = cap.read() if not ok: print(f"[WARN] Failed to read frame from camera {self.camera_cfg.id}, retrying...") time.sleep(0.2) continue ts = time.time() item = { "camera_id": self.camera_cfg.id, "camera_name": self.camera_cfg.name, "timestamp": ts, "frame": frame, } try: self.raw_frame_queue.put(item, timeout=1.0) except queue.Full: print(f"[WARN] Raw frame queue full, drop frame from camera {self.camera_cfg.id}") cap.release() print(f"[INFO] Stop capturing: id={self.camera_cfg.id}") # ========================= # 帧处理线程（抽帧 + 写mp4 + 调用用户函数 + 发WebSocket消息） # ========================= class FrameProcessorWorker(threading.Thread): def __init__( self, raw_frame_queue: "queue.Queue[Dict[str, Any]]", ws_send_queue: "queue.Queue[Dict[str, Any]]", stop_event: threading.Event, ): super().__init__(daemon=True) self.raw_frame_queue = raw_frame_queue self.ws_send_queue = ws_send_queue self.stop_event = stop_event # 每个摄像头独立维护视频写入状态 self.video_writers: Dict[int, cv2.VideoWriter] = {} self.video_frame_counts: Dict[int, int] = {} self.video_segment_start_ts: Dict[int, float] = {} self.video_segment_filenames: Dict[int, str] = {} os.makedirs(VIDEO_OUTPUT_DIR, exist_ok=True) # 控制 10fps 抽帧：记录每个摄像头上次处理时间 self.last_process_ts: Dict[int, float] = {} def _get_video_writer(self, camera_id: int, frame) -> Tuple[cv2.VideoWriter, str]: """ 获取（或新建）当前摄像头的 VideoWriter。 如果当前 segment 不存在，则新建一个，文件名由第一帧时间命名。 """ writer = self.video_writers.get(camera_id) if writer is not None: return writer, self.video_segment_filenames[camera_id] h, w = frame.shape[:2] fourcc = cv2.VideoWriter_fourcc(*"mp4v") start_ts = time.time() self.video_segment_start_ts[camera_id] = start_ts ts_str = time.strftime("%Y%m%d_%H%M%S", time.localtime(start_ts)) filename = f"{ts_str}_cam{camera_id}.mp4" filepath = os.path.join(VIDEO_OUTPUT_DIR, filename) writer = cv2.VideoWriter(filepath, fourcc, RTSP_TARGET_FPS, (w, h)) self.video_writers[camera_id] = writer self.video_frame_counts[camera_id] = 0 self.video_segment_filenames[camera_id] = filepath print(f"[INFO] Start new segment: camera={camera_id}, file={filepath}") return writer, filepath def _close_segment_if_needed(self, camera_id: int): """ 如果当前segment达到 FRAMES_PER_SEGMENT，则关闭并清理状态。 """ count = self.video_frame_counts.get(camera_id, 0) if count >= FRAMES_PER_SEGMENT: writer = self.video_writers.get(camera_id) if writer is not None: writer.release() print(f"[INFO] Close segment: camera={camera_id}, file={self.video_segment_filenames[camera_id]}") # 清空当前 segment 状态 self.video_writers.pop(camera_id, None) self.video_frame_counts.pop(camera_id, None) self.video_segment_start_ts.pop(camera_id, None) self.video_segment_filenames.pop(camera_id, None) def _encode_image_to_base64(self, image) -> str: ok, buf = cv2.imencode(".jpg", image) if not ok: raise RuntimeError("Failed to encode image to JPEG") return base64.b64encode(buf.tobytes()).decode("ascii") def run(self): print("[INFO] FrameProcessorWorker started") target_interval = 1.0 / RTSP_TARGET_FPS while not self.stop_event.is_set(): try: item = self.raw_frame_queue.get(timeout=0.5) except queue.Empty: continue camera_id = item["camera_id"] ts = item["timestamp"] frame = item["frame"] last_ts = self.last_process_ts.get(camera_id, 0.0) if ts - last_ts < target_interval: # 丢弃多余帧，保证约10fps self.raw_frame_queue.task_done() continue self.last_process_ts[camera_id] = ts # 1) 写入 mp4 (当前segment) writer, video_filepath = self._get_video_writer(camera_id, frame) writer.write(frame) self.video_frame_counts[camera_id] = self.video_frame_counts.get(camera_id, 0) + 1 # 2) 调用用户自定义处理逻辑 result = user_process_frame(frame, camera_id, ts) if result is not None and "image" in result and "type" in result: result_img = result["image"] result_type = int(result["type"]) # 3) 通过 WebSocket 发送帧结果 try: img_b64 = self._encode_image_to_base64(result_img) except Exception as e: print(f"[ERROR] Encode image failed: {e}") img_b64 = None if img_b64 is not None: msg = { "msg_type": "frame", "camera_id": camera_id, "timestamp": ts, "result_type": result_type, "image_base64": img_b64, } try: self.ws_send_queue.put(msg, timeout=1.0) except queue.Full: print("[WARN] ws_send_queue full, drop frame message") # 4) 如果 result_type != 0，通过 WebSocket 发送告警 if result_type != 0: alert_msg = { "msg_type": "alert", "camera_id": camera_id, "event_type": result_type, "video_file": video_filepath, "timestamp": ts, } try: self.ws_send_queue.put(alert_msg, timeout=1.0) except queue.Full: print("[WARN] ws_send_queue full, drop alert message") # 5) 检查是否需要切换到下一个 mp4 segment self._close_segment_if_needed(camera_id) self.raw_frame_queue.task_done() # 退出时，关闭所有 VideoWriter for cam_id, writer in list(self.video_writers.items()): writer.release() print(f"[INFO] Release writer on exit: camera={cam_id}") print("[INFO] FrameProcessorWorker stopped") # ========================= # 用户自定义函数 (TBD) # ========================= def user_process_frame(image, camera_id: int, timestamp: float) -> Dict[str, Any]: """ 你在这里实现算法逻辑： - image: numpy.ndarray, BGR - camera_id: 摄像头 id - timestamp: 捕获时间戳 (time.time()) 返回: - {"image": image, "type": int} """ # TODO: 替换为你的实际逻辑，例如模型推理 result_type = 0 # 示例：默认0 return { "image": image, "type": result_type, } # ========================= # 服务封装 # ========================= class RTSPService: def __init__(self, config_path: str): self.config_path = config_path self.cameras = self._load_config() self.stop_event = threading.Event() # 队列 self.raw_frame_queue: "queue.Queue[Dict[str, Any]]" = queue.Queue(maxsize=500) self.ws_send_queue: "queue.Queue[Dict[str, Any]]" = queue.Queue(maxsize=1000) # 线程 self.capture_workers = [] self.frame_processor = FrameProcessorWorker(self.raw_frame_queue, self.ws_send_queue, self.stop_event) self.ws_sender = WebSocketSender(self.ws_send_queue, self.stop_event) def _load_config(self): with open(self.config_path, "r", encoding="utf-8") as f: cfg = yaml.safe_load(f) cameras_cfg = cfg.get("cameras", []) cameras = [] for c in cameras_cfg: cameras.append( CameraConfig( id=int(c["id"]), name=str(c.get("name", f"cam_{c['id']}")), rtsp_url=str(c["rtsp_url"]), ) ) return cameras def start(self): print("[INFO] RTSPService starting...") # 启动 WebSocket 发送线程 self.ws_sender.start() # 启动帧处理线程 self.frame_processor.start() # 启动每个摄像头的抓流线程 for cam in self.cameras: w = RTSPCaptureWorker(cam, self.raw_frame_queue, self.stop_event) w.start() self.capture_workers.append(w) print("[INFO] RTSPService started") def stop(self): print("[INFO] RTSPService stopping...") self.stop_event.set() # 等待队列处理完（可选） try: self.raw_frame_queue.join() self.ws_send_queue.join() except Exception: pass for w in self.capture_workers: w.join(timeout=1.0) self.frame_processor.join(timeout=1.0) self.ws_sender.join(timeout=1.0) print("[INFO] RTSPService stopped") def main(): service = RTSPService(config_path="config.yaml") service.start() try: while True: time.sleep(1.0) except KeyboardInterrupt: print("[INFO] KeyboardInterrupt, shutting down...") finally: service.stop() if __name__ == "__main__": main()