Update app.py
Browse files
app.py
CHANGED
@@ -1,6 +1,5 @@
|
|
1 |
-
import os, cv2, glob, time, torch, shutil
|
2 |
-
import
|
3 |
-
import insightface, onnxruntime, gradio as gr
|
4 |
from moviepy.editor import VideoFileClip
|
5 |
from face_swapper import Inswapper, paste_to_whole
|
6 |
from face_analyser import detect_conditions, get_analysed_data, swap_options_list
|
@@ -8,7 +7,7 @@ from face_parsing import init_parsing_model, get_parsed_mask, mask_regions, mask
|
|
8 |
from face_enhancer import get_available_enhancer_names, load_face_enhancer_model, cv2_interpolations
|
9 |
from utils import merge_img_sequence_from_ref, open_directory, split_list_by_lengths, create_image_grid
|
10 |
|
11 |
-
#
|
12 |
parser = argparse.ArgumentParser(description="Free Face Swapper (2 faces)")
|
13 |
parser.add_argument("--out_dir", default=os.getcwd())
|
14 |
parser.add_argument("--batch_size", default=32)
|
@@ -21,7 +20,7 @@ DEF_OUTPUT_PATH, BATCH_SIZE = args.out_dir, int(args.batch_size)
|
|
21 |
device = "cuda" if USE_CUDA else "cpu"
|
22 |
EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
|
23 |
|
24 |
-
#
|
25 |
PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"] if USE_CUDA else ["CPUExecutionProvider"]
|
26 |
FACE_ANALYSER, FACE_SWAPPER, FACE_PARSER = None, None, None
|
27 |
FACE_ENHANCER_LIST = ["NONE"] + get_available_enhancer_names() + cv2_interpolations
|
@@ -43,12 +42,11 @@ def load_face_parser_model(path="./assets/pretrained_models/79999_iter.pth"):
|
|
43 |
if FACE_PARSER is None:
|
44 |
FACE_PARSER = init_parsing_model(path, device=device)
|
45 |
|
46 |
-
#
|
47 |
def run_face_swap(image_sequence, source_path, label, condition, age, distance, face_scale,
|
48 |
face_enhancer_name, enable_face_parser, mask_includes,
|
49 |
mask_soft_kernel, mask_soft_iterations, blur_amount, erode_amount,
|
50 |
enable_laplacian_blend, crop_mask):
|
51 |
-
global PREVIEW
|
52 |
yield f"### ⌛ {label}: phân tích khuôn mặt..."
|
53 |
source_data = (source_path, age) if condition != "Specific Face" else (([], []), distance)
|
54 |
|
@@ -57,18 +55,18 @@ def run_face_swap(image_sequence, source_path, label, condition, age, distance,
|
|
57 |
swap_condition=condition, detect_condition="best detection", scale=face_scale
|
58 |
)
|
59 |
|
60 |
-
#
|
61 |
yield f"### ⌛ {label}: hoán đổi..."
|
62 |
preds, matrs = [], []
|
63 |
for bp, bm in FACE_SWAPPER.batch_forward(whole_frame_list, analysed_targets, analysed_sources):
|
64 |
preds.extend(bp); matrs.extend(bm); EMPTY_CACHE()
|
65 |
|
66 |
-
#
|
67 |
if face_enhancer_name != "NONE":
|
68 |
enhancer_model, enhancer_runner = load_face_enhancer_model(name=face_enhancer_name, device=device)
|
69 |
preds = [cv2.resize(enhancer_runner(p, enhancer_model), (512, 512)) for p in preds]
|
70 |
|
71 |
-
#
|
72 |
masks = [None] * len(preds)
|
73 |
if enable_face_parser:
|
74 |
yield f"### ⌛ {label}: tạo mask..."
|
@@ -78,7 +76,7 @@ def run_face_swap(image_sequence, source_path, label, condition, age, distance,
|
|
78 |
masks.append(bm); EMPTY_CACHE()
|
79 |
masks = np.concatenate(masks, axis=0) if masks else [None] * len(preds)
|
80 |
|
81 |
-
#
|
82 |
yield f"### ⌛ {label}: dán khuôn mặt..."
|
83 |
sp_preds = split_list_by_lengths(preds, num_faces_per_frame)
|
84 |
sp_matrs = split_list_by_lengths(matrs, num_faces_per_frame)
|
@@ -98,7 +96,7 @@ def run_face_swap(image_sequence, source_path, label, condition, age, distance,
|
|
98 |
with concurrent.futures.ThreadPoolExecutor() as exe:
|
99 |
list(exe.map(post_process, range(len(image_sequence)), image_sequence))
|
100 |
|
101 |
-
#
|
102 |
def process(input_type, image_path, video_path, directory_path,
|
103 |
source_female_path, source_male_path,
|
104 |
output_path, output_name, keep_output_sequence,
|
@@ -119,7 +117,7 @@ def process(input_type, image_path, video_path, directory_path,
|
|
119 |
load_face_enhancer_model(name=face_enhancer_name, device=device)
|
120 |
if enable_face_parser: load_face_parser_model()
|
121 |
|
122 |
-
#
|
123 |
image_sequence, output_file = [], None
|
124 |
if input_type == "Image":
|
125 |
output_file = os.path.join(output_path, output_name + ".png")
|
@@ -145,7 +143,7 @@ def process(input_type, image_path, video_path, directory_path,
|
|
145 |
else:
|
146 |
yield "### ❌ Stream chưa hỗ trợ"; return
|
147 |
|
148 |
-
#
|
149 |
has_swap = False
|
150 |
if source_female_path and os.path.exists(source_female_path):
|
151 |
has_swap = True
|
@@ -181,50 +179,66 @@ def process(input_type, image_path, video_path, directory_path,
|
|
181 |
gr.update(value=PREVIEW, visible=True), gr.update(interactive=True), \
|
182 |
gr.update(interactive=True), gr.update(value=OUTPUT_FILE, visible=True)
|
183 |
|
184 |
-
#
|
185 |
-
|
|
|
|
|
186 |
gr.Markdown("# 🗿 Free Face Swapper (2 faces)")
|
|
|
187 |
with gr.Row():
|
188 |
-
with gr.
|
189 |
-
|
190 |
-
|
191 |
-
|
192 |
-
|
193 |
-
|
194 |
-
|
195 |
-
|
196 |
-
|
197 |
-
|
198 |
-
|
199 |
-
|
200 |
-
|
201 |
-
|
202 |
-
|
203 |
-
|
204 |
-
|
205 |
-
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
|
211 |
-
|
212 |
-
|
213 |
-
|
214 |
-
|
215 |
-
|
216 |
-
|
217 |
-
|
218 |
-
|
219 |
-
|
220 |
-
|
221 |
-
|
222 |
-
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
|
227 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
228 |
|
229 |
swap_inputs = [
|
230 |
input_type, image_input, video_input, direc_input,
|
|
|
1 |
+
import os, cv2, glob, time, torch, shutil, argparse, concurrent.futures
|
2 |
+
import numpy as np, insightface, onnxruntime, gradio as gr
|
|
|
3 |
from moviepy.editor import VideoFileClip
|
4 |
from face_swapper import Inswapper, paste_to_whole
|
5 |
from face_analyser import detect_conditions, get_analysed_data, swap_options_list
|
|
|
7 |
from face_enhancer import get_available_enhancer_names, load_face_enhancer_model, cv2_interpolations
|
8 |
from utils import merge_img_sequence_from_ref, open_directory, split_list_by_lengths, create_image_grid
|
9 |
|
10 |
+
# ---------------- ARGS ----------------
|
11 |
parser = argparse.ArgumentParser(description="Free Face Swapper (2 faces)")
|
12 |
parser.add_argument("--out_dir", default=os.getcwd())
|
13 |
parser.add_argument("--batch_size", default=32)
|
|
|
20 |
device = "cuda" if USE_CUDA else "cpu"
|
21 |
EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
|
22 |
|
23 |
+
# ---------------- MODELS ----------------
|
24 |
PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"] if USE_CUDA else ["CPUExecutionProvider"]
|
25 |
FACE_ANALYSER, FACE_SWAPPER, FACE_PARSER = None, None, None
|
26 |
FACE_ENHANCER_LIST = ["NONE"] + get_available_enhancer_names() + cv2_interpolations
|
|
|
42 |
if FACE_PARSER is None:
|
43 |
FACE_PARSER = init_parsing_model(path, device=device)
|
44 |
|
45 |
+
# ---------------- SWAP ----------------
|
46 |
def run_face_swap(image_sequence, source_path, label, condition, age, distance, face_scale,
|
47 |
face_enhancer_name, enable_face_parser, mask_includes,
|
48 |
mask_soft_kernel, mask_soft_iterations, blur_amount, erode_amount,
|
49 |
enable_laplacian_blend, crop_mask):
|
|
|
50 |
yield f"### ⌛ {label}: phân tích khuôn mặt..."
|
51 |
source_data = (source_path, age) if condition != "Specific Face" else (([], []), distance)
|
52 |
|
|
|
55 |
swap_condition=condition, detect_condition="best detection", scale=face_scale
|
56 |
)
|
57 |
|
58 |
+
# swap
|
59 |
yield f"### ⌛ {label}: hoán đổi..."
|
60 |
preds, matrs = [], []
|
61 |
for bp, bm in FACE_SWAPPER.batch_forward(whole_frame_list, analysed_targets, analysed_sources):
|
62 |
preds.extend(bp); matrs.extend(bm); EMPTY_CACHE()
|
63 |
|
64 |
+
# enhance
|
65 |
if face_enhancer_name != "NONE":
|
66 |
enhancer_model, enhancer_runner = load_face_enhancer_model(name=face_enhancer_name, device=device)
|
67 |
preds = [cv2.resize(enhancer_runner(p, enhancer_model), (512, 512)) for p in preds]
|
68 |
|
69 |
+
# mask
|
70 |
masks = [None] * len(preds)
|
71 |
if enable_face_parser:
|
72 |
yield f"### ⌛ {label}: tạo mask..."
|
|
|
76 |
masks.append(bm); EMPTY_CACHE()
|
77 |
masks = np.concatenate(masks, axis=0) if masks else [None] * len(preds)
|
78 |
|
79 |
+
# paste back
|
80 |
yield f"### ⌛ {label}: dán khuôn mặt..."
|
81 |
sp_preds = split_list_by_lengths(preds, num_faces_per_frame)
|
82 |
sp_matrs = split_list_by_lengths(matrs, num_faces_per_frame)
|
|
|
96 |
with concurrent.futures.ThreadPoolExecutor() as exe:
|
97 |
list(exe.map(post_process, range(len(image_sequence)), image_sequence))
|
98 |
|
99 |
+
# ---------------- PROCESS ----------------
|
100 |
def process(input_type, image_path, video_path, directory_path,
|
101 |
source_female_path, source_male_path,
|
102 |
output_path, output_name, keep_output_sequence,
|
|
|
117 |
load_face_enhancer_model(name=face_enhancer_name, device=device)
|
118 |
if enable_face_parser: load_face_parser_model()
|
119 |
|
120 |
+
# chuẩn bị input
|
121 |
image_sequence, output_file = [], None
|
122 |
if input_type == "Image":
|
123 |
output_file = os.path.join(output_path, output_name + ".png")
|
|
|
143 |
else:
|
144 |
yield "### ❌ Stream chưa hỗ trợ"; return
|
145 |
|
146 |
+
# swap nữ / nam
|
147 |
has_swap = False
|
148 |
if source_female_path and os.path.exists(source_female_path):
|
149 |
has_swap = True
|
|
|
179 |
gr.update(value=PREVIEW, visible=True), gr.update(interactive=True), \
|
180 |
gr.update(interactive=True), gr.update(value=OUTPUT_FILE, visible=True)
|
181 |
|
182 |
+
# ---------------- GUI ----------------
|
183 |
+
css = """footer{display:none !important}"""
|
184 |
+
|
185 |
+
with gr.Blocks(css=css) as interface:
|
186 |
gr.Markdown("# 🗿 Free Face Swapper (2 faces)")
|
187 |
+
|
188 |
with gr.Row():
|
189 |
+
with gr.Row():
|
190 |
+
with gr.Column(scale=0.4):
|
191 |
+
with gr.Tab("📄 Swap Condition"):
|
192 |
+
swap_option = gr.Dropdown(swap_options_list, value=swap_options_list[0], label="Condition")
|
193 |
+
age = gr.Number(value=25, label="Age", visible=False)
|
194 |
+
|
195 |
+
with gr.Tab("🎚️ Detection Settings"):
|
196 |
+
detect_condition_dropdown = gr.Dropdown(detect_conditions, value="best detection", label="Condition")
|
197 |
+
detection_size = gr.Number(value=640, label="Detection Size")
|
198 |
+
detection_threshold = gr.Number(value=0.6, label="Detection Threshold")
|
199 |
+
apply_detection_settings = gr.Button("Apply settings")
|
200 |
+
|
201 |
+
with gr.Tab("📤 Output Settings"):
|
202 |
+
output_directory = gr.Text(value=DEF_OUTPUT_PATH, label="Output Directory")
|
203 |
+
output_name = gr.Text(value="Result", label="Output Name")
|
204 |
+
keep_output_sequence = gr.Checkbox(label="Keep sequence", value=False)
|
205 |
+
|
206 |
+
with gr.Tab("🪄 Other Settings"):
|
207 |
+
face_scale = gr.Slider(0, 2, 1, label="Face Scale")
|
208 |
+
face_enhancer_name = gr.Dropdown(FACE_ENHANCER_LIST, value="NONE", label="Face Enhancer")
|
209 |
+
enable_face_parser_mask = gr.Checkbox(label="Enable Face Parsing", value=False)
|
210 |
+
mask_include = gr.Dropdown(mask_regions.keys(), value=list(mask_regions.keys()), multiselect=True)
|
211 |
+
mask_soft_kernel = gr.Number(value=17, visible=False)
|
212 |
+
mask_soft_iterations = gr.Number(value=10, label="Soft Iterations")
|
213 |
+
crop_top = gr.Slider(0, 511, 0, label="Crop Top")
|
214 |
+
crop_bott = gr.Slider(0, 511, 511, label="Crop Bottom")
|
215 |
+
crop_left = gr.Slider(0, 511, 0, label="Crop Left")
|
216 |
+
crop_right = gr.Slider(0, 511, 511, label="Crop Right")
|
217 |
+
erode_amount = gr.Slider(0, 1, 0.15, step=0.05, label="Mask Erode")
|
218 |
+
blur_amount = gr.Slider(0, 1, 0.1, step=0.05, label="Mask Blur")
|
219 |
+
enable_laplacian_blend = gr.Checkbox(value=True, label="Laplacian Blend")
|
220 |
+
|
221 |
+
gr.Markdown("### 👩 Source Female")
|
222 |
+
source_female_input = gr.Image(label="Source Female Face", type="filepath")
|
223 |
+
|
224 |
+
gr.Markdown("### 👨 Source Male")
|
225 |
+
source_male_input = gr.Image(label="Source Male Face", type="filepath")
|
226 |
+
|
227 |
+
input_type = gr.Radio(["Image", "Video", "Directory"], value="Image", label="Target Type")
|
228 |
+
image_input = gr.Image(label="Target Image", type="filepath", visible=True)
|
229 |
+
video_input = gr.Video(label="Target Video", visible=False)
|
230 |
+
direc_input = gr.Text(label="Directory Path", visible=False)
|
231 |
+
|
232 |
+
with gr.Column(scale=0.6):
|
233 |
+
info = gr.Markdown("...")
|
234 |
+
with gr.Row():
|
235 |
+
swap_button = gr.Button("✨ Swap", variant="primary")
|
236 |
+
cancel_button = gr.Button("⛔ Cancel")
|
237 |
+
preview_image = gr.Image(label="Output", interactive=False)
|
238 |
+
preview_video = gr.Video(label="Output", interactive=False, visible=False)
|
239 |
+
with gr.Row():
|
240 |
+
output_directory_button = gr.Button("📂", visible=False)
|
241 |
+
output_video_button = gr.Button("🎬", visible=False)
|
242 |
|
243 |
swap_inputs = [
|
244 |
input_type, image_input, video_input, direc_input,
|