SwitcherAI/uis/components/SwitcherAI_uis_components_face_selector.py

from typing import List, Optional, Tuple, Any, Dict
from time import sleep

import cv2
import gradio

import SwitcherAI.choices
import SwitcherAI.globals
from SwitcherAI import wording
from SwitcherAI.capturer import get_video_frame
from SwitcherAI.face_analyser import get_many_faces
from SwitcherAI.face_reference import clear_face_reference
from SwitcherAI.typing import Frame, FaceRecognition
from SwitcherAI.uis import core as ui
from SwitcherAI.uis.typing import ComponentName, Update
from SwitcherAI.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 SwitcherAI.globals.face_recognition
		}
		if is_image(SwitcherAI.globals.target_path):
			reference_frame = cv2.imread(SwitcherAI.globals.target_path)
			reference_face_gallery_args['value'] = extract_gallery_frames(reference_frame)
		if is_video(SwitcherAI.globals.target_path):
			reference_frame = get_video_frame(SwitcherAI.globals.target_path, SwitcherAI.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 = SwitcherAI.choices.face_recognition,
			value = SwitcherAI.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 = SwitcherAI.globals.reference_face_distance,
			maximum = 3,
			step = 0.05,
			visible = 'reference' in SwitcherAI.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':
		SwitcherAI.globals.face_recognition = face_recognition
		return gradio.update(visible = True), gradio.update(visible = True)
	if face_recognition == 'many':
		SwitcherAI.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 = []
	SwitcherAI.globals.reference_face_position = reference_face_position
	if is_image(SwitcherAI.globals.target_path):
		reference_frame = cv2.imread(SwitcherAI.globals.target_path)
		gallery_frames = extract_gallery_frames(reference_frame)
	if is_video(SwitcherAI.globals.target_path):
		reference_frame = get_video_frame(SwitcherAI.globals.target_path, SwitcherAI.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:
	SwitcherAI.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