File size: 9,526 Bytes
a1da63c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
from typing import List, Optional, Tuple

import gradio

import facefusion.choices
from facefusion import state_manager, wording
from facefusion.common_helper import calc_float_step
from facefusion.face_analyser import get_many_faces
from facefusion.face_selector import sort_and_filter_faces
from facefusion.face_store import clear_reference_faces, clear_static_faces
from facefusion.filesystem import is_image, is_video
from facefusion.typing import FaceSelectorAge, FaceSelectorGender, FaceSelectorMode, FaceSelectorOrder, VisionFrame
from facefusion.uis.core import get_ui_component, get_ui_components, register_ui_component
from facefusion.uis.typing import ComponentOptions
from facefusion.uis.ui_helper import convert_str_none
from facefusion.vision import get_video_frame, normalize_frame_color, read_static_image

FACE_SELECTOR_MODE_DROPDOWN : Optional[gradio.Dropdown] = None
FACE_SELECTOR_ORDER_DROPDOWN : Optional[gradio.Dropdown] = None
FACE_SELECTOR_AGE_DROPDOWN : Optional[gradio.Dropdown] = None
FACE_SELECTOR_GENDER_DROPDOWN : Optional[gradio.Dropdown] = None
REFERENCE_FACE_POSITION_GALLERY : Optional[gradio.Gallery] = None
REFERENCE_FACE_DISTANCE_SLIDER : Optional[gradio.Slider] = None


def render() -> None:
	global FACE_SELECTOR_MODE_DROPDOWN
	global FACE_SELECTOR_ORDER_DROPDOWN
	global FACE_SELECTOR_AGE_DROPDOWN
	global FACE_SELECTOR_GENDER_DROPDOWN
	global REFERENCE_FACE_POSITION_GALLERY
	global REFERENCE_FACE_DISTANCE_SLIDER

	reference_face_gallery_options : ComponentOptions =\
	{
		'label': wording.get('uis.reference_face_gallery'),
		'object_fit': 'cover',
		'columns': 8,
		'allow_preview': False,
		'visible': 'reference' in state_manager.get_item('face_selector_mode')
	}
	if is_image(state_manager.get_item('target_path')):
		reference_frame = read_static_image(state_manager.get_item('target_path'))
		reference_face_gallery_options['value'] = extract_gallery_frames(reference_frame)
	if is_video(state_manager.get_item('target_path')):
		reference_frame = get_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number'))
		reference_face_gallery_options['value'] = extract_gallery_frames(reference_frame)
	FACE_SELECTOR_MODE_DROPDOWN = gradio.Dropdown(
		label = wording.get('uis.face_selector_mode_dropdown'),
		choices = facefusion.choices.face_selector_modes,
		value = state_manager.get_item('face_selector_mode')
	)
	REFERENCE_FACE_POSITION_GALLERY = gradio.Gallery(**reference_face_gallery_options)
	with gradio.Row():
		FACE_SELECTOR_ORDER_DROPDOWN = gradio.Dropdown(
			label = wording.get('uis.face_selector_order_dropdown'),
			choices = facefusion.choices.face_selector_orders,
			value = state_manager.get_item('face_selector_order')
		)
		FACE_SELECTOR_AGE_DROPDOWN = gradio.Dropdown(
			label = wording.get('uis.face_selector_age_dropdown'),
			choices = [ 'none' ] + facefusion.choices.face_selector_ages,
			value = state_manager.get_item('face_selector_age') or 'none'
		)
		FACE_SELECTOR_GENDER_DROPDOWN = gradio.Dropdown(
			label = wording.get('uis.face_selector_gender_dropdown'),
			choices = [ 'none' ] + facefusion.choices.face_selector_genders,
			value = state_manager.get_item('face_selector_gender') or 'none'
		)
	REFERENCE_FACE_DISTANCE_SLIDER = gradio.Slider(
		label = wording.get('uis.reference_face_distance_slider'),
		value = state_manager.get_item('reference_face_distance'),
		step = calc_float_step(facefusion.choices.reference_face_distance_range),
		minimum = facefusion.choices.reference_face_distance_range[0],
		maximum = facefusion.choices.reference_face_distance_range[-1],
		visible = 'reference' in state_manager.get_item('face_selector_mode')
	)
	register_ui_component('face_selector_mode_dropdown', FACE_SELECTOR_MODE_DROPDOWN)
	register_ui_component('face_selector_order_dropdown', FACE_SELECTOR_ORDER_DROPDOWN)
	register_ui_component('face_selector_age_dropdown', FACE_SELECTOR_AGE_DROPDOWN)
	register_ui_component('face_selector_gender_dropdown', FACE_SELECTOR_GENDER_DROPDOWN)
	register_ui_component('reference_face_position_gallery', REFERENCE_FACE_POSITION_GALLERY)
	register_ui_component('reference_face_distance_slider', REFERENCE_FACE_DISTANCE_SLIDER)


