Introduce face masker models

This commit is contained in:
henryruhs 2024-12-16 16:12:40 +01:00
parent f42117ba61
commit d339d78a3a
10 changed files with 109 additions and 23 deletions

View File

@ -32,6 +32,8 @@ reference_face_distance =
reference_frame_number =
[face_masker]
face_occluder_model =
face_parser_model =
face_mask_types =
face_mask_blur =
face_mask_padding =

View File

@ -71,6 +71,8 @@ def apply_args(args : Args, apply_state_item : ApplyStateItem) -> None:
apply_state_item('reference_face_distance', args.get('reference_face_distance'))
apply_state_item('reference_frame_number', args.get('reference_frame_number'))
# face masker
apply_state_item('face_occluder_model', args.get('face_occluder_model'))
apply_state_item('face_parser_model', args.get('face_parser_model'))
apply_state_item('face_mask_types', args.get('face_mask_types'))
apply_state_item('face_mask_blur', args.get('face_mask_blur'))
apply_state_item('face_mask_padding', normalize_padding(args.get('face_mask_padding')))

View File

@ -2,7 +2,7 @@ import logging
from typing import List, Sequence
from facefusion.common_helper import create_float_range, create_int_range
from facefusion.typing import Angle, DownloadProvider, DownloadProviderSet, DownloadScope, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskRegion, FaceMaskType, FaceSelectorMode, FaceSelectorOrder, Gender, JobStatus, LogLevel, LogLevelSet, OutputAudioEncoder, OutputVideoEncoder, OutputVideoPreset, Race, Score, TempFrameFormat, UiWorkflow, VideoMemoryStrategy
from facefusion.typing import Angle, DownloadProvider, DownloadProviderSet, DownloadScope, ExecutionProvider, ExecutionProviderSet, FaceDetectorModel, FaceDetectorSet, FaceLandmarkerModel, FaceMaskRegion, FaceMaskType, FaceOccluderModel, FaceParserModel, FaceSelectorMode, FaceSelectorOrder, Gender, JobStatus, LogLevel, LogLevelSet, OutputAudioEncoder, OutputVideoEncoder, OutputVideoPreset, Race, Score, TempFrameFormat, UiWorkflow, VideoMemoryStrategy
face_detector_set : FaceDetectorSet =\
{
@ -15,8 +15,10 @@ face_detector_models : List[FaceDetectorModel] = list(face_detector_set.keys())
face_landmarker_models : List[FaceLandmarkerModel] = [ 'many', '2dfan4', 'peppa_wutz' ]
face_selector_modes : List[FaceSelectorMode] = [ 'many', 'one', 'reference' ]
face_selector_orders : List[FaceSelectorOrder] = [ 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best' ]
face_selector_genders : List[Gender] = ['female', 'male']
face_selector_races : List[Race] = ['white', 'black', 'latino', 'asian', 'indian', 'arabic']
face_selector_genders : List[Gender] = [ 'female', 'male' ]
face_selector_races : List[Race] = [ 'white', 'black', 'latino', 'asian', 'indian', 'arabic' ]
face_occluder_models : List[FaceOccluderModel] = [ 'xseg_groggy_5' ]
face_parser_models : List[FaceParserModel] = [ 'bisenet_resnet_18', 'bisenet_resnet_34' ]
face_mask_types : List[FaceMaskType] = [ 'box', 'occlusion', 'region' ]
face_mask_regions : List[FaceMaskRegion] = [ 'skin', 'left-eyebrow', 'right-eyebrow', 'left-eye', 'right-eye', 'glasses', 'nose', 'mouth', 'upper-lip', 'lower-lip' ]
temp_frame_formats : List[TempFrameFormat] = [ 'bmp', 'jpg', 'png' ]

View File

@ -5,7 +5,7 @@ import cv2
import numpy
from cv2.typing import Size
from facefusion import inference_manager
from facefusion import inference_manager, state_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
@ -30,7 +30,7 @@ FACE_MASK_REGIONS : Dict[FaceMaskRegion, int] =\
def create_static_model_set(download_scope : DownloadScope) -> ModelSet:
return\
{
'face_occluder':
'xseg_groggy_5':
{
'hashes':
{
@ -50,7 +50,27 @@ def create_static_model_set(download_scope : DownloadScope) -> ModelSet:
},
'size': (256, 256)
},
'face_parser':
'bisenet_resnet_18':
{
'hashes':
{
'face_parser':
{
'url': resolve_download_url('models-3.1.0', 'bisenet_resnet_18.hash'),
'path': resolve_relative_path('../.assets/models/bisenet_resnet_18.hash')
}
},
'sources':
{
'face_parser':
{
'url': resolve_download_url('models-3.1.0', 'bisenet_resnet_18.onnx'),
'path': resolve_relative_path('../.assets/models/bisenet_resnet_18.onnx')
}
},
'size': (512, 512)
},
'bisenet_resnet_34':
{
'hashes':
{
@ -83,17 +103,19 @@ def clear_inference_pool() -> None:
def collect_model_downloads() -> Tuple[DownloadSet, DownloadSet]:
model_hashes = {}
model_sources = {}
model_set = create_static_model_set('full')
model_hashes =\
{
'face_occluder': model_set.get('face_occluder').get('hashes').get('face_occluder'),
'face_parser': model_set.get('face_parser').get('hashes').get('face_parser')
}
model_sources =\
{
'face_occluder': model_set.get('face_occluder').get('sources').get('face_occluder'),
'face_parser': model_set.get('face_parser').get('sources').get('face_parser')
}
if state_manager.get_item('face_occluder_model') == 'xseg_groggy_5':
model_hashes['xseg_groggy_5'] = model_set.get('xseg_groggy_5').get('hashes').get('face_occluder')
model_sources['xseg_groggy_5'] = model_set.get('xseg_groggy_5').get('sources').get('face_occluder')
if state_manager.get_item('face_parser_model') == 'bisenet_resnet_18':
model_hashes['bisenet_resnet_18'] = model_set.get('bisenet_resnet_18').get('hashes').get('face_parser')
model_sources['bisenet_resnet_18'] = model_set.get('bisenet_resnet_18').get('sources').get('face_parser')
if state_manager.get_item('face_parser_model') == 'bisenet_resnet_34':
model_hashes['bisenet_resnet_34'] = model_set.get('bisenet_resnet_34').get('hashes').get('face_parser')
model_sources['bisenet_resnet_34'] = model_set.get('bisenet_resnet_34').get('sources').get('face_parser')
return model_hashes, model_sources
@ -118,7 +140,8 @@ def create_static_box_mask(crop_size : Size, face_mask_blur : float, face_mask_p
def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask:
model_size = create_static_model_set('full').get('face_occluder').get('size')
face_occluder_model = state_manager.get_item('face_occluder_model')
model_size = create_static_model_set('full').get(face_occluder_model).get('size')
prepare_vision_frame = cv2.resize(crop_vision_frame, model_size)
prepare_vision_frame = numpy.expand_dims(prepare_vision_frame, axis = 0).astype(numpy.float32) / 255
prepare_vision_frame = prepare_vision_frame.transpose(0, 1, 2, 3)
@ -130,7 +153,8 @@ def create_occlusion_mask(crop_vision_frame : VisionFrame) -> Mask:
def create_region_mask(crop_vision_frame : VisionFrame, face_mask_regions : List[FaceMaskRegion]) -> Mask:
model_size = create_static_model_set('full').get('face_parser').get('size')
face_parser_model = state_manager.get_item('face_parser_model')
model_size = create_static_model_set('full').get(face_parser_model).get('size')
prepare_vision_frame = cv2.resize(crop_vision_frame, model_size)
prepare_vision_frame = prepare_vision_frame[:, :, ::-1].astype(numpy.float32) / 255
prepare_vision_frame = numpy.subtract(prepare_vision_frame, numpy.array([ 0.485, 0.456, 0.406 ]).astype(numpy.float32))
@ -154,7 +178,8 @@ def create_mouth_mask(face_landmark_68 : FaceLandmark68) -> Mask:
def forward_occlude_face(prepare_vision_frame : VisionFrame) -> Mask:
face_occluder = get_inference_pool().get('face_occluder')
face_occluder_model = state_manager.get_item('face_occluder_model')
face_occluder = get_inference_pool().get(face_occluder_model)
with conditional_thread_semaphore():
occlusion_mask : Mask = face_occluder.run(None,
@ -166,7 +191,8 @@ def forward_occlude_face(prepare_vision_frame : VisionFrame) -> Mask:
def forward_parse_face(prepare_vision_frame : VisionFrame) -> Mask:
face_parser = get_inference_pool().get('face_parser')
face_parser_model = state_manager.get_item('face_parser_model')
face_parser = get_inference_pool().get(face_parser_model)
with conditional_thread_semaphore():
region_mask : Mask = face_parser.run(None,

View File

@ -132,6 +132,8 @@ def create_face_selector_program() -> ArgumentParser:
def create_face_masker_program() -> ArgumentParser:
program = ArgumentParser(add_help = False)
group_face_masker = program.add_argument_group('face masker')
group_face_masker.add_argument('--face-occluder-model', help = wording.get('help.face_occluder_model'), default = config.get_str_value('face_detector.face_occluder_model', 'xseg_groggy_5'), choices = facefusion.choices.face_occluder_models)
group_face_masker.add_argument('--face-parser-model', help = wording.get('help.face_parser_model'), default = config.get_str_value('face_detector.face_parser_model', 'bisenet_resnet_34'), choices = facefusion.choices.face_parser_models)
group_face_masker.add_argument('--face-mask-types', help = wording.get('help.face_mask_types').format(choices = ', '.join(facefusion.choices.face_mask_types)), default = config.get_str_list('face_masker.face_mask_types', 'box'), choices = facefusion.choices.face_mask_types, nargs = '+', metavar = 'FACE_MASK_TYPES')
group_face_masker.add_argument('--face-mask-blur', help = wording.get('help.face_mask_blur'), type = float, default = config.get_float_value('face_masker.face_mask_blur', '0.3'), choices = facefusion.choices.face_mask_blur_range, metavar = create_float_metavar(facefusion.choices.face_mask_blur_range))
group_face_masker.add_argument('--face-mask-padding', help = wording.get('help.face_mask_padding'), type = int, default = config.get_int_list('face_masker.face_mask_padding', '0 0 0 0'), nargs = '+')

View File

@ -101,6 +101,8 @@ FaceLandmarkerModel = Literal['many', '2dfan4', 'peppa_wutz']
FaceDetectorSet = Dict[FaceDetectorModel, List[str]]
FaceSelectorMode = Literal['many', 'one', 'reference']
FaceSelectorOrder = Literal['left-right', 'right-left', 'top-bottom', 'bottom-top', 'small-large', 'large-small', 'best-worst', 'worst-best']
FaceOccluderModel = Literal['xseg_groggy_5']
FaceParserModel = Literal['bisenet_resnet_18', 'bisenet_resnet_34']
FaceMaskType = Literal['box', 'occlusion', 'region']
FaceMaskRegion = Literal['skin', 'left-eyebrow', 'right-eyebrow', 'left-eye', 'right-eye', 'glasses', 'nose', 'mouth', 'upper-lip', 'lower-lip']
TempFrameFormat = Literal['bmp', 'jpg', 'png']
@ -231,6 +233,8 @@ StateKey = Literal\
'reference_face_position',
'reference_face_distance',
'reference_frame_number',
'face_occluder_model',
'face_parser_model',
'face_mask_types',
'face_mask_blur',
'face_mask_padding',
@ -292,6 +296,8 @@ State = TypedDict('State',
'reference_face_position' : int,
'reference_face_distance' : float,
'reference_frame_number' : int,
'face_occluder_model' : FaceOccluderModel,
'face_parser_model' : FaceParserModel,
'face_mask_types' : List[FaceMaskType],
'face_mask_blur' : float,
'face_mask_padding' : Padding,

View File

@ -3,11 +3,13 @@ from typing import List, Optional, Tuple
import gradio
import facefusion.choices
from facefusion import state_manager, wording
from facefusion import face_masker, state_manager, wording
from facefusion.common_helper import calc_float_step, calc_int_step
from facefusion.typing import FaceMaskRegion, FaceMaskType
from facefusion.typing import FaceMaskRegion, FaceMaskType, FaceOccluderModel, FaceParserModel
from facefusion.uis.core import register_ui_component
FACE_OCCLUDER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None
FACE_PARSER_MODEL_DROPDOWN : Optional[gradio.Dropdown] = None
FACE_MASK_TYPES_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None
FACE_MASK_REGIONS_CHECKBOX_GROUP : Optional[gradio.CheckboxGroup] = None
FACE_MASK_BLUR_SLIDER : Optional[gradio.Slider] = None
@ -18,6 +20,8 @@ FACE_MASK_PADDING_LEFT_SLIDER : Optional[gradio.Slider] = None
def render() -> None:
global FACE_OCCLUDER_MODEL_DROPDOWN
global FACE_PARSER_MODEL_DROPDOWN
global FACE_MASK_TYPES_CHECKBOX_GROUP
global FACE_MASK_REGIONS_CHECKBOX_GROUP
global FACE_MASK_BLUR_SLIDER
@ -26,8 +30,20 @@ def render() -> None:
global FACE_MASK_PADDING_BOTTOM_SLIDER
global FACE_MASK_PADDING_LEFT_SLIDER
has_box_mask = 'box' in state_manager.get_item('face_mask_types')
has_region_mask = 'region' in state_manager.get_item('face_mask_types')
with gradio.Row():
FACE_OCCLUDER_MODEL_DROPDOWN = gradio.Dropdown(
label = wording.get('uis.face_occluder_model_dropdown'),
choices = facefusion.choices.face_occluder_models,
value = state_manager.get_item('face_occluder_model')
)
FACE_PARSER_MODEL_DROPDOWN = gradio.Dropdown(
label = wording.get('uis.face_parser_model_dropdown'),
choices = facefusion.choices.face_parser_models,
value = state_manager.get_item('face_parser_model')
)
FACE_MASK_TYPES_CHECKBOX_GROUP = gradio.CheckboxGroup(
label = wording.get('uis.face_mask_types_checkbox_group'),
choices = facefusion.choices.face_mask_types,
@ -82,6 +98,8 @@ def render() -> None:
value = state_manager.get_item('face_mask_padding')[3],
visible = has_box_mask
)
register_ui_component('face_occluder_model_dropdown', FACE_OCCLUDER_MODEL_DROPDOWN)
register_ui_component('face_parser_model_dropdown', FACE_PARSER_MODEL_DROPDOWN)
register_ui_component('face_mask_types_checkbox_group', FACE_MASK_TYPES_CHECKBOX_GROUP)
register_ui_component('face_mask_regions_checkbox_group', FACE_MASK_REGIONS_CHECKBOX_GROUP)
register_ui_component('face_mask_blur_slider', FACE_MASK_BLUR_SLIDER)
@ -92,6 +110,8 @@ def render() -> None:
def listen() -> None:
FACE_OCCLUDER_MODEL_DROPDOWN.change(update_face_occluder_model, inputs = FACE_OCCLUDER_MODEL_DROPDOWN)
FACE_PARSER_MODEL_DROPDOWN.change(update_face_parser_model, inputs = FACE_PARSER_MODEL_DROPDOWN)
FACE_MASK_TYPES_CHECKBOX_GROUP.change(update_face_mask_types, inputs = FACE_MASK_TYPES_CHECKBOX_GROUP, outputs = [ FACE_MASK_TYPES_CHECKBOX_GROUP, FACE_MASK_REGIONS_CHECKBOX_GROUP, FACE_MASK_BLUR_SLIDER, FACE_MASK_PADDING_TOP_SLIDER, FACE_MASK_PADDING_RIGHT_SLIDER, FACE_MASK_PADDING_BOTTOM_SLIDER, FACE_MASK_PADDING_LEFT_SLIDER ])
FACE_MASK_REGIONS_CHECKBOX_GROUP.change(update_face_mask_regions, inputs = FACE_MASK_REGIONS_CHECKBOX_GROUP, outputs = FACE_MASK_REGIONS_CHECKBOX_GROUP)
FACE_MASK_BLUR_SLIDER.release(update_face_mask_blur, inputs = FACE_MASK_BLUR_SLIDER)
@ -100,6 +120,24 @@ def listen() -> None:
face_mask_padding_slider.release(update_face_mask_padding, inputs = face_mask_padding_sliders)
def update_face_occluder_model(face_occluder_model : FaceOccluderModel) -> gradio.Dropdown:
face_masker.clear_inference_pool()
state_manager.set_item('face_occluder_model', face_occluder_model)
if face_masker.pre_check():
return gradio.Dropdown(value = state_manager.get_item('face_occluder_model'))
return gradio.Dropdown()
def update_face_parser_model(face_parser_model : FaceParserModel) -> gradio.Dropdown:
face_masker.clear_inference_pool()
state_manager.set_item('face_parser_model', face_parser_model)
if face_masker.pre_check():
return gradio.Dropdown(value = state_manager.get_item('face_parser_model'))
return gradio.Dropdown()
def update_face_mask_types(face_mask_types : List[FaceMaskType]) -> Tuple[gradio.CheckboxGroup, gradio.CheckboxGroup, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider, gradio.Slider]:
face_mask_types = face_mask_types or facefusion.choices.face_mask_types
state_manager.set_item('face_mask_types', face_mask_types)

View File

@ -162,7 +162,9 @@ def listen() -> None:
'face_detector_model_dropdown',
'face_detector_size_dropdown',
'face_detector_angles_checkbox_group',
'face_landmarker_model_dropdown'
'face_landmarker_model_dropdown',
'face_occluder_model_dropdown',
'face_parser_model_dropdown'
]):
ui_component.change(clear_and_update_preview_image, inputs = PREVIEW_FRAME_SLIDER, outputs = PREVIEW_IMAGE)

View File

@ -52,6 +52,8 @@ ComponentName = Literal\
'face_selector_race_dropdown',
'face_swapper_model_dropdown',
'face_swapper_pixel_boost_dropdown',
'face_occluder_model_dropdown',
'face_parser_model_dropdown',
'frame_colorizer_blend_slider',
'frame_colorizer_model_dropdown',
'frame_colorizer_size_dropdown',

View File

@ -126,6 +126,8 @@ WORDING : Dict[str, Any] =\
'reference_face_distance': 'specify the similarity between the reference face and target face',
'reference_frame_number': 'specify the frame used to create the reference face',
# face masker
'face_occluder_model': 'choose the model responsible for occluding the face',
'face_parser_model': 'choose the model responsible for parsing the face',
'face_mask_types': 'mix and match different face mask types (choices: {choices})',
'face_mask_blur': 'specify the degree of blur applied to the box mask',
'face_mask_padding': 'apply top, right, bottom and left padding to the box mask',
@ -285,6 +287,8 @@ WORDING : Dict[str, Any] =\
'face_selector_race_dropdown': 'FACE SELECTOR RACE',
'face_swapper_model_dropdown': 'FACE SWAPPER MODEL',
'face_swapper_pixel_boost_dropdown': 'FACE SWAPPER PIXEL BOOST',
'face_occluder_model_dropdown': 'FACE OCCLUDER MODEL',
'face_parser_model_dropdown': 'FACE PARSER MODEL',
'frame_colorizer_blend_slider': 'FRAME COLORIZER BLEND',
'frame_colorizer_model_dropdown': 'FRAME COLORIZER MODEL',
'frame_colorizer_size_dropdown': 'FRAME COLORIZER SIZE',