diff --git a/facefusion/apis/core.py b/facefusion/apis/core.py index b0b8d64c..7c2037e8 100644 --- a/facefusion/apis/core.py +++ b/facefusion/apis/core.py @@ -6,6 +6,7 @@ from starlette.middleware import Middleware from starlette.middleware.cors import CORSMiddleware from starlette.routing import Route, WebSocketRoute +from facefusion import content_analyser from facefusion.apis.endpoints.assets import delete_assets, get_asset, get_assets, upload_asset from facefusion.apis.endpoints.capabilities import get_capabilities from facefusion.apis.endpoints.jobs import create_job, create_step, delete_job, delete_jobs, delete_step, get_job, get_jobs, update_job, update_jobs @@ -19,7 +20,7 @@ from facefusion.libraries import aom as aom_module, datachannel as datachannel_m def get_common_modules() -> List[ModuleType]: - return [ aom_module, datachannel_module, opus_module, vpx_module ] + return [ aom_module, content_analyser, datachannel_module, opus_module, vpx_module ] def pre_check() -> bool: diff --git a/facefusion/apis/stream_manager.py b/facefusion/apis/stream_manager.py index d83f6319..bb968846 100644 --- a/facefusion/apis/stream_manager.py +++ b/facefusion/apis/stream_manager.py @@ -9,12 +9,13 @@ import cv2 import numpy from starlette.websockets import WebSocket -from facefusion import rtc, rtc_store, state_manager, streamer +from facefusion import content_store, 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.content_analyser import analyse_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 is_vision_frame, read_static_images, to_buffer +from facefusion.vision import is_vision_frame, obscure_frame, read_static_images, to_buffer async def process_image(websocket : WebSocket) -> None: @@ -23,8 +24,11 @@ async def process_image(websocket : WebSocket) -> None: if is_vision_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) - output_vision_buffer = to_buffer(output_vision_frame) + if analyse_frame(capture_vision_frame): + output_vision_frame = obscure_frame(capture_vision_frame) + + output_vision_buffer = to_buffer(output_vision_frame) await websocket.send_bytes(output_vision_buffer) @@ -94,6 +98,7 @@ def process_video(session_id : SessionId, sdp_offer : SdpOffer) -> Optional[SdpA rtc_store.init_peers(session_id) rtc_store.get_peers(session_id).append(rtc_peer) + content_store.clear() threading.Thread(target = run_peer_loop, args = (session_id, rtc_peer), daemon = True).start() diff --git a/facefusion/apis/stream_video.py b/facefusion/apis/stream_video.py index 16ab4a30..69e16657 100644 --- a/facefusion/apis/stream_video.py +++ b/facefusion/apis/stream_video.py @@ -9,8 +9,9 @@ import numpy from facefusion import rtc, state_manager, streamer from facefusion.apis.stream_event import create_receive_event from facefusion.codecs import aom_decoder, aom_encoder, vpx_decoder, vpx_encoder +from facefusion.content_analyser import analyse_stream from facefusion.types import AomDecoder, AomEncoder, BitRate, Buffer, BufferPack, Resolution, RtcPeer, RtcPeerVideo, Time, VideoCodec, VisionFrame, VpxDecoder, VpxEncoder -from facefusion.vision import is_vision_frame, read_static_images +from facefusion.vision import is_vision_frame, obscure_frame, read_static_images def run_video_encode_loop(rtc_peer : RtcPeer, video_queue : Queue[Tuple[Time, Future[BufferPack]]]) -> None: @@ -73,7 +74,11 @@ def receive_video_frames(rtc_peer_video : RtcPeerVideo, video_queue : Queue[Tupl def process_video_frame(source_vision_frames : List[VisionFrame], input_vision_frame : VisionFrame) -> BufferPack: - output_vision_frame = streamer.process_stream_frame(source_vision_frames, input_vision_frame) + if analyse_stream(input_vision_frame): + output_vision_frame = obscure_frame(input_vision_frame) + else: + 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/content_analyser.py b/facefusion/content_analyser.py index 659150ac..a7ef7b73 100644 --- a/facefusion/content_analyser.py +++ b/facefusion/content_analyser.py @@ -3,14 +3,12 @@ from typing import Tuple import numpy -from facefusion import cli_progress, inference_manager, translator, video_manager +from facefusion import cli_progress, content_store, inference_manager, translator, video_manager from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url from facefusion.filesystem import resolve_relative_path from facefusion.thread_helper import conditional_thread_semaphore -from facefusion.types import Detection, DownloadScope, DownloadSet, Fps, InferencePool, ModelSet, VisionFrame -from facefusion.vision import detect_video_fps, fit_contain_frame, is_vision_frame, read_image - -STREAM_COUNTER = 0 +from facefusion.types import Detection, DownloadScope, DownloadSet, InferencePool, ModelSet, VisionFrame +from facefusion.vision import fit_contain_frame, is_vision_frame, read_image @lru_cache() @@ -134,17 +132,10 @@ def pre_check() -> bool: return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set) -def analyse_stream(vision_frame : VisionFrame, video_fps : Fps) -> bool: - global STREAM_COUNTER - - STREAM_COUNTER = STREAM_COUNTER + 1 - if STREAM_COUNTER % int(video_fps) == 0: - return analyse_frame(vision_frame) - return False - - def analyse_frame(vision_frame : VisionFrame) -> bool: - return detect_nsfw(vision_frame) + if is_vision_frame(vision_frame): + return detect_nsfw(vision_frame) + return False @lru_cache() @@ -155,12 +146,8 @@ def analyse_image(image_path : str) -> bool: @lru_cache() def analyse_video(video_path : str, trim_frame_start : int, trim_frame_end : int) -> bool: - video_fps = detect_video_fps(video_path) frame_range = range(trim_frame_start, trim_frame_end) video_reader = video_manager.get_reader(video_path, 'analyse_video') - rate = 0.0 - total = 0 - counter = 0 if trim_frame_start > 0: video_manager.seek_video_reader(video_reader, trim_frame_start) @@ -169,23 +156,25 @@ def analyse_video(video_path : str, trim_frame_start : int, trim_frame_end : int progress.set_title(translator.get('analysing')) progress.count(frame_range) - for frame_index in frame_range: + content_store.clear() + + for _ in frame_range: vision_frame = video_manager.read_video_frame(video_reader) - if frame_index % int(video_fps) == 0: - if is_vision_frame(vision_frame): - total += 1 + if content_store.tick() and analyse_frame(vision_frame): + content_store.set_hit() - if analyse_frame(vision_frame): - counter += 1 - - if counter > 0 and total > 0: - rate = counter / total * 100 - - progress.set_description('rate = ' + str(rate)) + progress.set_description('rate = ' + str(content_store.calculate_rate())) progress.update() - return bool(rate > 10.0) + return bool(content_store.calculate_rate() > 10.0) + + +def analyse_stream(vision_frame : VisionFrame) -> bool: + if content_store.tick() and analyse_frame(vision_frame): + content_store.set_hit() + + return content_store.get_hit() > 0 def detect_nsfw(vision_frame : VisionFrame) -> bool: diff --git a/facefusion/content_store.py b/facefusion/content_store.py new file mode 100644 index 00000000..c0b5b2f9 --- /dev/null +++ b/facefusion/content_store.py @@ -0,0 +1,32 @@ +from facefusion.types import ContentSet + +CONTENT_STORE : ContentSet =\ +{ + 'hit': 0, + 'total': 0 +} + + +def tick(step : int = 30) -> bool: + CONTENT_STORE['total'] += 1 + + return CONTENT_STORE.get('total') % step == 0 + + +def get_hit() -> int: + return CONTENT_STORE.get('hit') + + +def set_hit() -> None: + CONTENT_STORE['hit'] += 1 + + +def calculate_rate(step : int = 30) -> float: + if CONTENT_STORE.get('hit') and CONTENT_STORE.get('total'): + return CONTENT_STORE.get('hit') / CONTENT_STORE.get('total') * step * 100 + return 0.0 + + +def clear() -> None: + CONTENT_STORE['hit'] = 0 + CONTENT_STORE['total'] = 0 diff --git a/facefusion/core.py b/facefusion/core.py index 3c78163c..1ca676f9 100755 --- a/facefusion/core.py +++ b/facefusion/core.py @@ -8,7 +8,7 @@ from time import time import uvicorn import facefusion.apis.core -from facefusion import args_helper, benchmarker, cli_helper, content_analyser, hash_helper, logger, state_manager, translator +from facefusion import args_helper, benchmarker, cli_helper, content_analyser, content_store, hash_helper, logger, state_manager, translator from facefusion.args_helper import apply_args from facefusion.download import conditional_download_hashes, conditional_download_sources from facefusion.exit_helper import hard_exit, signal_exit @@ -103,7 +103,8 @@ def pre_check() -> bool: def common_pre_check() -> bool: content_analyser_content = inspect.getsource(content_analyser).encode() - return hash_helper.create_hash(content_analyser_content) == 'b51ff11f' + content_store_content = inspect.getsource(content_store).encode() + return hash_helper.create_hash(content_analyser_content) == 'c5d70a17' and hash_helper.create_hash(content_store_content) == '9cd9f029' def processors_pre_check() -> bool: diff --git a/facefusion/streamer.py b/facefusion/streamer.py index fba3921c..6d7aea4f 100644 --- a/facefusion/streamer.py +++ b/facefusion/streamer.py @@ -1,48 +1,10 @@ -import os -import subprocess -from collections import deque -from concurrent.futures import ThreadPoolExecutor -from typing import Deque, Iterator, List +from typing import List -import cv2 - -from facefusion import cli_progress, ffmpeg_builder, logger, state_manager, translator +from facefusion import logger, state_manager from facefusion.audio import create_empty_audio_frame -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 Buffer, Fps, StreamMode, VisionFrame -from facefusion.vision import extract_vision_mask, is_vision_frame, 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 cli_progress.create() as progress: - progress.set_title(translator.get('streaming')) - - with ThreadPoolExecutor(max_workers = state_manager.get_item('execution_thread_count')) as executor: - futures = [] - - while camera_capture and camera_capture.isOpened(): - _, capture_vision_frame = camera_capture.read() - if analyse_stream(capture_vision_frame, camera_fps): - camera_capture.release() - - if is_vision_frame(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_vision_frame = future_done.result() - capture_deque.append(capture_vision_frame) - futures.remove(future_done) - - while capture_deque: - progress.update() - yield capture_deque.popleft() +from facefusion.types import VisionFrame +from facefusion.vision import extract_vision_mask def process_stream_frame(source_vision_frames : List[VisionFrame], target_vision_frame : VisionFrame) -> VisionFrame: @@ -53,6 +15,7 @@ def process_stream_frame(source_vision_frames : List[VisionFrame], target_vision for processor_module in get_processors_modules(state_manager.get_item('processors')): logger.disable() + if processor_module.pre_process('stream'): logger.enable() temp_vision_frame, temp_vision_mask = processor_module.process_frame( @@ -67,34 +30,3 @@ def process_stream_frame(source_vision_frames : List[VisionFrame], target_vision logger.enable() return temp_vision_frame - - -def open_stream(stream_mode : StreamMode, stream_resolution : str, stream_fps : Fps) -> subprocess.Popen[Buffer]: - commands = ffmpeg_builder.chain( - ffmpeg_builder.capture_video(), - ffmpeg_builder.set_media_resolution(stream_resolution), - ffmpeg_builder.set_input_fps(stream_fps) - ) - - if stream_mode == 'udp': - commands.extend(ffmpeg_builder.set_input('-')) - commands.extend(ffmpeg_builder.set_stream_mode('udp')) - commands.extend(ffmpeg_builder.set_stream_quality(2000)) - commands.extend(ffmpeg_builder.set_output('udp://localhost:27000?pkt_size=1316')) - - if stream_mode == 'v4l2': - device_directory_path = '/sys/devices/virtual/video4linux' - commands.extend(ffmpeg_builder.set_input('-')) - commands.extend(ffmpeg_builder.set_stream_mode('v4l2')) - - if is_directory(device_directory_path): - device_names = os.listdir(device_directory_path) - - for device_name in device_names: - device_path = '/dev/' + device_name - commands.extend(ffmpeg_builder.set_output(device_path)) - - else: - logger.error(translator.get('stream_not_loaded').format(stream_mode = stream_mode), __name__) - - return open_ffmpeg(commands) diff --git a/facefusion/types.py b/facefusion/types.py index 29b24b89..af559d05 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -371,6 +371,12 @@ RtcPeer = TypedDict('RtcPeer', }) RtcStore : TypeAlias = Dict[SessionId, List[RtcPeer]] +ContentSet = TypedDict('ContentSet', +{ + 'hit' : int, + 'total' : int +}) + SdpAudioMedia = TypedDict('SdpAudioMedia', { 'codec': AudioCodec, diff --git a/tests/test_api_stream_manager.py b/tests/test_api_stream_manager.py index 6e1bd346..c54f6794 100644 --- a/tests/test_api_stream_manager.py +++ b/tests/test_api_stream_manager.py @@ -45,7 +45,8 @@ async def test_process_image() -> None: } ] - await process_image(websocket_mock) + with patch('facefusion.apis.stream_manager.analyse_frame', return_value = False): + await process_image(websocket_mock) websocket_mock.send_bytes.assert_called_once() diff --git a/tests/test_api_stream_video.py b/tests/test_api_stream_video.py index 58169fe9..c10f73b3 100644 --- a/tests/test_api_stream_video.py +++ b/tests/test_api_stream_video.py @@ -77,16 +77,17 @@ def test_run_video_encode_loop(video_codec : VideoCodec, payload_type : int) -> video_queue : Queue[Tuple[Time, Future[BufferPack]]] = Queue(maxsize = 30) - with ThreadPoolExecutor(max_workers = 1) as executor: - video_queue.put((0.1, executor.submit(process_video_frame, [], video_frame))) + with patch('facefusion.apis.stream_video.analyse_stream', return_value = False): + with ThreadPoolExecutor(max_workers = 1) as executor: + video_queue.put((0.1, executor.submit(process_video_frame, [], video_frame))) - with patch('facefusion.apis.stream_video.rtc.send_video') as send_video_mock: - encode_loop_thread = threading.Thread(target = run_video_encode_loop, args = (rtc_peer, video_queue), daemon = True) - encode_loop_thread.start() - empty_future : Future[BufferPack] = Future() - empty_future.set_result(BufferPack(buffer = bytes(), resolution = (0, 0))) - video_queue.put((0.0, empty_future)) - encode_loop_thread.join(timeout = 5.0) + with patch('facefusion.apis.stream_video.rtc.send_video') as send_video_mock: + encode_loop_thread = threading.Thread(target = run_video_encode_loop, args = (rtc_peer, video_queue), daemon = True) + encode_loop_thread.start() + empty_future : Future[BufferPack] = Future() + empty_future.set_result(BufferPack(buffer = bytes(), resolution = (0, 0))) + video_queue.put((0.0, empty_future)) + encode_loop_thread.join(timeout = 5.0) assert send_video_mock.called diff --git a/tests/test_content_store.py b/tests/test_content_store.py new file mode 100644 index 00000000..6ce0f7af --- /dev/null +++ b/tests/test_content_store.py @@ -0,0 +1,43 @@ +import pytest + +from facefusion.content_store import calculate_rate, clear, get_hit, set_hit, tick + + +@pytest.fixture(scope = 'function', autouse = True) +def before_each() -> None: + clear() + + +def test_get_hit() -> None: + assert get_hit() == 0 + + set_hit() + + assert get_hit() == 1 + + +def test_set_hit() -> None: + set_hit() + set_hit() + + assert get_hit() == 2 + + +def test_get_rate() -> None: + assert calculate_rate() == 0.0 + + for _ in range(100): + tick() + + set_hit() + + assert calculate_rate() == 30.0 + + +def test_clear() -> None: + tick() + set_hit() + clear() + + assert get_hit() == 0 + assert calculate_rate() == 0.0