File size: 7,581 Bytes
fcc8e7d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 |
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
|