modify: 实现视频流推理

2026-02-04 23:12:12 +08:00
parent 7880c9e15b
commit 2ff9a2ecb1
7 changed files with 258 additions and 74 deletions
@@ -0,0 +1,154 @@
+import sophon_chakcy.sail as sail
+import threading
+import time
+
+
+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):
+        """
+        处理视频流（视频文件或RTSP流）
+
+        Args:
+            detector: 检测器
+            video_path: 视频文件路径或 RTSP URL
+            max_frames: 最大处理帧数, 0表示处理所有帧
+            frame_skip: 跳帧处理, 0表示处理每一帧
+            reconnect_attempts: RTSP断流重连尝试次数
+            reconnect_delay: 重连延迟时间 (秒)
+        
+        Returns:
+            处理的帧数, 检测到的物体总数
+        """
+        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):
+                print(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=25"
+            self.encoder = sail.Encoder("results/output.mp4", self.dev_id, 'h264_bm', 'NV12', enc_params, 10)
+    
+    def run_task(self):     
+        if not self.decoder.is_opened():
+            print(f"错误: 无法打开视频流 {self.video_path}")
+            return 0, 0
+
+        frame_shape = self.decoder.get_frame_shape()
+        print(f"视频帧形状: {frame_shape}")
+
+        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:
+                    print(f"\n达到最大帧数限制: {self.max_frames}")
+                    break
+                    
+                # 读取一帧
+                bmimg = sail.BMImage()
+                ret = self.decoder.read(self.handle, bmimg)
+                if ret != 0:
+                    print(f"读取帧失败，错误码: {ret}")
+                    reconnect_count += 1
+
+                    if reconnect_count > self.reconnect_attempts:
+                        print("视频流无法继续读取")
+                        break
+
+                    print(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
+                        print("重连成功，继续处理...")
+                        continue
+                    else:
+                        print("重连失败")
+                
+                if self.frame_skip > 0 and frame_count % (self.frame_skip + 1) != 0:
+                    frame_count += 1
+                    if self.is_save:
+                        self.encoder.video_write(bmimg)
+                    continue
+
+                # 开始处理当前帧
+                start_time = time.time()
+                frame_count += 1
+
+                if frame_count % 10 == 0:
+                    print(f"处理第 {frame_count} 帧...")
+                    
+                # 检测
+                for detector in self.detectors:
+                    detections = detector.detect(bmimg)
+
+                    if len(detections) > 0:
+                        detections = detections[0]
+
+                        filtered_detections = []
+                        for det in detections:
+                            if len(det) >= 6:
+                                x1, y1, x2, y2, conf, cls_id = det[:6]
+                                if conf > detector.conf_thresh:
+                                    filtered_detections.append([x1, y1, x2, y2, conf, cls_id])
+
+                        if len(filtered_detections) > 0:
+                            result_img = detector.draw_detections(bmimg, filtered_detections, None)
+                        else:
+                            result_img = bmimg
+                    else:
+                        result_img = bmimg
+                
+                if self.is_save:
+                    self.encoder.video_write(result_img)
+
+        except KeyboardInterrupt:
+            print("\n检测到键盘中断，停止处理")
+        except Exception as e:
+            print(f"处理视频流时发生错误: {e}")
+
+        finally:
+            # 清理资源
+            self.decoder.release()
+            if self.is_save:
+                self.encoder.release()
+            print("视频解码器已释放")
+
+    def stop_task(self):
+        print("提前关闭")
+        self.stop_event.set()
+
+