mr2along commited on
Commit
fdd0003
·
verified ·
1 Parent(s): 5269076

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +70 -56
app.py CHANGED
@@ -1,6 +1,5 @@
1
- import os, cv2, glob, time, torch, shutil
2
- import argparse, numpy as np, concurrent.futures
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
- # ----------------- USER ARGS -----------------
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
- # ----------------- MODELS -----------------
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
- # ----------------- SWAP LOGIC -----------------
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
- # Swap
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
- # Enhance
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
- # Mask
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
- # Paste back
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
- # ----------------- PROCESS MAIN -----------------
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
- # --- chuẩn bị input ---
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
- # --- Swap nữ, nam ---
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
- # ----------------- GUI -----------------
185
- with gr.Blocks(css="footer{display:none !important}") as interface:
 
 
186
  gr.Markdown("# 🗿 Free Face Swapper (2 faces)")
 
187
  with gr.Row():
188
- with gr.Column(scale=0.4):
189
- swap_option = gr.Dropdown(swap_options_list, value=swap_options_list[0], label="Swap Condition")
190
- age = gr.Number(value=25, label="Age", visible=False)
191
- output_directory = gr.Text(value=DEF_OUTPUT_PATH, label="Output Dir")
192
- output_name = gr.Text(value="Result", label="Output Name")
193
- keep_output_sequence = gr.Checkbox(label="Keep sequence", value=False)
194
- face_scale = gr.Slider(0, 2, 1, label="Face Scale")
195
- face_enhancer_name = gr.Dropdown(FACE_ENHANCER_LIST, value="NONE", label="Face Enhancer")
196
- enable_face_parser_mask = gr.Checkbox(label="Enable Face Parsing", value=False)
197
- mask_include = gr.Dropdown(mask_regions.keys(), value=list(mask_regions.keys()), multiselect=True, label="Mask Include")
198
- mask_soft_kernel = gr.Number(value=17, label="Soft Kernel", visible=False)
199
- mask_soft_iterations = gr.Number(value=10, label="Soft Iterations")
200
- crop_top = gr.Slider(0, 511, 0, label="Crop Top")
201
- crop_bott = gr.Slider(0, 511, 511, label="Crop Bottom")
202
- crop_left = gr.Slider(0, 511, 0, label="Crop Left")
203
- crop_right = gr.Slider(0, 511, 511, label="Crop Right")
204
- erode_amount = gr.Slider(0, 1, 0.15, step=0.05, label="Mask Erode")
205
- blur_amount = gr.Slider(0, 1, 0.1, step=0.05, label="Mask Blur")
206
- enable_laplacian_blend = gr.Checkbox(label="Laplacian Blend", value=True)
207
-
208
- gr.Markdown("### 👩 Source Female")
209
- source_female_input = gr.Image(label="Source Female", type="filepath")
210
- gr.Markdown("### 👨 Source Male")
211
- source_male_input = gr.Image(label="Source Male", type="filepath")
212
-
213
- input_type = gr.Radio(["Image", "Video", "Directory"], value="Image", label="Target Type")
214
- image_input = gr.Image(label="Target Image", type="filepath", visible=True)
215
- video_input = gr.Video(label="Target Video", visible=False)
216
- direc_input = gr.Text(label="Target Directory", visible=False)
217
-
218
- with gr.Column(scale=0.6):
219
- info = gr.Markdown("...")
220
- with gr.Row():
221
- swap_button = gr.Button(" Swap", variant="primary")
222
- cancel_button = gr.Button("⛔ Cancel")
223
- preview_image = gr.Image(label="Output Image", interactive=False)
224
- preview_video = gr.Video(label="Output Video", interactive=False, visible=False)
225
- with gr.Row():
226
- output_directory_button = gr.Button("📂", visible=False)
227
- output_video_button = gr.Button("🎬", visible=False)
 
 
 
 
 
 
 
 
 
 
 
 
 
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,