roop/ui_tk.py

import os
import customtkinter as ctk
import webbrowser
import cv2
import roop.globals
import roop.metadata

from typing import Callable, Tuple
from PIL import Image, ImageOps
from roop.face_helper import get_many_faces, get_one_face, extract_face_images
from roop.capturer import get_video_frame, get_video_frame_total
from roop.utilities import is_image, is_video, resolve_relative_path, open_with_default_app, compute_cosine_distance, has_extension

ROOT = None
ROOT_HEIGHT = 550
ROOT_WIDTH = 600

PREVIEW = None
PREVIEW_MAX_HEIGHT = 700
PREVIEW_MAX_WIDTH = 1200
IMAGE_BUTTON_WIDTH = 200
IMAGE_BUTTON_HEIGHT = 200

RECENT_DIRECTORY_SOURCE = None
RECENT_DIRECTORY_TARGET = None
RECENT_DIRECTORY_OUTPUT = None

preview_label = None
preview_slider = None
source_label = None
target_label = None
status_label = None
FACE_BUTTONS = []
INPUT_FACES_DATA = None
OUTPUT_FACES_DATA = None


def init(start: Callable, destroy: Callable) -> ctk.CTk:
    global ROOT, PREVIEW, FACE_SELECT

    ROOT = create_root(start, destroy)
    PREVIEW = create_preview(ROOT)
    FACE_SELECT = create_select_faces_win(ROOT)
    return ROOT



def create_root(start: Callable, destroy: Callable) -> ctk.CTk:
    global source_button, target_button, status_label

    ctk.deactivate_automatic_dpi_awareness()
    ctk.set_appearance_mode('system')
    ctk.set_default_color_theme(resolve_relative_path('ui.json'))
    root = ctk.CTk()
    root.minsize(ROOT_WIDTH, ROOT_HEIGHT)
    root.title(f'{roop.metadata.name} {roop.metadata.version}')
    root.configure()
    root.protocol('WM_DELETE_WINDOW', lambda: destroy())

    base_x1 = 0.075
    base_x2 = 0.575
    base_y = 0.635
    
    source_button = ctk.CTkButton(root, text='Select image with face(s)', width=IMAGE_BUTTON_WIDTH, height=IMAGE_BUTTON_HEIGHT, compound='top', anchor='center', command=lambda: select_source_path())
    source_button.place(relx=base_x1, rely=0.05)

    target_button = ctk.CTkButton(root, text='Select target image/video', width=IMAGE_BUTTON_WIDTH, height=IMAGE_BUTTON_HEIGHT, compound='top', anchor='center', command=lambda: select_target_path())
    target_button.place(relx=base_x2, rely=0.05)

    enhance_label = ctk.CTkLabel(root, text='Select face enhancement engine', anchor='w')
    enhance_label.place(relx=base_x1, rely=0.49)
    enhance_label.configure(text_color=ctk.ThemeManager.theme.get('RoopDonate').get('text_color'))
    
    enhancer_cb = ctk.CTkComboBox(root, values=["None", "Codeformer", "DMDNet (unavailable)", "GFPGAN"], width=IMAGE_BUTTON_WIDTH, command=select_enhancer)
    enhancer_cb.set("None")
    enhancer_cb.place(relx=base_x1, rely=0.532)
    
    keep_fps_value = ctk.BooleanVar(value=roop.globals.keep_fps)
    keep_fps_checkbox = ctk.CTkSwitch(root, text='Keep fps', variable=keep_fps_value, command=lambda: setattr(roop.globals, 'keep_fps', not roop.globals.keep_fps))
    keep_fps_checkbox.place(relx=base_x1, rely=base_y)

    keep_frames_value = ctk.BooleanVar(value=roop.globals.keep_frames)
    keep_frames_switch = ctk.CTkSwitch(root, text='Keep frames', variable=keep_frames_value, command=lambda: setattr(roop.globals, 'keep_frames', keep_frames_value.get()))
    keep_frames_switch.place(relx=base_x1, rely=0.68)

    skip_audio_value = ctk.BooleanVar(value=roop.globals.skip_audio)
    skip_audio_switch = ctk.CTkSwitch(root, text='Skip audio', variable=skip_audio_value, command=lambda: setattr(roop.globals, 'skip_audio', skip_audio_value.get()))
    skip_audio_switch.place(relx=base_x2, rely=base_y)

    many_faces_value = ctk.BooleanVar(value=roop.globals.many_faces)
    many_faces_switch = ctk.CTkSwitch(root, text='Many faces', variable=many_faces_value, command=lambda: setattr(roop.globals, 'many_faces', many_faces_value.get()))
    many_faces_switch.place(relx=base_x2, rely=0.68)

    use_batch_value = ctk.BooleanVar(value=roop.globals.use_batch)
    use_batch_switch = ctk.CTkSwitch(root, text='Batch process folder', variable=use_batch_value, command=lambda: setattr(roop.globals, 'use_batch', use_batch_value.get()))
    use_batch_switch.place(relx=base_x1, rely=0.725)



    base_y = 0.84
  
    start_button = ctk.CTkButton(root, text='Start', command=lambda: select_output_path(start))
    start_button.place(relx=base_x1, rely=base_y, relwidth=0.15, relheight=0.05)

    stop_button = ctk.CTkButton(root, text='Destroy', command=lambda: destroy())
    stop_button.place(relx=0.35, rely=base_y, relwidth=0.15, relheight=0.05)

    preview_button = ctk.CTkButton(root, text='Preview', command=lambda: toggle_preview())
    preview_button.place(relx=0.55, rely=base_y, relwidth=0.15, relheight=0.05)

    result_button = ctk.CTkButton(root, text='Show Result', command=lambda: show_result())
    result_button.place(relx=0.75, rely=base_y, relwidth=0.15, relheight=0.05)

    status_label = ctk.CTkLabel(root, text=None, justify='center')
    status_label.place(relx=base_x1, rely=0.9, relwidth=0.8)

    return root

