content store (#1225)

* introduce content store for content analyser

* introduce content store for content analyser part2

* guard content store as well

* minor changes

* minor changes

* minor changes

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:

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

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)

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:

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

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:

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)

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,

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

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
 

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