ui/tabs/facemgr_tab.py

import os
import shutil
import cv2
import gradio as gr
import roop.utilities as util
import roop.globals
from roop.face_util import extract_face_images
from roop.capturer import get_video_frame, get_video_frame_total
from typing import List, Tuple, Optional
from roop.typing import Frame, Face, FaceSet

selected_face_index = -1
thumbs = []
images = []


def facemgr_tab() -> None:
    with gr.Tab("👨‍👩‍👧‍👦 Face Management"):
        with gr.Row():
            gr.Markdown("""
                        # Create blending facesets
                        Add multiple reference images into a faceset file.
                        """)
        with gr.Row():
            videoimagefst = gr.Image(label="Cut face from video frame", height=576, interactive=False, visible=True, format="jpeg")
        with gr.Row():
            frame_num_fst = gr.Slider(1, 1, value=1, label="Frame Number", info='0:00:00', step=1.0, interactive=False)
            fb_cutfromframe = gr.Button("Use faces from this frame", variant='secondary', interactive=False)
        with gr.Row():
            fb_facesetfile = gr.Files(label='Faceset', file_count='single', file_types=['.fsz'], interactive=True)
            fb_files = gr.Files(label='Input Files', file_count="multiple", file_types=["image", "video"], interactive=True)
        with gr.Row():
            with gr.Column():
                gr.Button("👀 Open Output Folder", size='sm').click(fn=lambda: util.open_folder(roop.globals.output_path))
            with gr.Column():
                gr.Markdown(' ')
        with gr.Row():
            faces = gr.Gallery(label="Faces in this Faceset", allow_preview=True, preview=True, height=128, object_fit="scale-down")
        with gr.Row():
            fb_remove = gr.Button("Remove selected", variant='secondary')
            fb_update = gr.Button("Create/Update Faceset file", variant='primary')
            fb_clear = gr.Button("Clear all", variant='stop')

    fb_facesetfile.change(fn=on_faceset_changed, inputs=[fb_facesetfile], outputs=[faces])
    fb_files.change(fn=on_fb_files_changed, inputs=[fb_files], outputs=[faces, videoimagefst, frame_num_fst, fb_cutfromframe])
    fb_update.click(fn=on_update_clicked, outputs=[fb_facesetfile])
    fb_remove.click(fn=on_remove_clicked, outputs=[faces])
    fb_clear.click(fn=on_clear_clicked, outputs=[faces, fb_files, fb_facesetfile])
    fb_cutfromframe.click(fn=on_cutfromframe_clicked, inputs=[fb_files, frame_num_fst], outputs=[faces])
    frame_num_fst.release(fn=on_frame_num_fst_changed, inputs=[fb_files, frame_num_fst], outputs=[videoimagefst])
    faces.select(fn=on_face_selected)


def on_faceset_changed(faceset, progress=gr.Progress()) -> List[Frame]:
    global thumbs, images

    if faceset is None:
        return thumbs

    thumbs.clear()
    filename = faceset.name
        
    if filename.lower().endswith('fsz'):
        progress(0, desc="Retrieving faces from Faceset File", )      
        unzipfolder = os.path.join(os.environ["TEMP"], 'faceset')
        if os.path.isdir(unzipfolder):
            shutil.rmtree(unzipfolder)
        util.mkdir_with_umask(unzipfolder)
        util.unzip(filename, unzipfolder)
        for file in os.listdir(unzipfolder):
            if file.endswith(".png"):
                SELECTION_FACES_DATA = extract_face_images(os.path.join(unzipfolder,file),  (False, 0), 0.5)
                if len(SELECTION_FACES_DATA) < 1:
                    gr.Warning(f"No face detected in {file}!")
                for f in SELECTION_FACES_DATA:
                    image = f[1]
                    images.append(image)
                    thumbs.append(util.convert_to_gradio(image))
        
        return thumbs


