import cv2 import sophon_chakcy.sail as sail import threading import time from logging import getLogger, Logger from typing import Optional, Tuple import os class VideoStreamInfo: @staticmethod def get_video_info(video_path: str) -> Tuple[bool, dict]: """ 获取视频流信息 Args: video_path: 视频文件路径或 RTSP URL timeout: 连接超时时间（秒） Returns: (success, info_dict) """ cap = None try: # 设置 OpenCV RTSP参数 if video_path.startswith("rtsp://"): # RTSP 流参数优化 cap = cv2.VideoCapture(video_path) cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = "rtsp_transport;tcp" else: cap = cv2.VideoCapture(video_path) if not cap.isOpened(): return False, {"error": "无法打开视频源"} # 获取基本信息 fps = cap.get(cv2.CAP_PROP_FPS) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # 尝试读取一帧确认能正常解码 success, frame = cap.read() if not success: return False, {"error": "无法读取视频帧"} # 获取编码器信息 fourcc = int(cap.get(cv2.CAP_PROP_FOURCC)) codec = chr(fourcc & 0xFF) + chr((fourcc >> 8) & 0xFF) + \ chr((fourcc >> 16) & 0xFF) + chr((fourcc >> 24) & 0xFF) info = { "fps": fps if fps > 0 else 25.0, # 默认25fps "width": width, "height": height, "frame_count": frame_count if frame_count > 0 else 0, "codec": codec, "is_rtsp": video_path.startswith("rtsp://"), "is_file": video_path.endswith(('.mp4', '.avi', '.mkv', '.mov', '.flv')), "success": True } return True, info except Exception as e: return False, {"error": str(e)} finally: if cap is not None: cap.release() class DetectVideoStream: def __init__(self, detectors, video_path: str, max_frames: int = 0, reconnect_attempts: int = 3, reconnect_delay: int = 5, is_save: bool = False, output: str = "./result/output.mp4", logger: Optional[Logger] = None, target_fps: float = 2.0, use_opencv_info: bool = True): """ 处理视频流（视频文件或RTSP流） Args: detector: 检测器 video_path: 视频文件路径或 RTSP URL max_frames: 最大处理帧数, 0表示处理所有帧 reconnect_attempts: RTSP断流重连尝试次数 reconnect_delay: 重连延迟时间 (秒) Returns: 处理的帧数, 检测到的物体总数 """ if logger is None: self.logger = getLogger() else: self.logger = logger self.detectors = detectors self.dev_id = self.detectors[0].dev_id self.handle = sail.Handle(self.dev_id) self.video_path = video_path self.max_frames = max_frames self.target_fps = target_fps # 使用OpenCV获取视频信息 self.video_info = {} if use_opencv_info: success, info = VideoStreamInfo.get_video_info(video_path) if success: self.video_info = info self.logger.info(f"视频信息: {info}") else: self.logger.warning(f"无法获取视频信息: {info.get('error', '未知错误')}") # 设置默认值 if video_path.startswith("rtsp://"): self.video_info = {"fps": 25.0, "is_rtsp": True} else: self.video_info = {"fps": 30.0, "is_rtsp": False} # 获取原始帧率 self.original_fps = self.video_info.get("fps", 25.0) self.logger.info(f"原始帧率: {self.original_fps:.2f}FPS, 目标帧率: {self.target_fps:.2f}FPS") if not video_path.startswith("rtsp"): self.reconnect_attempts = 0 else: self.reconnect_attempts = reconnect_attempts self.reconnect_delay = reconnect_delay self.decoder = sail.Decoder(video_path, True, self.dev_id) self.stop_event = threading.Event() self.is_save = is_save if self.is_save: import os from pathlib import Path result_dir = r"results" # 清理旧的目录 if not os.path.exists(result_dir): # 创建目录 result_dir_path = Path(result_dir) if not result_dir_path.exists(): result_dir_path.mkdir(parents=True, exist_ok=True) enc_params_dict = { "width": self.video_info.get("width", 1920), "height": self.video_info.get("height", 1080), "bitrate": 2000, "gop": 32, "gop_preset": 2, "framerate": self.target_fps } self.logger.info(f"输出：{output}") enc_params = ":".join({f"{key}={value}" for key, value in enc_params_dict.items()}) self.encoder = sail.Encoder(output, self.dev_id, 'h264_bm', 'NV12', enc_params, 10) self.discard_count = self._get_discard_frames_based_on_time() def _get_discard_frames_based_on_time(self) -> int: """ 基于时间计算需要丢弃的帧数 原理：根据经过的时间和原始帧率，计算应该跳过的帧数 """ if self.original_fps <= 0 or self.target_fps <= 0: return 0 if self.target_fps >= self.original_fps: return 0 # 目标帧率高于原始帧率，不丢弃 # 计算采样间隔 sample_interval = self.original_fps / self.target_fps # 需要丢弃的帧数 = 采样间隔 - 1（因为要保留1帧） discard_frames = int(sample_interval - 1) # 确保非负 discard_frames = max(0, discard_frames) return discard_frames def _read_and_discard_frames(self): """读取并丢弃指定数量的帧""" discarded = 0 for i in range(self.discard_count): try: # 读取一帧但不处理 bmimg = sail.BMImage() ret = self.decoder.read(self.handle, bmimg) if ret == 0: discarded += 1 else: self.logger.warning(f"丢弃第 {i+1} 帧失败，错误码: {ret}") break except Exception as e: self.logger.warning(f"丢弃帧时发生错误: {e}") break if discarded > 0: self.logger.debug(f"已丢弃 {discarded} 帧") def _get_discard_frame_interval(self) -> int: """计算跳帧间隔，而不是丢弃帧数""" if self.original_fps <= 0 or self.target_fps <= 0: return 1 if self.target_fps >= self.original_fps: return 1 # 计算采样间隔 sample_interval = max(1, int(round(self.original_fps / self.target_fps))) return sample_interval def run_task(self) -> bool: result_flog = True if not self.decoder.is_opened(): self.logger.error(f"错误: 无法打开视频流 {self.video_path}") result_flog = False frame_count = 0 reconnect_count = 0 skip_interval = self._get_discard_frame_interval() skip_counter = 0 # 跳帧计数器 try: while not self.stop_event.is_set(): if self.max_frames > 0 and frame_count >= self.max_frames: self.logger.warning(f"\n达到最大帧数限制: {self.max_frames}") result_flog = True break # 3. 计算需要丢弃的帧数 # if self.discard_count > 0: # self._read_and_discard_frames() # 读取一帧 bmimg = sail.BMImage() ret = self.decoder.read(self.handle, bmimg) if ret != 0: self.logger.error(f"读取帧失败，错误码: {ret}") reconnect_count += 1 if reconnect_count > self.reconnect_attempts: self.logger.error("视频流无法继续读取") result_flog = False break self.logger.info(f"尝试重连 ({reconnect_count}/{self.reconnect_attempts})...") time.sleep(self.reconnect_delay) self.decoder.release() self.decoder = sail.Decoder(self.video_path, True, self.dev_id) if self.decoder.is_opened(): reconnect_count = 0 self.logger.info("重连成功，继续处理...") continue else: self.logger.warning("重连失败") continue # 5. 重置重连计数器 reconnect_count = 0 # 跳帧逻辑：只处理每skip_interval帧中的第1帧 skip_counter += 1 if skip_counter % skip_interval != 1: continue # 跳过此帧 # 检测 for detector in self.detectors: filtered_detections = detector.detect(bmimg) bmimg = detector.draw_detections(bmimg, filtered_detections, None) if self.is_save: self.encoder.video_write(bmimg) # type: ignore frame_count += 1 except KeyboardInterrupt: self.logger.info("\n检测到键盘中断，停止处理") result_flog = True except Exception as e: self.logger.error(f"处理视频流时发生错误: {e}") result_flog = False finally: # 清理资源 self.decoder.release() if self.is_save: self.encoder.release() self.logger.info("视频解码器已释放") return result_flog def stop_task(self): self.logger.info("提前关闭") self.stop_event.set()