def create_preview(parent) -> ctk.CTkToplevel:
    global preview_label, preview_slider

def create_preview(parent: ctk.CTkToplevel) -> ctk.CTkToplevel:
    global preview_label, preview_slider

    preview = ctk.CTkToplevel(parent)
    preview.withdraw()
    preview.title('Preview')
    preview.configure()
    preview.protocol('WM_DELETE_WINDOW', lambda: toggle_preview())
    preview.resizable(width=False, height=False)

    preview_label = ctk.CTkLabel(preview, text=None)
    preview_label.pack(fill='both', expand=True)

    preview_slider = ctk.CTkSlider(preview, from_=0, to=0, command=lambda frame_value: update_preview(frame_value))

    return preview

def update_status(text: str) -> None:
    status_label.configure(text=text)
    ROOT.update()

def update_status(text: str) -> None:
    status_label.configure(text=text)
    ROOT.update()


def select_source_path() -> None:
    global RECENT_DIRECTORY_SOURCE, INPUT_FACES_DATA

    PREVIEW.withdraw()
    source_path = ctk.filedialog.askopenfilename(title='Select source image', initialdir=RECENT_DIRECTORY_SOURCE)
    image = None
    if is_image(source_path):
        roop.globals.source_path = source_path
        RECENT_DIRECTORY_SOURCE = os.path.dirname(roop.globals.source_path)
        INPUT_FACES_DATA = extract_face_images(roop.globals.source_path,  (False, 0))
        if len(INPUT_FACES_DATA) > 0:
            if len(INPUT_FACES_DATA) == 1:
                image = render_face_from_frame(INPUT_FACES_DATA[0][1], (IMAGE_BUTTON_WIDTH, IMAGE_BUTTON_HEIGHT))
                roop.globals.SELECTED_FACE_DATA_INPUT = INPUT_FACES_DATA[0][0]
            else:
                show_face_selection(INPUT_FACES_DATA, True)
        else:
            print('No face found!')
            roop.globals.source_path = None
    else:
        roop.globals.source_path = None
    source_button.configure(image=image)
    source_button._draw()


