tonyassi's picture
Upload 96 files
9ae3d29
from typing import List, Optional, Tuple, Any, Dict
from time import sleep
import cv2
import gradio
import DeepFakeAI.choices
import DeepFakeAI.globals
from DeepFakeAI import wording
from DeepFakeAI.capturer import get_video_frame
from DeepFakeAI.face_analyser import get_many_faces
from DeepFakeAI.face_reference import clear_face_reference
from DeepFakeAI.typing import Frame, FaceRecognition
from DeepFakeAI.uis import core as ui
from DeepFakeAI.uis.typing import ComponentName, Update
from DeepFakeAI.utilities import is_image, is_video
FACE_RECOGNITION_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_RECOGNITION_DROPDOWN
global REFERENCE_FACE_POSITION_GALLERY
global REFERENCE_FACE_DISTANCE_SLIDER
with gradio.Box():
reference_face_gallery_args: Dict[str, Any] = {
'label': wording.get('reference_face_gallery_label'),
'height': 120,
'object_fit': 'cover',
'columns': 10,
'allow_preview': False,
'visible': 'reference' in DeepFakeAI.globals.face_recognition
}
if is_image(DeepFakeAI.globals.target_path):
reference_frame = cv2.imread(DeepFakeAI.globals.target_path)
reference_face_gallery_args['value'] = extract_gallery_frames(reference_frame)
if is_video(DeepFakeAI.globals.target_path):
reference_frame = get_video_frame(DeepFakeAI.globals.target_path, DeepFakeAI.globals.reference_frame_number)
reference_face_gallery_args['value'] = extract_gallery_frames(reference_frame)
FACE_RECOGNITION_DROPDOWN = gradio.Dropdown(
label = wording.get('face_recognition_dropdown_label'),
choices = DeepFakeAI.choices.face_recognition,
value = DeepFakeAI.globals.face_recognition
)
REFERENCE_FACE_POSITION_GALLERY = gradio.Gallery(**reference_face_gallery_args)
REFERENCE_FACE_DISTANCE_SLIDER = gradio.Slider(
label = wording.get('reference_face_distance_slider_label'),
value = DeepFakeAI.globals.reference_face_distance,
maximum = 3,
step = 0.05,
visible = 'reference' in DeepFakeAI.globals.face_recognition
)
ui.register_component('face_recognition_dropdown', FACE_RECOGNITION_DROPDOWN)
ui.register_component('reference_face_position_gallery', REFERENCE_FACE_POSITION_GALLERY)
ui.register_component('reference_face_distance_slider', REFERENCE_FACE_DISTANCE_SLIDER)
def listen() -> None:
FACE_RECOGNITION_DROPDOWN.select(update_face_recognition, inputs = FACE_RECOGNITION_DROPDOWN, outputs = [ REFERENCE_FACE_POSITION_GALLERY, REFERENCE_FACE_DISTANCE_SLIDER ])
REFERENCE_FACE_POSITION_GALLERY.select(clear_and_update_face_reference_position)
REFERENCE_FACE_DISTANCE_SLIDER.change(update_reference_face_distance, inputs = REFERENCE_FACE_DISTANCE_SLIDER)
update_component_names : List[ComponentName] =\
[
'target_file',
'preview_frame_slider'
]
for component_name in update_component_names:
component = ui.get_component(component_name)
if component:
component.change(update_face_reference_position, outputs = REFERENCE_FACE_POSITION_GALLERY)
select_component_names : List[ComponentName] =\
[
'face_analyser_direction_dropdown',
'face_analyser_age_dropdown',
'face_analyser_gender_dropdown'
]
for component_name in select_component_names:
component = ui.get_component(component_name)
if component:
component.select(update_face_reference_position, outputs = REFERENCE_FACE_POSITION_GALLERY)
def update_face_recognition(face_recognition : FaceRecognition) -> Tuple[Update, Update]:
if face_recognition == 'reference':
DeepFakeAI.globals.face_recognition = face_recognition
return gradio.update(visible = True), gradio.update(visible = True)
if face_recognition == 'many':
DeepFakeAI.globals.face_recognition = face_recognition
return gradio.update(visible = False), gradio.update(visible = False)
def clear_and_update_face_reference_position(event: gradio.SelectData) -> Update:
clear_face_reference()
return update_face_reference_position(event.index)
def update_face_reference_position(reference_face_position : int = 0) -> Update:
sleep(0.2)
gallery_frames = []
DeepFakeAI.globals.reference_face_position = reference_face_position
if is_image(DeepFakeAI.globals.target_path):
reference_frame = cv2.imread(DeepFakeAI.globals.target_path)
gallery_frames = extract_gallery_frames(reference_frame)
if is_video(DeepFakeAI.globals.target_path):
reference_frame = get_video_frame(DeepFakeAI.globals.target_path, DeepFakeAI.globals.reference_frame_number)
gallery_frames = extract_gallery_frames(reference_frame)
if gallery_frames:
return gradio.update(value = gallery_frames)
return gradio.update(value = None)
def update_reference_face_distance(reference_face_distance : float) -> Update:
DeepFakeAI.globals.reference_face_distance = reference_face_distance
return gradio.update(value = reference_face_distance)
def extract_gallery_frames(reference_frame : Frame) -> List[Frame]:
crop_frames = []
faces = get_many_faces(reference_frame)
for face in faces:
start_x, start_y, end_x, end_y = map(int, face['bbox'])
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_frame = reference_frame[start_y:end_y, start_x:end_x]
crop_frames.append(ui.normalize_frame(crop_frame))
return crop_frames