def on_fb_files_changed(inputfiles, progress=gr.Progress()) -> Tuple[List[Frame], Optional[gr.Image], Optional[gr.Slider], Optional[gr.Button]]:
    global thumbs, images, total_frames, current_video_fps

    if inputfiles is None or len(inputfiles) < 1:
        return thumbs, None, None, None
    
    progress(0, desc="Retrieving faces from images", )
    slider = None
    video_image = None
    cut_button = None
    for f in inputfiles:
        source_path = f.name
        if util.has_image_extension(source_path):
            slider = gr.Slider(interactive=False)
            video_image = gr.Image(interactive=False)
            cut_button = gr.Button(interactive=False)
            roop.globals.source_path = source_path
            SELECTION_FACES_DATA = extract_face_images(roop.globals.source_path,  (False, 0), 0.5)
            for f in SELECTION_FACES_DATA:
                image = f[1]
                images.append(image)
                thumbs.append(util.convert_to_gradio(image))
        elif util.is_video(source_path) or source_path.lower().endswith('gif'):
            total_frames = get_video_frame_total(source_path)
            current_video_fps = util.detect_fps(source_path)
            cut_button = gr.Button(interactive=True)
            video_image, slider = display_video_frame(source_path, 1, total_frames)

    return thumbs, video_image, slider, cut_button
    

def display_video_frame(filename: str, frame_num: int, total: int=0) -> Tuple[gr.Image, gr.Slider]:
    global current_video_fps

    current_frame = get_video_frame(filename, frame_num)
    if current_video_fps == 0:
        current_video_fps = 1
    secs = (frame_num - 1) / current_video_fps
    minutes = secs / 60
    secs = secs % 60
    hours = minutes / 60
    minutes = minutes % 60
    milliseconds = (secs - int(secs)) * 1000
    timeinfo = f"{int(hours):0>2}:{int(minutes):0>2}:{int(secs):0>2}.{int(milliseconds):0>3}"
    if total > 0:
        return gr.Image(value=util.convert_to_gradio(current_frame), interactive=True), gr.Slider(info=timeinfo, minimum=1, maximum=total, interactive=True)  
    return gr.Image(value=util.convert_to_gradio(current_frame), interactive=True), gr.Slider(info=timeinfo, interactive=True)  


def on_face_selected(evt: gr.SelectData) -> None:
    global selected_face_index

    if evt is not None:
        selected_face_index = evt.index

def on_frame_num_fst_changed(inputfiles: List[gr.Files], frame_num: int) -> Frame:
    filename = inputfiles[0].name
    video_image, _ = display_video_frame(filename, frame_num, 0)
    return video_image


def on_cutfromframe_clicked(inputfiles: List[gr.Files], frame_num: int) -> List[Frame]:
    global thumbs

    filename = inputfiles[0].name
    SELECTION_FACES_DATA = extract_face_images(filename,  (True, frame_num), 0.5)
    for f in SELECTION_FACES_DATA:
        image = f[1]
        images.append(image)
        thumbs.append(util.convert_to_gradio(image))
    return thumbs


def on_remove_clicked() -> List[Frame]:
    global thumbs, images, selected_face_index

    if len(thumbs) > selected_face_index:
        f = thumbs.pop(selected_face_index)
        del f
        f = images.pop(selected_face_index)
        del f
    return thumbs

def on_clear_clicked() -> Tuple[List[Frame], None, None]:
    global thumbs, images

    thumbs.clear()
    images.clear()
    return thumbs, None, None


def on_update_clicked() -> Optional[str]:
    if len(images) < 1:
        gr.Warning(f"No faces to create faceset from!")
        return None

    imgnames = []
    for index,img in enumerate(images):
        filename = os.path.join(roop.globals.output_path, f'{index}.png')
        cv2.imwrite(filename, img)
        imgnames.append(filename)

    finalzip = os.path.join(roop.globals.output_path, 'faceset.fsz')        
    util.zip(imgnames, finalzip)
    return finalzip