def select_target_path() -> None:
    global RECENT_DIRECTORY_TARGET, OUTPUT_FACES_DATA

    PREVIEW.withdraw()
    target_path = ctk.filedialog.askopenfilename(title='Select target image or video', initialdir=RECENT_DIRECTORY_TARGET)
    image = None
    if is_image(target_path) and not target_path.lower().endswith(('gif')):
        roop.globals.target_path = target_path
        RECENT_DIRECTORY_TARGET = os.path.dirname(roop.globals.target_path)
        if roop.globals.many_faces:
            roop.globals.SELECTED_FACE_DATA_OUTPUT = None
            image = render_image_preview(target_path, (IMAGE_BUTTON_WIDTH, IMAGE_BUTTON_HEIGHT))
        else:
            OUTPUT_FACES_DATA = extract_face_images(roop.globals.target_path, (False, 0))
            if len(OUTPUT_FACES_DATA) > 0:
                if len(OUTPUT_FACES_DATA) == 1:
                    image = render_face_from_frame(OUTPUT_FACES_DATA[0][1], (IMAGE_BUTTON_WIDTH, IMAGE_BUTTON_HEIGHT))
                    roop.globals.SELECTED_FACE_DATA_OUTPUT = OUTPUT_FACES_DATA[0][0]
                    if roop.globals.SELECTED_FACE_DATA_INPUT is not None:
                        emb1 = roop.globals.SELECTED_FACE_DATA_INPUT.embedding
                        emb2 = roop.globals.SELECTED_FACE_DATA_OUTPUT.embedding
                        dist = compute_cosine_distance(emb1, emb2)
                        print(f'Similarity Distance between Source->Target={dist}')
                else:
                    show_face_selection(OUTPUT_FACES_DATA, False)
            else:
                print('No face found!')
                roop.globals.target_path = None

    elif is_video(target_path) or target_path.lower().endswith(('gif')):
        roop.globals.target_path = target_path
        RECENT_DIRECTORY_TARGET = os.path.dirname(roop.globals.target_path)
        if roop.globals.many_faces:
            roop.globals.SELECTED_FACE_DATA_OUTPUT = None
            image = render_video_preview(target_path, (IMAGE_BUTTON_WIDTH, IMAGE_BUTTON_HEIGHT))
        else:
            max_frame = get_video_frame_total(roop.globals.target_path)
            dialog = ctk.CTkInputDialog(text=f"Please input frame number with target face (1 - {max_frame})", title="Extract Face from Video")
            selected_frame = dialog.get_input()
            try:
                selected_frame = int(selected_frame)
            except:
                selected_frame = 1
            
            selected_frame = max(selected_frame, 1)
            selected_frame = min(selected_frame, max_frame)
            OUTPUT_FACES_DATA = extract_face_images(roop.globals.target_path, (True, selected_frame))
            if len(OUTPUT_FACES_DATA) > 0:
                if len(OUTPUT_FACES_DATA) == 1:
                    image = render_face_from_frame(OUTPUT_FACES_DATA[0][1], (IMAGE_BUTTON_WIDTH, IMAGE_BUTTON_HEIGHT))
                    roop.globals.SELECTED_FACE_DATA_OUTPUT = OUTPUT_FACES_DATA[0][0]
                else:
                    show_face_selection(OUTPUT_FACES_DATA, False)
            else:
                roop.globals.target_path = None
        
    else:
        roop.globals.target_path = None

    target_button.configure(image=image)
    target_button._draw()



def select_output_path(start):
    global RECENT_DIRECTORY_OUTPUT


def params_gen_func(proc, frame):
    return {"original_frame": frame, "blend_ratio": 0.5, "swap_mode": "all", "input_face_datas": [roop.globals.SELECTED_FACE_DATA_INPUT], "target_face_datas": [roop.globals.SELECTED_FACE_DATA_OUTPUT]}


