modify: 多线程管理器
This commit is contained in:
@@ -1,6 +1,9 @@
|
||||
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:
|
||||
@@ -11,7 +14,9 @@ class DetectVideoStream:
|
||||
frame_skip: int = 0,
|
||||
reconnect_attempts: int = 3,
|
||||
reconnect_delay: int = 5,
|
||||
is_save: bool = False):
|
||||
is_save: bool = False,
|
||||
output: str = "./result/output.mp4",
|
||||
logger: Optional[Logger] = None):
|
||||
"""
|
||||
处理视频流（视频文件或RTSP流）
|
||||
|
||||
@@ -26,6 +31,10 @@ class DetectVideoStream:
|
||||
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
|
||||
@@ -48,7 +57,7 @@ class DetectVideoStream:
|
||||
result_dir = r"results"
|
||||
# 清理旧的目录
|
||||
if os.path.exists(result_dir):
|
||||
print(f"清理旧的构建目录: {result_dir}")
|
||||
self.logger.info(f"清理旧的构建目录: {result_dir}")
|
||||
shutil.rmtree(result_dir, ignore_errors=True)
|
||||
|
||||
# 创建目录
|
||||
@@ -56,16 +65,14 @@ class DetectVideoStream:
|
||||
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
|
||||
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)
|
||||
|
||||
frame_shape = self.decoder.get_frame_shape()
|
||||
print(f"视频帧形状: {frame_shape}")
|
||||
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
|
||||
@@ -73,82 +80,83 @@ class DetectVideoStream:
|
||||
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}")
|
||||
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:
|
||||
print(f"读取帧失败，错误码: {ret}")
|
||||
self.logger.error(f"读取帧失败，错误码: {ret}")
|
||||
reconnect_count += 1
|
||||
|
||||
if reconnect_count > self.reconnect_attempts:
|
||||
print("视频流无法继续读取")
|
||||
self.logger.error("视频流无法继续读取")
|
||||
break
|
||||
|
||||
print(f"尝试重连 ({reconnect_count}/{self.reconnect_attempts})...")
|
||||
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
|
||||
print("重连成功，继续处理...")
|
||||
self.logger.info("重连成功，继续处理...")
|
||||
continue
|
||||
else:
|
||||
print("重连失败")
|
||||
self.logger.warning("重连失败")
|
||||
# time.sleep(0.01)
|
||||
|
||||
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)
|
||||
self.logger.info(f"处理第 {frame_count} 帧...")
|
||||
|
||||
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
|
||||
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(result_img)
|
||||
self.encoder.video_write(bmimg) # type: ignore
|
||||
# time.sleep(0.01)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n检测到键盘中断，停止处理")
|
||||
self.logger.info("\n检测到键盘中断，停止处理")
|
||||
result_flog = True
|
||||
except Exception as e:
|
||||
print(f"处理视频流时发生错误: {e}")
|
||||
self.logger.error(f"处理视频流时发生错误: {e}")
|
||||
result_flog = False
|
||||
|
||||
finally:
|
||||
# 清理资源
|
||||
self.decoder.release()
|
||||
if self.is_save:
|
||||
self.encoder.release()
|
||||
print("视频解码器已释放")
|
||||
self.logger.info("视频解码器已释放")
|
||||
|
||||
return result_flog
|
||||
|
||||
def stop_task(self):
|
||||
print("提前关闭")
|
||||
self.logger.info("提前关闭")
|
||||
self.stop_event.set()
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
import os
|
||||
import time
|
||||
import logging
|
||||
import sys
|
||||
from logging.handlers import RotatingFileHandler
|
||||
|
||||
|
||||
class LogManager:
|
||||
"""日志管理器"""
|
||||
|
||||
@staticmethod
|
||||
def setup_logging(task_name=None):
|
||||
"""配置日志系统"""
|
||||
# 创建日志目录
|
||||
log_dir = "logs"
|
||||
if not os.path.exists(log_dir):
|
||||
os.makedirs(log_dir)
|
||||
|
||||
# 生成日志文件名
|
||||
if task_name:
|
||||
log_file = os.path.join(log_dir, f"{task_name}_{time.strftime('%Y%m%d_%H%M%S')}.log")
|
||||
logger_name = task_name
|
||||
else:
|
||||
log_file = os.path.join(log_dir, f"detection_{time.strftime('%Y%m%d_%H%M%S')}.log")
|
||||
logger_name = "detection"
|
||||
|
||||
# 获取或创建特定名称的日志记录器
|
||||
logger = logging.getLogger(logger_name)
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
# 清除现有处理器（避免重复添加）
|
||||
if logger.handlers:
|
||||
logger.handlers.clear()
|
||||
|
||||
# 文件处理器 - 轮转日志，最大10MB，保留5个备份
|
||||
file_handler = RotatingFileHandler(
|
||||
log_file,
|
||||
maxBytes=10*1024*1024,
|
||||
backupCount=5,
|
||||
encoding='utf-8'
|
||||
)
|
||||
file_formatter = logging.Formatter(
|
||||
'%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||
datefmt='%Y-%m-%d %H:%M:%S'
|
||||
)
|
||||
file_handler.setFormatter(file_formatter)
|
||||
logger.addHandler(file_handler)
|
||||
|
||||
# 控制台处理器 - 只输出警告及以上级别
|
||||
console_handler = logging.StreamHandler(sys.stdout)
|
||||
console_formatter = logging.Formatter(
|
||||
'%(levelname)s - %(message)s'
|
||||
)
|
||||
console_handler.setFormatter(console_formatter)
|
||||
console_handler.setLevel(logging.WARNING) # 只显示WARNING及以上
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
# 防止日志传递给父logger（如根logger）
|
||||
logger.propagate = False
|
||||
|
||||
return log_file, logger # 同时返回logger实例
|
||||
@@ -0,0 +1,20 @@
|
||||
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
|
||||
@@ -0,0 +1,110 @@
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
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
|
||||
|
||||
|
||||
@dataclass
|
||||
class TaskArgs:
|
||||
task_name: str = field(
|
||||
metadata={"description": "任务名称"}
|
||||
)
|
||||
algorithms: list = field(
|
||||
metadata={"description": "算法列表"}
|
||||
)
|
||||
input: str = field(
|
||||
metadata={"description": "输入(一般为rtsp流地址或视频文件)"}
|
||||
)
|
||||
output: str = field(
|
||||
metadata={"description": "输出(一般为rtsp流地址或视频文件)"}
|
||||
)
|
||||
is_vfs: bool = field(
|
||||
default=False,
|
||||
metadata={"description": "是否是虚拟文件系统(一般在生成环境中使用)"}
|
||||
)
|
||||
dev_id: int = field(
|
||||
default=0,
|
||||
metadata={"description": "设备ID"}
|
||||
)
|
||||
conf_thresh: float = field(
|
||||
default=0.25,
|
||||
metadata={"description": "置信度阈值"}
|
||||
)
|
||||
nms_thresh: float = field(
|
||||
default=0.25,
|
||||
metadata={"description": "NMS阈值"}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> "TaskArgs":
|
||||
data["algorithms"] = [item.strip() for item in data["algorithms"].split('*') if item.strip()]
|
||||
print(data["algorithms"])
|
||||
return cls(**data)
|
||||
|
||||
|
||||
class TaskManager:
|
||||
def __init__(self,
|
||||
max_workers=None,
|
||||
thread_name_prefix:str="task_"):
|
||||
self.executor = ThreadPoolExecutor(
|
||||
max_workers=max_workers,
|
||||
thread_name_prefix=thread_name_prefix
|
||||
)
|
||||
self._task_dict: Dict[str, DetectVideoStream] = {}
|
||||
self._results: List[bool] = []
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def submit(self, task_args: TaskArgs):
|
||||
"""提交任务"""
|
||||
if task_args.task_name in self._task_dict.keys():
|
||||
print(f"'{task_args.task_name}' 任务已存在")
|
||||
return
|
||||
|
||||
log_file, logger = LogManager.setup_logging(task_args.task_name)
|
||||
logger.info(f"日志文件 {log_file}")
|
||||
algorithms = [get_detector(algorithm, task_args.is_vfs) for algorithm in task_args.algorithms]
|
||||
algorithms = [
|
||||
algorithm[0](
|
||||
algorithm[1],
|
||||
task_args.is_vfs,
|
||||
logger=logger,
|
||||
dev_id=task_args.dev_id,
|
||||
conf_thresh=task_args.conf_thresh,
|
||||
nms_thresh=task_args.nms_thresh)
|
||||
for
|
||||
algorithm
|
||||
in
|
||||
algorithms
|
||||
]
|
||||
detect_video_stream = DetectVideoStream(algorithms,
|
||||
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)
|
||||
return future
|
||||
|
||||
def _task_callback(self, future):
|
||||
"""任务完成回调"""
|
||||
try:
|
||||
result = future.result()
|
||||
with self._lock:
|
||||
self._results.append(result)
|
||||
except Exception as e:
|
||||
print(f"任务执行出错: {e}")
|
||||
|
||||
def stop_task(self, task_name):
|
||||
self._task_dict[task_name].stop_task()
|
||||
|
||||
def shutdown(self, wait=True):
|
||||
"""关闭线程池"""
|
||||
self.executor.shutdown(wait=wait)
|
||||
return self._results
|
||||
|
||||
Reference in New Issue
Block a user