sophon_ai_project/src/sophon_chakcy/utils/detect_video_stream.py

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