def select_output_path(start: Callable[[], None]) -> None:
    global RECENT_DIRECTORY_OUTPUT

    if roop.globals.use_batch:
        roop.globals.target_folder_path = ctk.filedialog.askdirectory(title='Select folder to batch process')
        print(f'Batch process folder set to {roop.globals.target_folder_path}')
    else:
        roop.globals.target_folder_path = None

    if roop.globals.target_folder_path is not None:
        output_path = ctk.filedialog.askdirectory(title='Select output folder')
    elif is_image(roop.globals.target_path) and has_extension(roop.globals.target_path, ['gif']) == False:
        output_path = ctk.filedialog.asksaveasfilename(title='Save image output file', defaultextension='.png', initialfile='output.png', initialdir=RECENT_DIRECTORY_OUTPUT)
    elif is_video(roop.globals.target_path) or has_extension(roop.globals.target_path, ['gif']) == True:
        output_path = ctk.filedialog.asksaveasfilename(title='Save video output file', defaultextension='.mp4', initialfile='output.mp4', initialdir=RECENT_DIRECTORY_OUTPUT)
    else:
        output_path = None

    if output_path:
        roop.globals.output_path = output_path
        RECENT_DIRECTORY_OUTPUT = os.path.dirname(roop.globals.output_path)
        start()


def select_enhancer(choice):
    roop.globals.selected_enhancer = choice


def show_result():
    open_with_default_app(roop.globals.output_path)
    


def render_image_preview(image_path: str, size: Tuple[int, int]) -> ctk.CTkImage:
    image = Image.open(image_path)
    if size:
        image = ImageOps.fit(image, size, Image.LANCZOS)
    return ctk.CTkImage(image, size=image.size)


def render_face_from_frame(face, size: Tuple[int, int] = None) -> ctk.CTkImage:
    image = Image.fromarray(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))
    image = ImageOps.fit(image, (IMAGE_BUTTON_WIDTH, IMAGE_BUTTON_HEIGHT), Image.LANCZOS)
    return ctk.CTkImage(image, size=image.size)


def render_video_preview(video_path: str, size: Tuple[int, int], frame_number: int = 0) -> ctk.CTkImage:
    capture = cv2.VideoCapture(video_path)
    if frame_number:
        capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
    has_frame, frame = capture.read()
    if has_frame:
        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        if size:
            image = ImageOps.fit(image, size, Image.LANCZOS)
        return ctk.CTkImage(image, size=image.size)
    capture.release()
    cv2.destroyAllWindows()


def toggle_preview() -> None:
    if PREVIEW.state() == 'normal':
        PREVIEW.withdraw()
    elif roop.globals.source_path and roop.globals.target_path:
        init_preview()
        update_preview()
        PREVIEW.deiconify()


def init_preview() -> None:
    if is_image(roop.globals.target_path):
        preview_slider.pack_forget()
    if is_video(roop.globals.target_path):
        video_frame_total = get_video_frame_total(roop.globals.target_path)
        preview_slider.configure(to=video_frame_total)
        preview_slider.pack(fill='x')
        preview_slider.set(0)


def update_preview(frame_number: int = 0) -> None:
    if roop.globals.source_path and roop.globals.target_path:
        temp_frame = get_video_frame(roop.globals.target_path, frame_number)
        swap_mode = "selected"
        processors = "faceswap"
        if roop.globals.selected_enhancer == 'GFPGAN':
            processors += ",gfpgan"
        elif roop.globals.selected_enhancer == 'Codeformer':
            processors += ",codeformer"

        temp_frame, _ = roop.globals.IMAGE_CHAIN_PROCESSOR.run_chain(temp_frame,  
                                                        {"swap_mode": swap_mode,
                                                            "input_face_datas": [roop.globals.SELECTED_FACE_DATA_INPUT], "target_face_datas": [roop.globals.SELECTED_FACE_DATA_OUTPUT]},
                                                            processors)
        image = Image.fromarray(cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB))
        image = ImageOps.contain(image, (PREVIEW_MAX_WIDTH, PREVIEW_MAX_HEIGHT), Image.LANCZOS)
        image = ctk.CTkImage(image, size=image.size)
        preview_label.configure(image=image)


