modify: 多线程任务调度, 性能待优化
This commit is contained in:
@@ -1,22 +1,85 @@
|
||||
import cv2
|
||||
import sophon_chakcy.sail as sail
|
||||
import threading
|
||||
import time
|
||||
from queue import Queue
|
||||
from logging import getLogger, Logger
|
||||
from typing import Optional
|
||||
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,
|
||||
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):
|
||||
logger: Optional[Logger] = None,
|
||||
target_fps: float = 2.0,
|
||||
use_opencv_info: bool = True):
|
||||
"""
|
||||
处理视频流（视频文件或RTSP流）
|
||||
|
||||
@@ -24,7 +87,6 @@ class DetectVideoStream:
|
||||
detector: 检测器
|
||||
video_path: 视频文件路径或 RTSP URL
|
||||
max_frames: 最大处理帧数, 0表示处理所有帧
|
||||
frame_skip: 跳帧处理, 0表示处理每一帧
|
||||
reconnect_attempts: RTSP断流重连尝试次数
|
||||
reconnect_delay: 重连延迟时间 (秒)
|
||||
|
||||
@@ -37,36 +99,120 @@ class DetectVideoStream:
|
||||
self.logger = logger
|
||||
self.detectors = detectors
|
||||
self.dev_id = self.detectors[0].dev_id
|
||||
self.handle = self.detectors[0].handle
|
||||
self.handle = sail.Handle(self.dev_id)
|
||||
self.video_path = video_path
|
||||
self.max_frames = max_frames
|
||||
self.frame_skip = frame_skip
|
||||
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
|
||||
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)
|
||||
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 = 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)
|
||||
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
|
||||
@@ -76,6 +222,8 @@ class DetectVideoStream:
|
||||
|
||||
frame_count = 0
|
||||
reconnect_count = 0
|
||||
skip_interval = self._get_discard_frame_interval()
|
||||
skip_counter = 0 # 跳帧计数器
|
||||
|
||||
try:
|
||||
while not self.stop_event.is_set():
|
||||
@@ -83,63 +231,53 @@ class DetectVideoStream:
|
||||
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()
|
||||
decoder = sail.Decoder(self.video_path, True, self.dev_id)
|
||||
if decoder.is_opened():
|
||||
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("重连失败")
|
||||
# time.sleep(0.01)
|
||||
|
||||
if self.frame_skip > 0 and frame_count % (self.frame_skip + 1) != 0:
|
||||
frame_count += 1
|
||||
continue
|
||||
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()))
|
||||
# 5. 重置重连计数器
|
||||
reconnect_count = 0
|
||||
# 跳帧逻辑：只处理每skip_interval帧中的第1帧
|
||||
skip_counter += 1
|
||||
if skip_counter % skip_interval != 1:
|
||||
continue # 跳过此帧
|
||||
|
||||
# 检测
|
||||
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)
|
||||
filtered_detections = detector.detect(bmimg)
|
||||
bmimg = detector.draw_detections(bmimg, filtered_detections, None)
|
||||
|
||||
if self.is_save:
|
||||
self.encoder.video_write(bmimg) # type: ignore
|
||||
# time.sleep(0.01)
|
||||
|
||||
frame_count += 1
|
||||
|
||||
except KeyboardInterrupt:
|
||||
self.logger.info("\n检测到键盘中断，停止处理")
|
||||
@@ -160,3 +298,4 @@ class DetectVideoStream:
|
||||
def stop_task(self):
|
||||
self.logger.info("提前关闭")
|
||||
self.stop_event.set()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user