diff --git a/facefusion/apis/stream_manager.py b/facefusion/apis/stream_manager.py index a4dd32d2..62268ee7 100644 --- a/facefusion/apis/stream_manager.py +++ b/facefusion/apis/stream_manager.py @@ -12,16 +12,17 @@ from starlette.websockets import WebSocket from facefusion import rtc, rtc_store, state_manager, streamer from facefusion.apis.stream_audio import receive_audio_frames, run_audio_encode_loop from facefusion.apis.stream_video import receive_video_frames, run_video_encode_loop -from facefusion.audio import create_empty_audio_frame from facefusion.libraries import datachannel as datachannel_module from facefusion.types import AudioCodec, AudioFrame, BufferPack, PeerConnection, RtcPeer, RtcPeerAudio, SdpAnswer, SdpOffer, SessionId, Time, VideoCodec, VisionFrame +from facefusion.vision import read_static_images async def process_image(websocket : WebSocket) -> None: capture_vision_frame = await anext(receive_vision_frames(websocket), None) if numpy.any(capture_vision_frame): - output_vision_frame = streamer.process_frame(create_empty_audio_frame(), capture_vision_frame) + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) + output_vision_frame = streamer.process_stream_frame(source_vision_frames, capture_vision_frame) is_success, output_frame_buffer = cv2.imencode('.jpg', output_vision_frame) if is_success: diff --git a/facefusion/apis/stream_video.py b/facefusion/apis/stream_video.py index 947a3cf5..810af061 100644 --- a/facefusion/apis/stream_video.py +++ b/facefusion/apis/stream_video.py @@ -6,11 +6,11 @@ from typing import Optional, Tuple import cv2 import numpy -from facefusion import rtc, streamer +from facefusion import rtc, state_manager, streamer from facefusion.apis.stream_event import create_receive_event -from facefusion.audio import create_empty_audio_frame from facefusion.codecs import aom_decoder, aom_encoder, vpx_decoder, vpx_encoder from facefusion.types import AomDecoder, AomEncoder, BitRate, Buffer, BufferPack, Resolution, RtcPeer, RtcPeerVideo, Time, VideoCodec, VisionFrame, VpxDecoder, VpxEncoder +from facefusion.vision import read_static_images def run_video_encode_loop(rtc_peer : RtcPeer, video_queue : Queue[Tuple[Time, Future[BufferPack]]]) -> None: @@ -72,7 +72,8 @@ def receive_video_frames(rtc_peer_video : RtcPeerVideo, video_queue : Queue[Tupl def process_video_frame(input_vision_frame : VisionFrame) -> BufferPack: - output_vision_frame = streamer.process_frame(create_empty_audio_frame(), input_vision_frame) + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) + output_vision_frame = streamer.process_stream_frame(source_vision_frames, input_vision_frame) output_resolution : Resolution = (output_vision_frame.shape[1], output_vision_frame.shape[0]) output_buffer = cv2.cvtColor(output_vision_frame, cv2.COLOR_BGR2YUV_I420).tobytes() return BufferPack(buffer = output_buffer, resolution = output_resolution) diff --git a/facefusion/streamer.py b/facefusion/streamer.py index 6cfccdd5..413cf4dc 100644 --- a/facefusion/streamer.py +++ b/facefusion/streamer.py @@ -2,7 +2,7 @@ import os import subprocess from collections import deque from concurrent.futures import ThreadPoolExecutor -from typing import Deque, Iterator +from typing import Deque, Iterator, List import cv2 import numpy @@ -14,30 +14,30 @@ from facefusion.content_analyser import analyse_stream from facefusion.ffmpeg import open_ffmpeg from facefusion.filesystem import is_directory from facefusion.processors.core import get_processors_modules -from facefusion.types import AudioFrame, Buffer, Fps, StreamMode, VisionFrame +from facefusion.types import Fps, StreamMode, VisionFrame from facefusion.vision import extract_vision_mask, read_static_images def multi_process_capture(camera_capture : cv2.VideoCapture, camera_fps : Fps) -> Iterator[VisionFrame]: capture_deque : Deque[VisionFrame] = deque() + source_vision_frames = read_static_images(state_manager.get_item('source_paths')) with tqdm(desc = translator.get('streaming'), unit = 'frame', disable = state_manager.get_item('log_level') in [ 'warn', 'error' ]) as progress: with ThreadPoolExecutor(max_workers = state_manager.get_item('execution_thread_count')) as executor: futures = [] while camera_capture and camera_capture.isOpened(): - _, capture_frame = camera_capture.read() - if analyse_stream(capture_frame, camera_fps): + _, capture_vision_frame = camera_capture.read() + if analyse_stream(capture_vision_frame, camera_fps): camera_capture.release() - if numpy.any(capture_frame): - audio_frame = create_empty_audio_frame() - future = executor.submit(process_frame, audio_frame, capture_frame) + if numpy.any(capture_vision_frame): + future = executor.submit(process_stream_frame, source_vision_frames, capture_vision_frame) futures.append(future) for future_done in [ future for future in futures if future.done() ]: - capture_frame = future_done.result() - capture_deque.append(capture_frame) + capture_vision_frame = future_done.result() + capture_deque.append(capture_vision_frame) futures.remove(future_done) while capture_deque: @@ -45,10 +45,10 @@ def multi_process_capture(camera_capture : cv2.VideoCapture, camera_fps : Fps) - yield capture_deque.popleft() -def process_frame(stream_audio_frame : AudioFrame, stream_vision_frame : VisionFrame) -> VisionFrame: - source_vision_frames = read_static_images(state_manager.get_item('source_paths')) +def process_stream_frame(source_vision_frames : List[VisionFrame], target_vision_frame : VisionFrame) -> VisionFrame: + source_audio_frame = create_empty_audio_frame() source_voice_frame = create_empty_audio_frame() - temp_vision_frame = stream_vision_frame.copy() + temp_vision_frame = target_vision_frame.copy() temp_vision_mask = extract_vision_mask(temp_vision_frame) for processor_module in get_processors_modules(state_manager.get_item('processors')): @@ -58,9 +58,9 @@ def process_frame(stream_audio_frame : AudioFrame, stream_vision_frame : VisionF temp_vision_frame, temp_vision_mask = processor_module.process_frame( { 'source_vision_frames': source_vision_frames, - 'source_audio_frame': stream_audio_frame, + 'source_audio_frame': source_audio_frame, 'source_voice_frame': source_voice_frame, - 'target_vision_frames': [ stream_vision_frame ], + 'target_vision_frames': [ target_vision_frame ], 'temp_vision_frame': temp_vision_frame, 'temp_vision_mask': temp_vision_mask })