def create_select_faces_win(parent) -> ctk.CTkToplevel:
    global scrollable_frame

    face_win = ctk.CTkToplevel(parent)
    face_win.minsize(800, 400)
    face_win.title('Select Face(s)')
    face_win.configure()
    face_win.withdraw()
    face_win.protocol('WM_DELETE_WINDOW', lambda: cancel_face_selection())
    scrollable_frame = ctk.CTkScrollableFrame(face_win, orientation='horizontal', label_text='Choose face by clicking on it', width=(IMAGE_BUTTON_WIDTH + 40)*3, height=IMAGE_BUTTON_HEIGHT+32)
    scrollable_frame.grid(row=0, column=0, padx=20, pady=20)
    scrollable_frame.place(relx=0.05, rely=0.05)
    cancel_button = ctk.CTkButton(face_win, text='Cancel', command=lambda: cancel_face_selection())
    cancel_button.place(relx=0.05, rely=0.85, relwidth=0.075, relheight=0.075)

    return face_win

def cancel_face_selection() -> None:
    toggle_face_selection();
    ROOT.wm_attributes('-disabled', False)
    ROOT.focus()

def select_face(index, is_input) -> None:
    global source_button, target_button, INPUT_FACES_DATA, OUTPUT_FACES_DATA

    if is_input:
        roop.globals.SELECTED_FACE_DATA_INPUT = INPUT_FACES_DATA[index][0]
        image = render_face_from_frame(INPUT_FACES_DATA[index][1], (IMAGE_BUTTON_WIDTH, IMAGE_BUTTON_HEIGHT))
        source_button.configure(image=image)
        source_button._draw()
    else:
        roop.globals.SELECTED_FACE_DATA_OUTPUT = OUTPUT_FACES_DATA[index][0]
        image = render_face_from_frame(OUTPUT_FACES_DATA[index][1], (IMAGE_BUTTON_WIDTH, IMAGE_BUTTON_HEIGHT))
        target_button.configure(image=image)
        target_button._draw()
        if roop.globals.SELECTED_FACE_DATA_INPUT is not None:
            emb1 = roop.globals.SELECTED_FACE_DATA_INPUT.embedding
            emb2 = roop.globals.SELECTED_FACE_DATA_OUTPUT.embedding
            dist = compute_cosine_distance(emb1, emb2)
            print(f'Similarity Distance between Source->Target={dist}')

    toggle_face_selection();
    ROOT.wm_attributes('-disabled', False)
    ROOT.focus()



def toggle_face_selection() -> None:
    if FACE_SELECT.state() == 'normal':
        FACE_SELECT.withdraw()
    else:
        FACE_SELECT.deiconify()


def show_face_selection(faces, is_input):
    global FACE_BUTTONS, scrollable_frame

    ROOT.wm_attributes('-disabled', True)

    if len(FACE_BUTTONS) > 0:
        for b in FACE_BUTTONS:
            try:
                # b.place_forget()
                b.destroy()
            except:
                continue
        FACE_BUTTONS.clear()

    for i,face in enumerate(faces):
        image = render_face_from_frame(face[1], (128, 128))
        score = face[0]['det_score']
        age = face[0]['age']
        button_text = f'Score: {score} - Sex: {face[0].sex} - Age: {age}'
        face_button = ctk.CTkButton(scrollable_frame, text=button_text, width=128, height=128, compound='top', anchor='center', command=lambda faceindex=i: select_face(index=faceindex, is_input=is_input))
        face_button.grid(row=0, column=i, pady=5, padx=5)
        face_button.configure(image=image)
        face_button._draw()
        FACE_BUTTONS.append(face_button)

    FACE_SELECT.deiconify()