mirror of
https://github.com/facefusion/facefusion.git
synced 2026-09-17 13:02:15 +02:00
Henry Ruhs
1362afa8b7
* introduce content store for content analyser * introduce content store for content analyser part2 * guard content store as well * minor changes * minor changes * minor changes
236 lines
6.8 KiB
Python
236 lines
6.8 KiB
Python
from functools import lru_cache
|
|
from typing import Tuple
|
|
|
|
import numpy
|
|
|
|
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, InferencePool, ModelSet, VisionFrame
|
|
from facefusion.vision import fit_contain_frame, is_vision_frame, read_image
|
|
|
|
|
|
@lru_cache()
|
|
def create_static_model_set(download_scope : DownloadScope) -> ModelSet:
|
|
return\
|
|
{
|
|
'nsfw_1':
|
|
{
|
|
'__metadata__':
|
|
{
|
|
'vendor': 'EraX',
|
|
'license': 'Apache-2.0',
|
|
'year': 2024
|
|
},
|
|
'hashes':
|
|
{
|
|
'content_analyser':
|
|
{
|
|
'url': resolve_download_url('models-3.3.0', 'nsfw_1.hash'),
|
|
'path': resolve_relative_path('../.assets/models/nsfw_1.hash')
|
|
}
|
|
},
|
|
'sources':
|
|
{
|
|
'content_analyser':
|
|
{
|
|
'url': resolve_download_url('models-3.3.0', 'nsfw_1.onnx'),
|
|
'path': resolve_relative_path('../.assets/models/nsfw_1.onnx')
|
|
}
|
|
},
|
|
'size': (640, 640),
|
|
'mean': (0.0, 0.0, 0.0),
|
|
'standard_deviation': (1.0, 1.0, 1.0)
|
|
},
|
|
'nsfw_2':
|
|
{
|
|
'__metadata__':
|
|
{
|
|
'vendor': 'Marqo',
|
|
'license': 'Apache-2.0',
|
|
'year': 2024
|
|
},
|
|
'hashes':
|
|
{
|
|
'content_analyser':
|
|
{
|
|
'url': resolve_download_url('models-3.3.0', 'nsfw_2.hash'),
|
|
'path': resolve_relative_path('../.assets/models/nsfw_2.hash')
|
|
}
|
|
},
|
|
'sources':
|
|
{
|
|
'content_analyser':
|
|
{
|
|
'url': resolve_download_url('models-3.3.0', 'nsfw_2.onnx'),
|
|
'path': resolve_relative_path('../.assets/models/nsfw_2.onnx')
|
|
}
|
|
},
|
|
'size': (384, 384),
|
|
'mean': (0.5, 0.5, 0.5),
|
|
'standard_deviation': (0.5, 0.5, 0.5)
|
|
},
|
|
'nsfw_3':
|
|
{
|
|
'__metadata__':
|
|
{
|
|
'vendor': 'Freepik',
|
|
'license': 'MIT',
|
|
'year': 2025
|
|
},
|
|
'hashes':
|
|
{
|
|
'content_analyser':
|
|
{
|
|
'url': resolve_download_url('models-3.3.0', 'nsfw_3.hash'),
|
|
'path': resolve_relative_path('../.assets/models/nsfw_3.hash')
|
|
}
|
|
},
|
|
'sources':
|
|
{
|
|
'content_analyser':
|
|
{
|
|
'url': resolve_download_url('models-3.3.0', 'nsfw_3.onnx'),
|
|
'path': resolve_relative_path('../.assets/models/nsfw_3.onnx')
|
|
}
|
|
},
|
|
'size': (448, 448),
|
|
'mean': (0.48145466, 0.4578275, 0.40821073),
|
|
'standard_deviation': (0.26862954, 0.26130258, 0.27577711)
|
|
}
|
|
}
|
|
|
|
|
|
def get_inference_pool() -> InferencePool:
|
|
model_names = [ 'nsfw_1', 'nsfw_2', 'nsfw_3' ]
|
|
_, model_source_set = collect_model_downloads()
|
|
|
|
return inference_manager.get_inference_pool(__name__, model_names, model_source_set)
|
|
|
|
|
|
def clear_inference_pool() -> None:
|
|
model_names = [ 'nsfw_1', 'nsfw_2', 'nsfw_3' ]
|
|
inference_manager.clear_inference_pool(__name__, model_names)
|
|
|
|
|
|
def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]:
|
|
model_set = create_static_model_set('full')
|
|
model_hash_set = {}
|
|
model_source_set = {}
|
|
|
|
for content_analyser_model in [ 'nsfw_1', 'nsfw_2', 'nsfw_3' ]:
|
|
model_hash_set[content_analyser_model] = model_set.get(content_analyser_model).get('hashes').get('content_analyser')
|
|
model_source_set[content_analyser_model] = model_set.get(content_analyser_model).get('sources').get('content_analyser')
|
|
|
|
return model_hash_set, model_source_set
|
|
|
|
|
|
def pre_check() -> bool:
|
|
model_hash_set, model_source_set = collect_model_downloads()
|
|
|
|
return conditional_download_hashes(model_hash_set) and conditional_download_sources(model_source_set)
|
|
|
|
|
|
def analyse_frame(vision_frame : VisionFrame) -> bool:
|
|
if is_vision_frame(vision_frame):
|
|
return detect_nsfw(vision_frame)
|
|
return False
|
|
|
|
|
|
@lru_cache()
|
|
def analyse_image(image_path : str) -> bool:
|
|
vision_frame = read_image(image_path)
|
|
return analyse_frame(vision_frame)
|
|
|
|
|
|
@lru_cache()
|
|
def analyse_video(video_path : str, trim_frame_start : int, trim_frame_end : int) -> bool:
|
|
frame_range = range(trim_frame_start, trim_frame_end)
|
|
video_reader = video_manager.get_reader(video_path, 'analyse_video')
|
|
|
|
if trim_frame_start > 0:
|
|
video_manager.seek_video_reader(video_reader, trim_frame_start)
|
|
|
|
with cli_progress.create(unit = 'frame') as progress:
|
|
progress.set_title(translator.get('analysing'))
|
|
progress.count(frame_range)
|
|
|
|
content_store.clear()
|
|
|
|
for _ in frame_range:
|
|
vision_frame = video_manager.read_video_frame(video_reader)
|
|
|
|
if content_store.tick() and analyse_frame(vision_frame):
|
|
content_store.set_hit()
|
|
|
|
progress.set_description('rate = ' + str(content_store.calculate_rate()))
|
|
progress.update()
|
|
|
|
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:
|
|
is_nsfw_1 = detect_with_nsfw_1(vision_frame)
|
|
is_nsfw_2 = detect_with_nsfw_2(vision_frame)
|
|
is_nsfw_3 = detect_with_nsfw_3(vision_frame)
|
|
|
|
return is_nsfw_1 and is_nsfw_2 or is_nsfw_1 and is_nsfw_3 or is_nsfw_2 and is_nsfw_3
|
|
|
|
|
|
def detect_with_nsfw_1(vision_frame : VisionFrame) -> bool:
|
|
detect_vision_frame = prepare_detect_frame(vision_frame, 'nsfw_1')
|
|
detection = forward_nsfw(detect_vision_frame, 'nsfw_1')
|
|
detection_score = numpy.max(numpy.amax(detection[:, 4:], axis = 1))
|
|
return bool(detection_score > 0.2)
|
|
|
|
|
|
def detect_with_nsfw_2(vision_frame : VisionFrame) -> bool:
|
|
detect_vision_frame = prepare_detect_frame(vision_frame, 'nsfw_2')
|
|
detection = forward_nsfw(detect_vision_frame, 'nsfw_2')
|
|
detection_score = detection[0] - detection[1]
|
|
return bool(detection_score > 0.25)
|
|
|
|
|
|
def detect_with_nsfw_3(vision_frame : VisionFrame) -> bool:
|
|
detect_vision_frame = prepare_detect_frame(vision_frame, 'nsfw_3')
|
|
detection = forward_nsfw(detect_vision_frame, 'nsfw_3')
|
|
detection_score = (detection[2] + detection[3]) - (detection[0] + detection[1])
|
|
return bool(detection_score > 10.5)
|
|
|
|
|
|
def forward_nsfw(vision_frame : VisionFrame, model_name : str) -> Detection:
|
|
content_analyser = get_inference_pool().get(model_name)
|
|
|
|
with conditional_thread_semaphore():
|
|
detection = content_analyser.run(None,
|
|
{
|
|
'input': vision_frame
|
|
})[0]
|
|
|
|
if model_name in [ 'nsfw_2', 'nsfw_3' ]:
|
|
return detection[0]
|
|
|
|
return detection
|
|
|
|
|
|
def prepare_detect_frame(temp_vision_frame : VisionFrame, model_name : str) -> VisionFrame:
|
|
model_set = create_static_model_set('full').get(model_name)
|
|
model_size = model_set.get('size')
|
|
model_mean = model_set.get('mean')
|
|
model_standard_deviation = model_set.get('standard_deviation')
|
|
|
|
detect_vision_frame = fit_contain_frame(temp_vision_frame, model_size)
|
|
detect_vision_frame = detect_vision_frame[:, :, ::-1] / 255.0
|
|
detect_vision_frame -= model_mean
|
|
detect_vision_frame /= model_standard_deviation
|
|
detect_vision_frame = numpy.expand_dims(detect_vision_frame.transpose(2, 0, 1), axis = 0).astype(numpy.float32)
|
|
return detect_vision_frame
|