chakcy/sophon_ai_project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

modify: 多线程任务调度, 性能待优化

Browse Source
This commit is contained in:
chakcy
2026-02-07 19:32:38 +08:00
parent 91d8456d4c
commit 375d3518de
425 changed files with 582760 additions and 80820 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/sophon_chakcy/utils/detect_video_stream.py
+188 -49
View File
@@ -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()
src/sophon_chakcy/utils/multi_decoder_manager.py
+6
View File
@@ -0,0 +1,6 @@
import sophon_chakcy.sail as sail
class MultiDecoderManager:
def __init__(self):
pass
src/sophon_chakcy/utils/push_video_stream.py
-20
View File
@@ -1,20 +0,0 @@
import sophon_chakcy.sail as sail
from queue import Queue
class PushVideoStream:
def __init__(self, output, dev_id, width, height):
self.dev_id = dev_id
enc_params_dict = {
"width": width,
"height": height,
"bitrate": 2000,
"gop": 32,
"gop_preset": 2,
"framerate": 12
}
self.bmimg_queue = Queue(maxsize=30)
enc_params = ":".join({f"{key}:{value}" for key, value in enc_params_dict})
self.encoder = sail.Encoder(output, self.dev_id, 'h264_bm', 'NV12', enc_params, 10)
def
src/sophon_chakcy/utils/task_manager.py
+176 -14
View File
@@ -1,11 +1,11 @@
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import ThreadPoolExecutor, wait, ALL_COMPLETED, FIRST_COMPLETED
import threading
from dataclasses import dataclass, field
from .log_manager import LogManager
from .detect_video_stream import DetectVideoStream
from logging import getLogger
from ..get_detector import get_detector
from typing import List, Dict
from typing import List, Dict, Any, Optional, Tuple, Set
import time
@dataclass
@@ -54,10 +54,11 @@ class TaskManager:
max_workers=max_workers,
thread_name_prefix=thread_name_prefix
)
self._task_dict: Dict[str, DetectVideoStream] = {}
self._task_dict: Dict[str, Dict[str, Any]] = {}
self._results: List[bool] = []
self._lock = threading.Lock()
def submit(self, task_args: TaskArgs):
"""提交任务"""
if task_args.task_name in self._task_dict.keys():
@@ -84,27 +85,188 @@ class TaskManager:
task_args.input,
is_save=True,
logger=logger,
frame_skip=3,
output=task_args.output)
self._task_dict[task_args.task_name] = detect_video_stream
future = self.executor.submit(detect_video_stream.run_task)
future.add_done_callback(self._task_callback)
future.add_done_callback(lambda f: self._task_callback(f, task_args.task_name))
with self._lock:
self._task_dict[task_args.task_name] = {
'future': future,
'detect_video_stream': detect_video_stream,
'logger': logger,
'start_time': time.time(),
'status': 'running'
}
return future
def _task_callback(self, future):
def _task_callback(self, future, task_name):
"""任务完成回调"""
try:
result = future.result()
with self._lock:
if task_name in self._task_dict:
self._task_dict[task_name]['status'] = 'completed'
self._task_dict[task_name]['end_time'] = time.time()
self._results.append(result)
except Exception as e:
print(f"任务执行出错: {e}")
print(f"任务 '{task_name}' 执行出错: {e}")
with self._lock:
if task_name in self._task_dict:
self._task_dict[task_name]['status'] = 'failed'
self._task_dict[task_name]['error'] = str(e)
self._task_dict[task_name]['end_time'] = time.time()
def stop_task(self, task_name):
self._task_dict[task_name].stop_task()
def wait_for_task(self, task_name: str, timeout: Optional[float] = None) -> bool:
"""
等待指定任务完成
Args:
task_name: 任务名称
timeout: 超时时间（秒），None表示无限等待
Returns:
bool: 任务是否成功完成
Raises:
KeyError: 任务不存在
TimeoutError: 等待超时
"""
with self._lock:
if task_name not in self._task_dict:
raise KeyError(f"任务 '{task_name}' 不存在")
future = self._task_dict[task_name]['future']
try:
return future.result(timeout=timeout)
except Exception as e:
# 如果future已经完成但有异常，这里会抛出
raise e
def shutdown(self, wait=True):
def wait_for_all(self, timeout: Optional[float] = None,
return_when: str = ALL_COMPLETED) -> Tuple[Set, Set]:
"""
等待所有任务完成
Args:
timeout: 超时时间（秒），None表示无限等待
return_when: 何时返回，可选值：
- ALL_COMPLETED: 所有任务完成时返回
- FIRST_COMPLETED: 第一个任务完成时返回
- FIRST_EXCEPTION: 第一个任务异常时返回
Returns:
Tuple[Set, Set]: (已完成的任务集合, 未完成的任务集合)
"""
with self._lock:
futures = [task_info['future'] for task_info in self._task_dict.values()]
if not futures:
return set(), set()
done, not_done = wait(futures, timeout=timeout, return_when=return_when)
return done, not_done
def wait_for_any(self, timeout: Optional[float] = None) -> str:
"""
等待任意一个任务完成
Args:
timeout: 超时时间（秒）
Returns:
str: 完成的任务名称，如果超时返回None
"""
done, _ = self.wait_for_all(timeout=timeout, return_when=FIRST_COMPLETED)
if done:
# 找出是哪个任务完成了
for task_name, task_info in self._task_dict.items():
if task_info['future'] in done:
return task_name
return None
def get_task_status(self, task_name: str) -> Dict[str, any]:
"""
获取任务状态
Args:
task_name: 任务名称
Returns:
Dict: 包含任务状态信息的字典
"""
with self._lock:
if task_name not in self._task_dict:
raise KeyError(f"任务 '{task_name}' 不存在")
task_info = self._task_dict[task_name].copy()
# 添加额外信息
future = task_info['future']
task_info['done'] = future.done()
task_info['running'] = future.running()
task_info['cancelled'] = future.cancelled()
if task_info['done'] and not future.cancelled():
try:
task_info['result'] = future.result(timeout=0)
except Exception as e:
task_info['exception'] = str(e)
# 计算运行时间
if 'start_time' in task_info:
if 'end_time' in task_info:
task_info['duration'] = task_info['end_time'] - task_info['start_time']
else:
task_info['duration'] = time.time() - task_info['start_time']
return task_info
def get_all_task_status(self) -> Dict[str, Dict[str, any]]:
"""
获取所有任务状态
Returns:
Dict[str, Dict]: 所有任务的状态信息字典
"""
with self._lock:
task_names = list(self._task_dict.keys())
return {name: self.get_task_status(name) for name in task_names}
def stop_task(self, task_name: str, wait: bool = True):
"""停止指定任务"""
with self._lock:
if task_name not in self._task_dict:
raise KeyError(f"任务 '{task_name}' 不存在")
task_info = self._task_dict[task_name]
task_info['detect_video_stream'].stop_task()
# 取消future
future = task_info['future']
if not future.done():
future.cancel()
if wait and not future.done():
try:
future.result(timeout=5) # 等待最多5秒
except:
pass
task_info['status'] = 'stopped'
task_info['end_time'] = time.time()
def shutdown(self, wait: bool = True):
"""关闭线程池"""
# 先停止所有任务
task_names = list(self._task_dict.keys())
for task_name in task_names:
try:
self.stop_task(task_name, wait=False)
except:
pass
# 关闭线程池
self.executor.shutdown(wait=wait)
return self._results
return self._results