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