65 lines
2.5 KiB
Python
65 lines
2.5 KiB
Python
import asyncio
|
|
from typing import AsyncGenerator, List
|
|
from dailyai.pipeline.frame_processor import FrameProcessor
|
|
|
|
from dailyai.pipeline.frames import EndParallelPipeQueueFrame, EndStreamQueueFrame, QueueFrame
|
|
|
|
"""
|
|
This class manages a pipe of FrameProcessors, and runs them in sequence. The "source"
|
|
and "sink" queues are managed by the caller. You can use this class stand-alone to
|
|
perform specialized processing, or you can use the Transport's run_pipeline method to
|
|
instantiate and run a pipeline with the Transport's sink and source queues.
|
|
"""
|
|
|
|
class Pipeline:
|
|
|
|
def __init__(
|
|
self,
|
|
processors: List[FrameProcessor],
|
|
source: asyncio.Queue | None = None,
|
|
sink: asyncio.Queue[QueueFrame] | None = None,
|
|
):
|
|
self.processors = processors
|
|
self.source: asyncio.Queue[QueueFrame] | None = source
|
|
self.sink: asyncio.Queue[QueueFrame] | None = sink
|
|
|
|
def set_source(self, source: asyncio.Queue[QueueFrame]):
|
|
self.source = source
|
|
|
|
def set_sink(self, sink: asyncio.Queue[QueueFrame]):
|
|
self.sink = sink
|
|
|
|
async def get_next_source_frame(self) -> AsyncGenerator[QueueFrame, None]:
|
|
if self.source is None:
|
|
raise ValueError("Source queue not set")
|
|
yield await self.source.get()
|
|
|
|
async def run_pipeline(self):
|
|
if self.source is None or self.sink is None:
|
|
raise ValueError("Source or sink queue not set")
|
|
|
|
try:
|
|
while True:
|
|
frame_generators = [self.get_next_source_frame()]
|
|
for processor in self.processors:
|
|
next_frame_generators = []
|
|
for frame_generator in frame_generators:
|
|
async for frame in frame_generator:
|
|
next_frame_generators.append(processor.process_frame(frame))
|
|
frame_generators = next_frame_generators
|
|
|
|
for frame_generator in frame_generators:
|
|
async for frame in frame_generator:
|
|
await self.sink.put(frame)
|
|
if isinstance(
|
|
frame, EndStreamQueueFrame
|
|
) or isinstance(
|
|
frame, EndParallelPipeQueueFrame
|
|
):
|
|
return
|
|
except asyncio.CancelledError:
|
|
# this means there's been an interruption, do any cleanup necessary here.
|
|
for processor in self.processors:
|
|
await processor.interrupted()
|
|
pass
|