sophon_ai_project/src/sophon_chakcy/utils/detect_video_stream.py

import sophon_chakcy.sail as sail
import threading
import time
from queue import Queue
from logging import getLogger, Logger
from typing import Optional


class DetectVideoStream:
    def __init__(self, 
                 detectors, 
                 video_path: str, 
                 max_frames: int = 0, 
                 frame_skip: int = 0, 
                 reconnect_attempts: int = 3, 
                 reconnect_delay: int = 5,
                 is_save: bool = False,
                 output: str = "./result/output.mp4",
                 logger: Optional[Logger] = None):
        """
        处理视频流（视频文件或RTSP流）

        Args:
            detector: 检测器
            video_path: 视频文件路径或 RTSP URL
            max_frames: 最大处理帧数, 0表示处理所有帧
            frame_skip: 跳帧处理, 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 = self.detectors[0].handle
        self.video_path = video_path
        self.max_frames = max_frames
        self.frame_skip = frame_skip
        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
            import shutil

            result_dir = r"results"
             # 清理旧的目录
            if os.path.exists(result_dir):
                self.logger.info(f"清理旧的构建目录: {result_dir}")
                shutil.rmtree(result_dir, ignore_errors=True)
            
            # 创建目录
            result_dir_path = Path(result_dir)
            if not result_dir_path.exists():
                result_dir_path.mkdir(parents=True, exist_ok=True)

            enc_params = f"width=1920:height=1080:bitrate=2000:gop=32:gop_preset=2:framerate=12"
            self.encoder = sail.Encoder(output, self.dev_id, 'h264_bm', 'NV12', enc_params, 10)

    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

        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
                    
                # 读取一帧
                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("视频流无法继续读取")
                        break

                    self.logger.info(f"尝试重连 ({reconnect_count}/{self.reconnect_attempts})...")
                    time.sleep(self.reconnect_delay)
                    self.decoder.release()
                    decoder = sail.Decoder(self.video_path, True, self.dev_id)
                    if decoder.is_opened():
                        reconnect_count = 0
                        self.logger.info("重连成功，继续处理...")
                        continue
                    else:
                        self.logger.warning("重连失败")
                # time.sleep(0.01)
                
                if self.frame_skip > 0 and frame_count % (self.frame_skip + 1) != 0:
                    frame_count += 1
                    continue

                # 开始处理当前帧
                frame_count += 1

                if frame_count % 10 == 0:
                    self.logger.info(f"处理第 {frame_count} 帧...")

                filtered_detections_queue = Queue()    
                
                def thread_detect(bmimg, detector):
                    filtered_detections = detector.detect(bmimg)
                    filtered_detections_queue.put((detector, filtered_detections.copy()))
                
                # 检测
                threads = []
                for detector in self.detectors:
                    t = threading.Thread(target=thread_detect, args=(bmimg, detector, ))
                    t.start()
                    threads.append(t)
                for t in threads:
                    t.join(timeout=2)
                
                while not filtered_detections_queue.empty():
                    detector, filtered_detections=filtered_detections_queue.get()
                    if len(filtered_detections) > 0:
                        bmimg = detector.draw_detections(bmimg, filtered_detections, None)
                        
                if self.is_save:
                    self.encoder.video_write(bmimg) # type: ignore
                    # time.sleep(0.01)

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