added VisionImageFrame and VisionImageFrameAggregator

This commit is contained in:
Aleix Conchillo Flaqué
2024-04-10 09:18:54 -07:00
parent 2f9899af5a
commit 3c20f9153d
5 changed files with 68 additions and 17 deletions

View File

@@ -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)

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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