zimoyin
94 lines
4.1 KiB
Python
94 lines
4.1 KiB
Python
from loguru import logger
|
|
from typing import List, Iterator
|
|
import asyncio
|
|
from .base_processor import BaseProcessor
|
|
from .pipeline_data import PipelineData
|
|
|
|
|
|
class Pipeline:
|
|
"""管道核心类，管理多个处理器，负责数据流转"""
|
|
|
|
def __init__(self):
|
|
self.processors: List[BaseProcessor] = [] # 同步处理器列表
|
|
self.async_processors: List[BaseProcessor] = [] # 异步处理器列表
|
|
# 初始化日志记录器
|
|
self.logger = logger
|
|
self.logger.debug("管道已初始化")
|
|
|
|
def add_processor(self, processor: BaseProcessor) -> None:
|
|
"""向管道添加同步处理器（按添加顺序执行）"""
|
|
self.processors.append(processor)
|
|
self.logger.info(f"管道已添加同步处理器: {processor}")
|
|
|
|
def add_processor_async(self, processor: BaseProcessor) -> None:
|
|
"""向管道添加异步处理器（异步执行，不阻塞数据流）"""
|
|
self.async_processors.append(processor)
|
|
self.logger.info(f"管道已添加异步处理器: {processor}")
|
|
|
|
def remove_processor(self, processor_name: str) -> bool:
|
|
"""移除指定名称的处理器"""
|
|
# 检查同步处理器列表
|
|
for idx, p in enumerate(self.processors):
|
|
if p.name == processor_name:
|
|
self.processors.pop(idx)
|
|
self.logger.info(f"管道已移除同步处理器: {p}")
|
|
return True
|
|
# 检查异步处理器列表
|
|
for idx, p in enumerate(self.async_processors):
|
|
if p.name == processor_name:
|
|
self.async_processors.pop(idx)
|
|
self.logger.info(f"管道已移除异步处理器: {p}")
|
|
return True
|
|
self.logger.warning(f"未找到处理器: {processor_name}")
|
|
return False
|
|
|
|
def run(self, data_iterator: Iterator[PipelineData]) -> Iterator[PipelineData]:
|
|
"""
|
|
运行管道：迭代输入数据，依次经过所有同步处理器，同时异步执行异步处理器
|
|
:param data_iterator: 数据源迭代器（如视频+YOLO的迭代器）
|
|
:return: 处理后的数据包迭代器
|
|
"""
|
|
|
|
# 创建异步任务列表，用于管理异步处理器
|
|
async def run_async_processors(data: PipelineData):
|
|
tasks = []
|
|
for async_processor in self.async_processors:
|
|
task = asyncio.create_task(self._process_async(async_processor, data))
|
|
tasks.append(task)
|
|
if tasks:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
for data in data_iterator:
|
|
# 依次执行每个同步处理器的处理逻辑
|
|
for processor in self.processors:
|
|
try:
|
|
data = processor.process(data)
|
|
except Exception as e:
|
|
self.logger.error(f"同步处理器 {processor.name} 执行出错: {e}")
|
|
|
|
# 异步执行异步处理器，不阻塞数据流
|
|
if self.async_processors:
|
|
# 在独立的事件循环中运行异步处理器（如果可用）
|
|
try:
|
|
# 尝试获取当前事件循环
|
|
loop = asyncio.get_running_loop()
|
|
# 在当前事件循环中调度异步处理器
|
|
asyncio.create_task(run_async_processors(data.copy() if hasattr(data, 'copy') else data))
|
|
except RuntimeError:
|
|
# 如果没有运行中的事件循环，则创建一个新的
|
|
asyncio.run(run_async_processors(data.copy() if hasattr(data, 'copy') else data))
|
|
|
|
yield data # 返回处理后的数据包
|
|
|
|
async def _process_async(self, processor, data):
|
|
"""异步执行单个处理器"""
|
|
try:
|
|
# 对数据进行异步处理（如果处理器支持）
|
|
if hasattr(processor, 'process_async'):
|
|
await processor.process_async(data)
|
|
else:
|
|
# 如果处理器不支持异步处理，则同步处理
|
|
processor.process(data)
|
|
except Exception as e:
|
|
self.logger.error(f"异步处理器 {processor.name} 执行出错: {e}")
|