mirror of
https://github.com/facefusion/facefusion.git
synced 2026-07-15 06:27:21 +02:00
Resolve final stream TODOs (#1155)
* refactor a lot * refactor a lot * fix test * fix test
This commit is contained in:
Henry Ruhs
@@ -1,5 +1,3 @@
|
||||
import ctypes
|
||||
import time
|
||||
from functools import partial
|
||||
from queue import Queue
|
||||
from typing import Optional, Tuple
|
||||
@@ -13,6 +11,7 @@ from facefusion.types import AudioCodec, AudioFrame, Buffer, OpusDecoder, RtcPee
|
||||
|
||||
|
||||
def run_audio_encode_loop(rtc_peer : RtcPeer, audio_queue : Queue[Tuple[Time, AudioFrame]]) -> None:
|
||||
audio_codec = rtc_peer.get('audio').get('codec')
|
||||
temp_audio_time, temp_audio_frame = audio_queue.get()
|
||||
audio_encoder = opus_encoder.create(48000, 2)
|
||||
|
||||
@@ -20,7 +19,7 @@ def run_audio_encode_loop(rtc_peer : RtcPeer, audio_queue : Queue[Tuple[Time, Au
|
||||
audio_buffer = opus_encoder.encode(audio_encoder, temp_audio_frame.tobytes(), 960)
|
||||
|
||||
if audio_buffer:
|
||||
audio_timestamp = int(temp_audio_time * 48000)
|
||||
audio_timestamp = rtc.convert_time_to_timestamp(audio_codec, temp_audio_time)
|
||||
rtc.send_audio(rtc_peer, audio_buffer, audio_timestamp)
|
||||
|
||||
temp_audio_time, temp_audio_frame = audio_queue.get()
|
||||
@@ -59,10 +58,10 @@ def destroy_audio_decoder(audio_codec : AudioCodec, audio_decoder : OpusDecoder)
|
||||
opus_decoder.destroy(audio_decoder)
|
||||
|
||||
|
||||
def handle_audio_frame(audio_codec : AudioCodec, audio_decoder : OpusDecoder, audio_queue : Queue[Tuple[Time, AudioFrame]], track : int, data : ctypes.c_void_p, size : int, info : ctypes.c_void_p, pointer : ctypes.c_void_p) -> None:
|
||||
audio_buffer = ctypes.string_at(data, size)
|
||||
def handle_audio_frame(audio_codec : AudioCodec, audio_decoder : OpusDecoder, audio_queue : Queue[Tuple[Time, AudioFrame]], audio_buffer : Buffer, audio_timestamp : int) -> None:
|
||||
audio_frame = decode_audio_frame(audio_codec, audio_decoder, audio_buffer)
|
||||
|
||||
if audio_frame:
|
||||
temp_audio_frame = numpy.frombuffer(audio_frame, dtype = numpy.float32)
|
||||
audio_queue.put((time.monotonic(), temp_audio_frame))
|
||||
audio_frame = numpy.frombuffer(audio_frame, dtype = numpy.float32)
|
||||
audio_time = rtc.convert_timestamp_to_time(audio_codec, audio_timestamp)
|
||||
audio_queue.put((audio_time, audio_frame))
|
||||
|
||||
@@ -10,7 +10,7 @@ def create_receive_event(track : int, frame_handler : FrameHandler) -> threading
|
||||
datachannel_library = datachannel_module.create_static_library()
|
||||
receive_event = threading.Event()
|
||||
|
||||
frame_callback = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p)(frame_handler)
|
||||
frame_callback = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p)(partial(dispatch_frame, frame_handler))
|
||||
close_callback = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_void_p)(partial(dispatch_event, receive_event))
|
||||
datachannel_library.rtcSetFrameCallback(track, frame_callback)
|
||||
datachannel_library.rtcSetClosedCallback(track, close_callback)
|
||||
@@ -20,5 +20,11 @@ def create_receive_event(track : int, frame_handler : FrameHandler) -> threading
|
||||
return receive_event
|
||||
|
||||
|
||||
def dispatch_frame(frame_handler : FrameHandler, track : int, data : ctypes.c_void_p, size : int, info : ctypes.c_void_p, pointer : ctypes.c_void_p) -> None:
|
||||
frame_buffer = ctypes.string_at(data, size)
|
||||
frame_timestamp = ctypes.cast(info, ctypes.POINTER(ctypes.c_uint32)).contents.value
|
||||
frame_handler(frame_buffer, frame_timestamp)
|
||||
|
||||
|
||||
def dispatch_event(event : threading.Event, track : int, pointer : ctypes.c_void_p) -> None:
|
||||
event.set()
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
import ctypes
|
||||
import time
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from functools import partial
|
||||
from queue import Queue
|
||||
@@ -26,32 +24,26 @@ def run_video_encode_loop(rtc_peer : RtcPeer, video_queue : Queue[Tuple[Time, Fu
|
||||
temp_resolution : Resolution = video_resolution
|
||||
temp_bitrate : BitRate = 8000
|
||||
video_encoder = create_video_encoder(video_codec, temp_resolution, temp_bitrate)
|
||||
temp_video_time = video_time
|
||||
frame_index = 0
|
||||
|
||||
while video_buffer:
|
||||
encode_start = time.monotonic()
|
||||
sender_bitrate = calculate_sender_bitrate(rtc_peer, temp_bitrate)
|
||||
sender_bitrate = rtc_peer.get('sender_bitrate').value
|
||||
|
||||
if video_resolution[0] - temp_resolution[0] or video_resolution[1] - temp_resolution[1]:
|
||||
temp_resolution = video_resolution
|
||||
update_video_encoder_resolution(video_codec, video_encoder, temp_resolution)
|
||||
|
||||
if sender_bitrate - temp_bitrate:
|
||||
if sender_bitrate > 0 and sender_bitrate - temp_bitrate:
|
||||
temp_bitrate = sender_bitrate
|
||||
update_video_encoder_bitrate(video_codec, video_encoder, temp_bitrate)
|
||||
|
||||
__video_buffer__ = encode_video_frame(video_codec, video_encoder, video_buffer, temp_resolution, frame_index)
|
||||
|
||||
if __video_buffer__:
|
||||
video_timestamp = int(video_time * 90000)
|
||||
video_timestamp = rtc.convert_time_to_timestamp(video_codec, video_time)
|
||||
rtc.send_video(rtc_peer, __video_buffer__, video_timestamp)
|
||||
|
||||
encode_time = time.monotonic() - encode_start
|
||||
frame_interval = video_time - temp_video_time
|
||||
temp_video_time = video_time
|
||||
|
||||
receiver_bitrate = calculate_receiver_bitrate(rtc_peer, encode_time, frame_interval)
|
||||
receiver_bitrate = rtc_peer.get('receiver_bitrate').value
|
||||
rtc.adapt_receiver_bitrate(rtc_peer, receiver_bitrate)
|
||||
frame_index += 1
|
||||
|
||||
@@ -86,31 +78,6 @@ def process_video_frame(input_vision_frame : VisionFrame) -> BufferPack:
|
||||
return BufferPack(buffer = output_buffer, resolution = output_resolution)
|
||||
|
||||
|
||||
def calculate_receiver_bitrate(rtc_peer : RtcPeer, encode_time : float, frame_interval : float) -> BitRate:
|
||||
min_bitrate : BitRate = 500
|
||||
max_bitrate : BitRate = 8000
|
||||
bitrate : BitRate = rtc_peer.get('receiver_bitrate').value
|
||||
|
||||
if frame_interval > 0:
|
||||
scale = frame_interval / encode_time
|
||||
bitrate = int(bitrate * scale)
|
||||
bitrate = max(min_bitrate, min(max_bitrate, bitrate))
|
||||
|
||||
return bitrate
|
||||
|
||||
|
||||
#todo: does not feel final as this is an clamp and not calculate
|
||||
def calculate_sender_bitrate(rtc_peer : RtcPeer, bitrate : BitRate) -> BitRate:
|
||||
min_bitrate : BitRate = 500
|
||||
max_bitrate : BitRate = 8000
|
||||
peer_bitrate : BitRate = rtc_peer.get('sender_bitrate').value
|
||||
|
||||
if peer_bitrate > 0:
|
||||
bitrate = max(min_bitrate, min(max_bitrate, peer_bitrate))
|
||||
|
||||
return bitrate
|
||||
|
||||
|
||||
def decode_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, input_buffer : Buffer) -> Optional[VisionFrame]:
|
||||
if video_codec == 'av1':
|
||||
aom_pointer = aom_decoder.decode(video_decoder, input_buffer)
|
||||
@@ -199,11 +166,10 @@ def update_video_encoder_bitrate(video_codec : VideoCodec, video_encoder : VpxEn
|
||||
return False
|
||||
|
||||
|
||||
#todo: we can remove the dead args or pass audio buffer
|
||||
def handle_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, video_queue : Queue[Tuple[Time, Future[BufferPack]]], video_executor : ThreadPoolExecutor, track : int, data : ctypes.c_void_p, size : int, info : ctypes.c_void_p, pointer : ctypes.c_void_p) -> None:
|
||||
video_buffer = ctypes.string_at(data, size)
|
||||
def handle_video_frame(video_codec : VideoCodec, video_decoder : VpxDecoder | AomDecoder, video_queue : Queue[Tuple[Time, Future[BufferPack]]], video_executor : ThreadPoolExecutor, video_buffer : Buffer, video_timestamp : int) -> None:
|
||||
vision_frame = decode_video_frame(video_codec, video_decoder, video_buffer)
|
||||
|
||||
if numpy.any(vision_frame) and video_queue.qsize() < video_queue.maxsize:
|
||||
video_future = video_executor.submit(process_video_frame, vision_frame)
|
||||
video_queue.put((time.monotonic(), video_future))
|
||||
video_time = rtc.convert_timestamp_to_time(video_codec, video_timestamp)
|
||||
video_queue.put((video_time, video_future))
|
||||
|
||||
+13
-1
@@ -2,7 +2,7 @@ import ctypes
|
||||
from typing import List, Optional
|
||||
|
||||
from facefusion.libraries import datachannel as datachannel_module
|
||||
from facefusion.types import AudioCodec, BitRate, Buffer, MediaDirection, PeerConnection, RtcAudioTrack, RtcPeer, RtcTrackInit, RtcVideoTrack, SdpAnswer, SdpOffer, VideoCodec
|
||||
from facefusion.types import AudioCodec, BitRate, Buffer, MediaDirection, PeerConnection, RtcAudioTrack, RtcPeer, RtcTrackInit, RtcVideoTrack, SdpAnswer, SdpOffer, Time, VideoCodec
|
||||
|
||||
|
||||
def create_peer_connection() -> PeerConnection:
|
||||
@@ -239,6 +239,18 @@ def get_payload_type(sdp_offer : SdpOffer, codec : AudioCodec | VideoCodec) -> i
|
||||
return 0
|
||||
|
||||
|
||||
def convert_time_to_timestamp(codec : AudioCodec | VideoCodec, time : Time) -> int:
|
||||
if codec == 'opus':
|
||||
return int(time * 48000)
|
||||
return int(time * 90000)
|
||||
|
||||
|
||||
def convert_timestamp_to_time(codec : AudioCodec | VideoCodec, timestamp : int) -> Time:
|
||||
if codec == 'opus':
|
||||
return timestamp / 48000
|
||||
return timestamp / 90000
|
||||
|
||||
|
||||
@ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_uint, ctypes.c_void_p)
|
||||
def handle_sender_bitrate(_ : int, bitrate : BitRate, pointer : int) -> None:
|
||||
ctypes.cast(pointer, ctypes.POINTER(ctypes.c_uint)).contents.value = bitrate // 1000
|
||||
|
||||
+2
-2
@@ -102,8 +102,6 @@ VideoCodec : TypeAlias = Literal['av1', 'vp8', 'vp9']
|
||||
AomVideoCodec : TypeAlias = Literal['av1']
|
||||
VxpVideoCodec : TypeAlias = Literal['vp8', 'vp9']
|
||||
|
||||
FrameHandler : TypeAlias = Callable[..., None]
|
||||
|
||||
AomEncoder : TypeAlias = ctypes.Array[ctypes.c_char]
|
||||
AomDecoder : TypeAlias = ctypes.Array[ctypes.c_char]
|
||||
OpusEncoder : TypeAlias = ctypes.c_void_p
|
||||
@@ -118,6 +116,8 @@ BufferPack = TypedDict('BufferPack',
|
||||
'resolution' : Resolution
|
||||
})
|
||||
|
||||
FrameHandler : TypeAlias = Callable[[Buffer, int], None]
|
||||
|
||||
Args : TypeAlias = Dict[str, Any]
|
||||
|
||||
Choice : TypeAlias = Union[int | str]
|
||||
|
||||
@@ -14,7 +14,7 @@ from facefusion.download import conditional_download
|
||||
from facefusion.ffmpeg import read_audio_buffer
|
||||
from facefusion.hash_helper import create_hash
|
||||
from facefusion.libraries import datachannel as datachannel_module, opus as opus_module
|
||||
from facefusion.types import AudioCodec, AudioFrame, FrameHandler, RtcPeer, RtcPeerAudio
|
||||
from facefusion.types import AudioCodec, AudioFrame, Buffer, FrameHandler, RtcPeer, RtcPeerAudio
|
||||
from .assert_helper import get_test_example_file, get_test_examples_directory
|
||||
|
||||
|
||||
@@ -37,8 +37,11 @@ def before_each() -> None:
|
||||
rtc_store.clear()
|
||||
|
||||
|
||||
def set_ready_event(ready_event : threading.Event, track : int, close_callback : FrameHandler) -> None:
|
||||
def dispatch_frame(buffer : Buffer, track : int, frame_handler : FrameHandler) -> threading.Event:
|
||||
frame_handler(buffer, 0)
|
||||
ready_event = threading.Event()
|
||||
ready_event.set()
|
||||
return ready_event
|
||||
|
||||
|
||||
def test_run_audio_encode_loop() -> None:
|
||||
@@ -48,6 +51,12 @@ def test_run_audio_encode_loop() -> None:
|
||||
rtc_peer : RtcPeer =\
|
||||
{
|
||||
'peer_connection': peer_connection,
|
||||
'audio':
|
||||
{
|
||||
'sender_track': 0,
|
||||
'receiver_track': 0,
|
||||
'codec': 'opus'
|
||||
},
|
||||
'video':
|
||||
{
|
||||
'sender_track': 0,
|
||||
@@ -82,11 +91,7 @@ def test_receive_audio_frames(audio_codec : AudioCodec) -> None:
|
||||
audio_frame = numpy.frombuffer(audio_buffer, dtype = numpy.int16).astype(numpy.float32) / 32768.0
|
||||
audio_queue : Queue[Tuple[float, AudioFrame]] = Queue(maxsize = 300)
|
||||
|
||||
datachannel_mock = MagicMock()
|
||||
ready_event = threading.Event()
|
||||
datachannel_mock.rtcSetClosedCallback.side_effect = partial(set_ready_event, ready_event)
|
||||
|
||||
with patch('facefusion.libraries.datachannel.create_static_library', return_value = datachannel_mock):
|
||||
with patch('facefusion.apis.stream_audio.create_receive_event', side_effect = partial(dispatch_frame, bytes([ 0 ]))):
|
||||
with patch('facefusion.apis.stream_audio.decode_audio_frame', return_value = audio_frame.tobytes()):
|
||||
rtc_peer_audio : RtcPeerAudio =\
|
||||
{
|
||||
@@ -94,12 +99,7 @@ def test_receive_audio_frames(audio_codec : AudioCodec) -> None:
|
||||
'receiver_track': 0,
|
||||
'codec': audio_codec
|
||||
}
|
||||
audio_receiver_thread = threading.Thread(target = receive_audio_frames, args = (rtc_peer_audio, audio_queue), daemon = True)
|
||||
audio_receiver_thread.start()
|
||||
ready_event.wait(timeout = 5.0)
|
||||
datachannel_mock.rtcSetFrameCallback.call_args[0][1](0, bytes([ 0 ]), 1, None, None)
|
||||
datachannel_mock.rtcSetClosedCallback.call_args[0][1](0, None)
|
||||
audio_receiver_thread.join(timeout = 5.0)
|
||||
receive_audio_frames(rtc_peer_audio, audio_queue)
|
||||
|
||||
_, temp_audio_frame = audio_queue.get_nowait()
|
||||
|
||||
@@ -113,7 +113,7 @@ def test_handle_audio_frame() -> None:
|
||||
audio_queue : Queue[Tuple[float, AudioFrame]] = Queue(maxsize = 300)
|
||||
|
||||
with patch('facefusion.apis.stream_audio.decode_audio_frame', return_value = audio_frame.tobytes()):
|
||||
handle_audio_frame('opus', audio_decoder_mock, audio_queue, 0, ctypes.c_void_p(), 1, ctypes.c_void_p(), ctypes.c_void_p())
|
||||
handle_audio_frame('opus', audio_decoder_mock, audio_queue, bytes(), 0)
|
||||
|
||||
_, temp_audio_frame = audio_queue.get_nowait()
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@ from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from functools import partial
|
||||
from queue import Queue
|
||||
from typing import Tuple
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
import cv2
|
||||
import numpy
|
||||
@@ -18,7 +18,7 @@ from facefusion.common_helper import is_linux, is_macos, is_windows
|
||||
from facefusion.download import conditional_download
|
||||
from facefusion.hash_helper import create_hash
|
||||
from facefusion.libraries import aom as aom_module, datachannel as datachannel_module, vpx as vpx_module
|
||||
from facefusion.types import BufferPack, FrameHandler, RtcPeer, RtcPeerVideo, Time, VideoCodec
|
||||
from facefusion.types import Buffer, BufferPack, FrameHandler, RtcPeer, RtcPeerVideo, Time, VideoCodec
|
||||
from facefusion.vision import read_video_frame
|
||||
from .assert_helper import get_test_example_file, get_test_examples_directory
|
||||
|
||||
@@ -44,8 +44,11 @@ def before_each() -> None:
|
||||
rtc_store.clear()
|
||||
|
||||
|
||||
def set_ready_event(ready_event : threading.Event, track : int, close_callback : FrameHandler) -> None:
|
||||
def dispatch_frame(buffer : Buffer, track : int, frame_handler : FrameHandler) -> threading.Event:
|
||||
frame_handler(buffer, 0)
|
||||
ready_event = threading.Event()
|
||||
ready_event.set()
|
||||
return ready_event
|
||||
|
||||
|
||||
@pytest.mark.parametrize('video_codec, payload_type', [ ('av1', 35), ('vp8', 96), ('vp9', 98) ])
|
||||
@@ -57,6 +60,12 @@ def test_run_video_encode_loop(video_codec : VideoCodec, payload_type : int) ->
|
||||
rtc_peer : RtcPeer =\
|
||||
{
|
||||
'peer_connection': peer_connection,
|
||||
'audio':
|
||||
{
|
||||
'sender_track': 0,
|
||||
'receiver_track': 0,
|
||||
'codec': 'opus'
|
||||
},
|
||||
'video':
|
||||
{
|
||||
'sender_track': video_sender_track,
|
||||
@@ -101,12 +110,8 @@ def test_receive_video_frames(video_codec : VideoCodec) -> None:
|
||||
video_frame = read_video_frame(get_test_example_file('target-240p.mp4'))
|
||||
video_queue : Queue[Tuple[Time, Future[BufferPack]]] = Queue(maxsize = 30)
|
||||
|
||||
datachannel_mock = MagicMock()
|
||||
ready_event = threading.Event()
|
||||
datachannel_mock.rtcSetClosedCallback.side_effect = partial(set_ready_event, ready_event)
|
||||
|
||||
with ThreadPoolExecutor(max_workers = 1) as executor:
|
||||
with patch('facefusion.libraries.datachannel.create_static_library', return_value = datachannel_mock):
|
||||
with patch('facefusion.apis.stream_video.create_receive_event', side_effect = partial(dispatch_frame, bytes([ 0 ]))):
|
||||
with patch('facefusion.apis.stream_video.decode_video_frame', return_value = video_frame):
|
||||
with patch('facefusion.apis.stream_video.process_video_frame', return_value = BufferPack(buffer = video_frame.tobytes(), resolution = (426, 226))):
|
||||
rtc_peer_video : RtcPeerVideo =\
|
||||
@@ -115,12 +120,7 @@ def test_receive_video_frames(video_codec : VideoCodec) -> None:
|
||||
'receiver_track': 0,
|
||||
'codec': video_codec
|
||||
}
|
||||
video_receiver_thread = threading.Thread(target = receive_video_frames, args = (rtc_peer_video, video_queue, executor), daemon = True)
|
||||
video_receiver_thread.start()
|
||||
ready_event.wait(timeout = 5.0)
|
||||
datachannel_mock.rtcSetFrameCallback.call_args[0][1](0, bytes([ 0 ]), 1, None, None)
|
||||
datachannel_mock.rtcSetClosedCallback.call_args[0][1](0, None)
|
||||
video_receiver_thread.join(timeout = 5.0)
|
||||
receive_video_frames(rtc_peer_video, video_queue, executor)
|
||||
_, video_future = video_queue.get_nowait()
|
||||
|
||||
video_buffer = video_future.result().get('buffer')
|
||||
@@ -267,7 +267,7 @@ def test_handle_video_frame(video_codec : VideoCodec) -> None:
|
||||
with ThreadPoolExecutor(max_workers = 1) as executor:
|
||||
with patch('facefusion.apis.stream_video.decode_video_frame', return_value = video_frame):
|
||||
with patch('facefusion.apis.stream_video.process_video_frame', return_value = BufferPack(buffer = video_frame.tobytes(), resolution = (426, 226))):
|
||||
handle_video_frame(video_codec, video_decoder, video_queue, executor, 0, ctypes.c_void_p(), 1, ctypes.c_void_p(), ctypes.c_void_p())
|
||||
handle_video_frame(video_codec, video_decoder, video_queue, executor, bytes(), 0)
|
||||
_, video_future = video_queue.get_nowait()
|
||||
|
||||
video_buffer = video_future.result().get('buffer')
|
||||
|
||||
Reference in New Issue
Block a user