def listen() -> None:
	FACE_SELECTOR_MODE_DROPDOWN.change(update_face_selector_mode, inputs = FACE_SELECTOR_MODE_DROPDOWN, outputs = [ REFERENCE_FACE_POSITION_GALLERY, REFERENCE_FACE_DISTANCE_SLIDER ])
	FACE_SELECTOR_ORDER_DROPDOWN.change(update_face_selector_order, inputs = FACE_SELECTOR_ORDER_DROPDOWN, outputs = REFERENCE_FACE_POSITION_GALLERY)
	FACE_SELECTOR_AGE_DROPDOWN.change(update_face_selector_age, inputs = FACE_SELECTOR_AGE_DROPDOWN, outputs = REFERENCE_FACE_POSITION_GALLERY)
	FACE_SELECTOR_GENDER_DROPDOWN.change(update_face_selector_gender, inputs = FACE_SELECTOR_GENDER_DROPDOWN, outputs = REFERENCE_FACE_POSITION_GALLERY)
	REFERENCE_FACE_POSITION_GALLERY.select(clear_and_update_reference_face_position)
	REFERENCE_FACE_DISTANCE_SLIDER.release(update_reference_face_distance, inputs = REFERENCE_FACE_DISTANCE_SLIDER)

	for ui_component in get_ui_components(
	[
		'target_image',
		'target_video'
	]):
		for method in [ 'upload', 'change', 'clear' ]:
			getattr(ui_component, method)(update_reference_face_position)
			getattr(ui_component, method)(update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY)

	for ui_component in get_ui_components(
	[
		'face_detector_model_dropdown',
		'face_detector_size_dropdown',
		'face_detector_angles_checkbox_group'
	]):
		ui_component.change(clear_and_update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY)

	face_detector_score_slider = get_ui_component('face_detector_score_slider')
	if face_detector_score_slider:
		face_detector_score_slider.release(clear_and_update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY)

	preview_frame_slider = get_ui_component('preview_frame_slider')
	if preview_frame_slider:
		preview_frame_slider.release(update_reference_frame_number, inputs = preview_frame_slider)
		preview_frame_slider.release(update_reference_position_gallery, outputs = REFERENCE_FACE_POSITION_GALLERY)


def update_face_selector_mode(face_selector_mode : FaceSelectorMode) -> Tuple[gradio.Gallery, gradio.Slider]:
	state_manager.set_item('face_selector_mode', face_selector_mode)
	if face_selector_mode == 'many':
		return gradio.Gallery(visible = False), gradio.Slider(visible = False)
	if face_selector_mode == 'one':
		return gradio.Gallery(visible = False), gradio.Slider(visible = False)
	if face_selector_mode == 'reference':
		return gradio.Gallery(visible = True), gradio.Slider(visible = True)


def update_face_selector_order(face_analyser_order : FaceSelectorOrder) -> gradio.Gallery:
	state_manager.set_item('face_selector_order', convert_str_none(face_analyser_order))
	return update_reference_position_gallery()


def update_face_selector_age(face_selector_age : FaceSelectorAge) -> gradio.Gallery:
	state_manager.set_item('face_selector_age', convert_str_none(face_selector_age))
	return update_reference_position_gallery()


def update_face_selector_gender(face_analyser_gender : FaceSelectorGender) -> gradio.Gallery:
	state_manager.set_item('face_selector_gender', convert_str_none(face_analyser_gender))
	return update_reference_position_gallery()


def clear_and_update_reference_face_position(event : gradio.SelectData) -> gradio.Gallery:
	clear_reference_faces()
	clear_static_faces()
	update_reference_face_position(event.index)
	return update_reference_position_gallery()


def update_reference_face_position(reference_face_position : int = 0) -> None:
	state_manager.set_item('reference_face_position', reference_face_position)


def update_reference_face_distance(reference_face_distance : float) -> None:
	state_manager.set_item('reference_face_distance', reference_face_distance)


def update_reference_frame_number(reference_frame_number : int) -> None:
	state_manager.set_item('reference_frame_number', reference_frame_number)


def clear_and_update_reference_position_gallery() -> gradio.Gallery:
	clear_reference_faces()
	clear_static_faces()
	return update_reference_position_gallery()


def update_reference_position_gallery() -> gradio.Gallery:
	gallery_vision_frames = []
	if is_image(state_manager.get_item('target_path')):
		temp_vision_frame = read_static_image(state_manager.get_item('target_path'))
		gallery_vision_frames = extract_gallery_frames(temp_vision_frame)
	if is_video(state_manager.get_item('target_path')):
		temp_vision_frame = get_video_frame(state_manager.get_item('target_path'), state_manager.get_item('reference_frame_number'))
		gallery_vision_frames = extract_gallery_frames(temp_vision_frame)
	if gallery_vision_frames:
		return gradio.Gallery(value = gallery_vision_frames)
	return gradio.Gallery(value = None)


def extract_gallery_frames(temp_vision_frame : VisionFrame) -> List[VisionFrame]:
	gallery_vision_frames = []
	faces = sort_and_filter_faces(get_many_faces([ temp_vision_frame ]))

	for face in faces:
		start_x, start_y, end_x, end_y = map(int, face.bounding_box)
		padding_x = int((end_x - start_x) * 0.25)
		padding_y = int((end_y - start_y) * 0.25)
		start_x = max(0, start_x - padding_x)
		start_y = max(0, start_y - padding_y)
		end_x = max(0, end_x + padding_x)
		end_y = max(0, end_y + padding_y)
		crop_vision_frame = temp_vision_frame[start_y:end_y, start_x:end_x]
		crop_vision_frame = normalize_frame_color(crop_vision_frame)
		gallery_vision_frames.append(crop_vision_frame)
	return gallery_vision_frames