File size: 5,611 Bytes
162d5c8
 
 
 
 
cf5e7f4
162d5c8
 
62a21bd
cf5e7f4
ac35a95
62a21bd
 
162d5c8
 
62a21bd
d91a673
62a21bd
e2846c4
09ede70
e2846c4
ac35a95
 
ad67495
cf5e7f4
 
e2846c4
149eeaf
e2846c4
149eeaf
e2846c4
ad67495
162d5c8
 
9ed41df
 
cf5e7f4
62a21bd
149eeaf
 
cf5e7f4
 
 
 
 
 
 
149eeaf
cf5e7f4
 
 
 
 
 
 
 
 
 
 
9ed41df
62a21bd
 
cf5e7f4
162d5c8
 
 
 
 
62a21bd
e2846c4
 
62a21bd
e2846c4
62a21bd
 
 
 
 
 
 
 
 
 
 
 
0d27fd9
62a21bd
 
 
162d5c8
 
d91a673
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
162d5c8
768e92a
ac35a95
 
62a21bd
768e92a
 
ac35a95
d91a673
 
cf5e7f4
d91a673
 
cf5e7f4
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
from typing import List
import av
import asyncio
from collections import deque
import threading
import cv2

import numpy as np
import ray
from ray.util.queue import Queue
from webrtc_av_queue_actor import WebRtcAVQueueActor
import pydub
import torch

class StreamlitAVQueue:
    def __init__(self, audio_bit_rate=16000):
        self._output_channels = 2
        self._audio_bit_rate = audio_bit_rate
        self._listening = True
        self._looking = False
        self._lock = threading.Lock()
        self.queue_actor = WebRtcAVQueueActor.options(
            name="WebRtcAVQueueActor", 
            get_if_exists=True,
            ).remote()
        self._out_video_frame = None 
        
    def set_looking_listening(self, looking, listening: bool):
        with self._lock:
            self._looking = looking
            self._listening = listening
            
    async def queued_video_frames_callback(
                self,
                frames: List[av.VideoFrame],
            ) -> av.VideoFrame:
        updated_frames = []
        try:
            with self._lock:
                should_look = self._looking
            next_out_video_frame = await self.queue_actor.get_out_video_frame.remote()
            if next_out_video_frame is not None:
                self._out_video_frame = next_out_video_frame
            for i, frame in enumerate(frames):
                user_image = frame.to_ndarray(format="rgb24")
                if should_look:
                    shared_tensor_ref = ray.put(user_image)
                    await self.queue_actor.enqueue_in_video_frame.remote(shared_tensor_ref)
                if self._out_video_frame is not None:
                    frame = self._out_video_frame
                    # resize user image to 1/4 size
                    user_frame = cv2.resize(user_image, (user_image.shape[1]//4, user_image.shape[0]//4), interpolation=cv2.INTER_AREA)
                    x_user = 0
                    y_user = frame.shape[0] - user_frame.shape[0]
                    final_frame = frame.copy()
                    final_frame[y_user:y_user+user_frame.shape[0], x_user:x_user+user_frame.shape[1]] = user_frame
                    frame = av.VideoFrame.from_ndarray(final_frame, format="rgb24")

                updated_frames.append(frame)
                # print (f"tesnor len: {len(shared_tensor)}, tensor shape: {shared_tensor.shape}, tensor type:{shared_tensor.dtype} tensor ref: {shared_tensor_ref}")
        except Exception as e:
            print (e)
        return updated_frames

    async def queued_audio_frames_callback(
                self,
                frames: List[av.AudioFrame],
            ) -> av.AudioFrame:
        try:
            with self._lock:
                should_listed = self._listening
            sound_chunk = pydub.AudioSegment.empty()
            if len(frames) > 0 and should_listed:
                for frame in frames:
                    sound = pydub.AudioSegment(
                        data=frame.to_ndarray().tobytes(),
                        sample_width=frame.format.bytes,
                        frame_rate=frame.sample_rate,
                        channels=len(frame.layout.channels),
                    )
                    sound = sound.set_channels(1)
                    sound = sound.set_frame_rate(self._audio_bit_rate)
                    sound_chunk += sound
                shared_buffer = np.array(sound_chunk.get_array_of_samples())
                shared_buffer_ref = ray.put(shared_buffer)
                await self.queue_actor.enqueue_in_audio_frame.remote(shared_buffer_ref)
        except Exception as e:
            print (e)
            
        # return empty frames to avoid echo
        new_frames = []
        try:
            for frame in frames:
                required_samples = frame.samples
                # print (f"frame: {frame.format.name}, {frame.layout.name}, {frame.sample_rate}, {frame.samples}")
                assert frame.format.bytes == 2
                assert frame.format.name == 's16'
                frame_as_bytes = await self.queue_actor.get_out_audio_frame.remote()
                if frame_as_bytes:
                    # print(f"frame_as_bytes: {len(frame_as_bytes)}")
                    assert len(frame_as_bytes) == frame.samples * frame.format.bytes
                    samples = np.frombuffer(frame_as_bytes, dtype=np.int16)
                else:
                    samples = np.zeros((required_samples * 2 * 1), dtype=np.int16)
                if self._output_channels == 2:
                    samples = np.vstack((samples, samples)).reshape((-1,), order='F')
                samples = samples.reshape(1, -1)
                layout = 'stereo' if self._output_channels == 2 else 'mono'
                new_frame = av.AudioFrame.from_ndarray(samples, format='s16', layout=layout)
                new_frame.sample_rate = frame.sample_rate
                new_frames.append(new_frame)
        except Exception as e:
            print (e)
        return new_frames

    async def get_in_audio_frames_async(self) -> List[av.AudioFrame]:
        shared_buffers = await self.queue_actor.get_in_audio_frames.remote()
        return shared_buffers

    async def get_video_frames_async(self) -> List[av.AudioFrame]:
        shared_tensors = await self.queue_actor.get_in_video_frames.remote()
        return shared_tensors
    
    def get_out_audio_queue(self)->Queue:
        return self.queue_actor.get_out_audio_queue.remote()
    
    def get_out_video_queue(self)->Queue:
        return self.queue_actor.get_out_video_queue.remote()