added VisionImageFrame and VisionImageFrameAggregator
This commit is contained in:
@@ -5,7 +5,7 @@ import os
|
||||
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from dailyai.pipeline.aggregators import FrameProcessor, UserResponseAggregator
|
||||
from dailyai.pipeline.aggregators import FrameProcessor, UserResponseAggregator, VisionImageFrameAggregator
|
||||
|
||||
from dailyai.pipeline.frames import Frame, TextFrame, UserImageRequestFrame
|
||||
from dailyai.pipeline.pipeline import Pipeline
|
||||
@@ -59,6 +59,8 @@ async def main(room_url: str, token):
|
||||
|
||||
image_requester = UserImageRequester()
|
||||
|
||||
vision_aggregator = VisionImageFrameAggregator()
|
||||
|
||||
moondream = MoondreamService()
|
||||
|
||||
tts = ElevenLabsTTSService(
|
||||
@@ -73,7 +75,7 @@ async def main(room_url: str, token):
|
||||
transport.render_participant_video(participant["id"], framerate=0)
|
||||
image_requester.set_participant_id(participant["id"])
|
||||
|
||||
pipeline = Pipeline([user_response, image_requester, moondream, tts])
|
||||
pipeline = Pipeline([user_response, image_requester, vision_aggregator, moondream, tts])
|
||||
|
||||
await transport.run(pipeline)
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ from dailyai.pipeline.frames import (
|
||||
EndFrame,
|
||||
EndPipeFrame,
|
||||
Frame,
|
||||
ImageFrame,
|
||||
LLMMessagesFrame,
|
||||
LLMResponseEndFrame,
|
||||
LLMResponseStartFrame,
|
||||
@@ -14,6 +15,7 @@ from dailyai.pipeline.frames import (
|
||||
TranscriptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
VisionImageFrame,
|
||||
)
|
||||
from dailyai.pipeline.pipeline import Pipeline
|
||||
from dailyai.services.ai_services import AIService
|
||||
@@ -463,3 +465,37 @@ class GatedAggregator(FrameProcessor):
|
||||
self.accumulator = []
|
||||
else:
|
||||
self.accumulator.append(frame)
|
||||
|
||||
|
||||
class VisionImageFrameAggregator(FrameProcessor):
|
||||
"""This aggregator waits for a consecutive TextFrame and an
|
||||
ImageFrame. After the ImageFrame arrives it will output a VisionImageFrame.
|
||||
|
||||
>>> from dailyai.pipeline.frames import ImageFrame
|
||||
|
||||
>>> async def print_frames(aggregator, frame):
|
||||
... async for frame in aggregator.process_frame(frame):
|
||||
... print(frame)
|
||||
|
||||
>>> aggregator = VisionImageFrameAggregator()
|
||||
>>> asyncio.run(print_frames(aggregator, TextFrame("What do you see?")))
|
||||
>>> asyncio.run(print_frames(aggregator, ImageFrame(image=bytes([]), size=(0, 0))))
|
||||
VisionImageFrame, text: What do you see?, image size: 0x0, buffer size: 0 B
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._describe_text = None
|
||||
|
||||
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
|
||||
if isinstance(frame, TextFrame):
|
||||
self._describe_text = frame.text
|
||||
elif isinstance(frame, ImageFrame):
|
||||
if self._describe_text:
|
||||
yield VisionImageFrame(self._describe_text, frame.image, frame.size)
|
||||
self._describe_text = None
|
||||
else:
|
||||
yield frame
|
||||
else:
|
||||
yield frame
|
||||
|
||||
@@ -79,8 +79,10 @@ class ImageFrame(Frame):
|
||||
|
||||
@dataclass()
|
||||
class URLImageFrame(ImageFrame):
|
||||
"""An image. Will be shown by the transport if the transport's camera is
|
||||
enabled."""
|
||||
"""An image with an associated URL. Will be shown by the transport if the
|
||||
transport's camera is enabled.
|
||||
|
||||
"""
|
||||
url: str | None
|
||||
|
||||
def __init__(self, url, image, size):
|
||||
@@ -91,6 +93,22 @@ class URLImageFrame(ImageFrame):
|
||||
return f"{self.__class__.__name__}, url: {self.url}, image size: {self.size[0]}x{self.size[1]}, buffer size: {len(self.image)} B"
|
||||
|
||||
|
||||
@dataclass()
|
||||
class VisionImageFrame(ImageFrame):
|
||||
"""An image with an associated text to ask for a description of it. Will be shown by the
|
||||
transport if the transport's camera is enabled.
|
||||
|
||||
"""
|
||||
text: str | None
|
||||
|
||||
def __init__(self, text, image, size):
|
||||
super().__init__(image, size)
|
||||
self.text = text
|
||||
|
||||
def __str__(self):
|
||||
return f"{self.__class__.__name__}, text: {self.text}, image size: {self.size[0]}x{self.size[1]}, buffer size: {len(self.image)} B"
|
||||
|
||||
|
||||
@dataclass()
|
||||
class UserImageFrame(ImageFrame):
|
||||
"""An image associated to a user. Will be shown by the transport if the transport's camera is
|
||||
|
||||
@@ -15,6 +15,7 @@ from dailyai.pipeline.frames import (
|
||||
TextFrame,
|
||||
TranscriptionFrame,
|
||||
URLImageFrame,
|
||||
VisionImageFrame,
|
||||
)
|
||||
|
||||
from abc import abstractmethod
|
||||
@@ -108,19 +109,13 @@ class VisionService(AIService):
|
||||
self._describe_text = None
|
||||
|
||||
@abstractmethod
|
||||
async def run_vision(self, describe_text: str, frame: ImageFrame) -> str:
|
||||
async def run_vision(self, frame: VisionImageFrame) -> str:
|
||||
pass
|
||||
|
||||
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
|
||||
if isinstance(frame, TextFrame):
|
||||
self._describe_text = frame.text
|
||||
elif isinstance(frame, ImageFrame):
|
||||
if self._describe_text:
|
||||
description = await self.run_vision(self._describe_text, frame)
|
||||
self._describe_text = None
|
||||
yield TextFrame(description)
|
||||
else:
|
||||
yield frame
|
||||
if isinstance(frame, VisionImageFrame):
|
||||
description = await self.run_vision(frame)
|
||||
yield TextFrame(description)
|
||||
else:
|
||||
yield frame
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from dailyai.pipeline.frames import ImageFrame
|
||||
from dailyai.pipeline.frames import ImageFrame, VisionImageFrame
|
||||
from dailyai.services.ai_services import VisionService
|
||||
|
||||
from PIL import Image
|
||||
@@ -42,11 +42,11 @@ class MoondreamService(VisionService):
|
||||
).to(device=device, dtype=dtype)
|
||||
self._model.eval()
|
||||
|
||||
async def run_vision(self, describe_text: str, frame: ImageFrame) -> str:
|
||||
async def run_vision(self, frame: VisionImageFrame) -> str:
|
||||
image = Image.frombytes("RGB", (frame.size[0], frame.size[1]), frame.image)
|
||||
image_embeds = self._model.encode_image(image)
|
||||
description = self._model.answer_question(
|
||||
image_embeds=image_embeds,
|
||||
question=describe_text,
|
||||
question=frame.text,
|
||||
tokenizer=self._tokenizer)
|
||||
return description
|
||||
|
||||
Reference in New Issue
Block a user