GPU测试

benchmark/decode_thread.py

@@ -0,0 +1,299 @@
+"""
+视频解码模块
+"""
+
+import os
+import time
+import threading
+import queue
+from typing import Optional, Dict, Tuple
+from dataclasses import dataclass
+from enum import Enum
+
+import cv2
+import numpy as np
+
+from .utils import setup_logging, RateCounter
+
+logger = setup_logging()
+
+
+class SourceType(Enum):
+    RTSP = "rtsp"
+    FILE = "file"
+    SYNTHETIC = "synthetic"
+
+
+@dataclass
+class DecodeStats:
+    """解码统计信息"""
+    total_frames: int = 0
+    dropped_frames: int = 0
+    decode_errors: int = 0
+    reconnect_count: int = 0
+    avg_decode_time_ms: float = 0.0
+    current_fps: float = 0.0
+
+
+class DecodeThread(threading.Thread):
+    """视频解码线程"""
+    
+    def __init__(
+        self,
+        source_id: str,
+        source: str,
+        frame_queue: queue.Queue,
+        target_fps: float,
+        source_type: str = "rtsp",
+        resolution: Tuple[int, int] = (640, 480)
+    ):
+        super().__init__(daemon=True)
+        
+        self.source_id = source_id
+        self.source = source
+        self.frame_queue = frame_queue
+        self.target_fps = target_fps
+        self.source_type = SourceType(source_type)
+        self.resolution = resolution
+        
+        self._running = False
+        self._lock = threading.Lock()
+        self._cap: Optional[cv2.VideoCapture] = None
+        
+        self.stats = DecodeStats()
+        self._rate_counter = RateCounter(window_size=int(target_fps * 2))
+        self._decode_times = []
+        
+        self._frame_interval = 1.0 / target_fps
+        self._last_frame_time = 0.0
+    
+    def run(self):
+        self._running = True
+        logger.info(f"[{self.source_id}] 解码线程启动，目标帧率: {self.target_fps} FPS")
+        
+        while self._running:
+            try:
+                if self.source_type == SourceType.SYNTHETIC:
+                    self._decode_synthetic()
+                elif self.source_type == SourceType.FILE:
+                    self._decode_file()
+                else:
+                    self._decode_rtsp()
+            except Exception as e:
+                logger.error(f"[{self.source_id}] 解码异常: {e}")
+                self.stats.decode_errors += 1
+                time.sleep(1.0)
+        
+        self._cleanup()
+        logger.info(f"[{self.source_id}] 解码线程停止")
+    
+    def _decode_synthetic(self):
+        """合成图像源解码"""
+        while self._running:
+            current_time = time.time()
+            elapsed = current_time - self._last_frame_time
+            if elapsed < self._frame_interval:
+                time.sleep(self._frame_interval - elapsed)
+            
+            start_time = time.perf_counter()
+            frame = self._generate_synthetic_frame()
+            
+            decode_time = (time.perf_counter() - start_time) * 1000
+            self._decode_times.append(decode_time)
+            if len(self._decode_times) > 100:
+                self._decode_times.pop(0)
+            
+            self._put_frame(frame)
+            
+            self._last_frame_time = time.time()
+            self.stats.total_frames += 1
+            self._rate_counter.tick()
+    
+    def _generate_synthetic_frame(self) -> np.ndarray:
+        """生成合成帧"""
+        h, w = self.resolution
+        frame = np.random.randint(50, 200, (h, w, 3), dtype=np.uint8)
+        
+        num_objects = np.random.randint(0, 5)
+        for _ in range(num_objects):
+            x1 = np.random.randint(0, w - 50)
+            y1 = np.random.randint(0, h - 100)
+            x2 = x1 + np.random.randint(30, 80)
+            y2 = y1 + np.random.randint(60, 150)
+            color = tuple(np.random.randint(0, 255, 3).tolist())
+            cv2.rectangle(frame, (x1, y1), (x2, y2), color, -1)
+        
+        return frame
+    
+    def _decode_file(self):
+        """本地视频文件解码"""
+        if not os.path.exists(self.source):
+            logger.error(f"[{self.source_id}] 视频文件不存在: {self.source}")
+            return
+        
+        self._cap = cv2.VideoCapture(self.source)
+        if not self._cap.isOpened():
+            logger.error(f"[{self.source_id}] 无法打开视频文件: {self.source}")
+            return
+        
+        source_fps = self._cap.get(cv2.CAP_PROP_FPS)
+        frame_skip = max(1, int(source_fps / self.target_fps))
+        frame_count = 0
+        
+        while self._running:
+            start_time = time.perf_counter()
+            
+            ret, frame = self._cap.read()
+            if not ret:
+                self._cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+                continue
+            
+            frame_count += 1
+            if frame_count % frame_skip != 0:
+                continue
+            
+            decode_time = (time.perf_counter() - start_time) * 1000
+            self._decode_times.append(decode_time)
+            if len(self._decode_times) > 100:
+                self._decode_times.pop(0)
+            
+            self._put_frame(frame)
+            
+            self.stats.total_frames += 1
+            self._rate_counter.tick()
+            
+            elapsed = time.perf_counter() - start_time
+            sleep_time = self._frame_interval - elapsed
+            if sleep_time > 0:
+                time.sleep(sleep_time)
+    
+    def _decode_rtsp(self):
+        """RTSP 流解码"""
+        max_retries = 3
+        retry_count = 0
+        
+        while self._running and retry_count < max_retries:
+            try:
+                self._cap = cv2.VideoCapture(self.source, cv2.CAP_FFMPEG)
+                if not self._cap.isOpened():
+                    raise RuntimeError(f"无法连接 RTSP: {self.source}")
+                
+                self._cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
+                logger.info(f"[{self.source_id}] RTSP 连接成功")
+                retry_count = 0
+                
+                while self._running:
+                    start_time = time.perf_counter()
+                    
+                    ret, frame = self._cap.read()
+                    if not ret:
+                        logger.warning(f"[{self.source_id}] RTSP 读取失败")
+                        break
+                    
+                    decode_time = (time.perf_counter() - start_time) * 1000
+                    self._decode_times.append(decode_time)
+                    if len(self._decode_times) > 100:
+                        self._decode_times.pop(0)
+                    
+                    self._put_frame(frame)
+                    
+                    self.stats.total_frames += 1
+                    self._rate_counter.tick()
+                    
+                    elapsed = time.perf_counter() - start_time
+                    sleep_time = self._frame_interval - elapsed
+                    if sleep_time > 0:
+                        time.sleep(sleep_time)
+                
+            except Exception as e:
+                logger.error(f"[{self.source_id}] RTSP 错误: {e}")
+                retry_count += 1
+                self.stats.reconnect_count += 1
+                time.sleep(5.0)
+            finally:
+                if self._cap:
+                    self._cap.release()
+                    self._cap = None
+        
+        if retry_count >= max_retries:
+            logger.error(f"[{self.source_id}] RTSP 重连失败，切换到合成源")
+            self.source_type = SourceType.SYNTHETIC
+    
+    def _put_frame(self, frame: np.ndarray):
+        """将帧放入队列"""
+        timestamp = time.time()
+        
+        try:
+            if self.frame_queue.full():
+                try:
+                    self.frame_queue.get_nowait()
+                    self.stats.dropped_frames += 1
+                except queue.Empty:
+                    pass
+            
+            self.frame_queue.put_nowait((frame, timestamp, self.source_id))
+        except queue.Full:
+            self.stats.dropped_frames += 1
+    
+    def stop(self):
+        self._running = False
+    
+    def _cleanup(self):
+        with self._lock:
+            if self._cap:
+                self._cap.release()
+                self._cap = None
+    
+    def get_stats(self) -> DecodeStats:
+        self.stats.current_fps = self._rate_counter.get_rate()
+        if self._decode_times:
+            self.stats.avg_decode_time_ms = sum(self._decode_times) / len(self._decode_times)
+        return self.stats
+
+
+class FrameQueueManager:
+    """帧队列管理器"""
+    
+    def __init__(self, queue_size: int = 2):
+        self.queue_size = queue_size
+        self.queues: Dict[str, queue.Queue] = {}
+        self.decode_threads: Dict[str, DecodeThread] = {}
+    
+    def add_source(
+        self,
+        source_id: str,
+        source: str,
+        target_fps: float,
+        source_type: str = "synthetic",
+        resolution: Tuple[int, int] = (640, 480)
+    ) -> queue.Queue:
+        frame_queue = queue.Queue(maxsize=self.queue_size)
+        self.queues[source_id] = frame_queue
+        
+        decode_thread = DecodeThread(
+            source_id=source_id,
+            source=source,
+            frame_queue=frame_queue,
+            target_fps=target_fps,
+            source_type=source_type,
+            resolution=resolution
+        )
+        self.decode_threads[source_id] = decode_thread
+        
+        return frame_queue
+    
+    def start_all(self):
+        for thread in self.decode_threads.values():
+            thread.start()
+    
+    def stop_all(self):
+        for thread in self.decode_threads.values():
+            thread.stop()
+        for thread in self.decode_threads.values():
+            thread.join(timeout=3.0)
+    
+    def get_all_stats(self) -> Dict[str, DecodeStats]:
+        return {source_id: thread.get_stats() for source_id, thread in self.decode_threads.items()}
+    
+    def get_total_dropped_frames(self) -> int:
+        return sum(thread.stats.dropped_frames for thread in self.decode_threads.values())