Update app.py

app.py

@@ -5,18 +5,18 @@ import torch
 import argparse
 import insightface
 import onnxruntime
-import numpy as np
 import gradio as gr
 from tqdm import tqdm
 
 from face_swapper import Inswapper, paste_to_whole
 from face_analyser import analyse_face
-from face_enhancer import load_face_enhancer_model, cv2_interpolations
-from utils import create_image_grid
-
-## ------------------------------ USER ARGS ------------------------------
+from face_enhancer import (
+    load_face_enhancer_model,
+    get_available_enhancer_names,
+)
 
-parser = argparse.ArgumentParser(description="Free Face Swapper (Male/Female mode)")
+# ------------------------------ ARGS ------------------------------
+parser = argparse.ArgumentParser(description="Face Swapper (Male+Female with Enhancers)")
 parser.add_argument("--out_dir", default=os.getcwd())
 parser.add_argument("--batch_size", default=32)
 parser.add_argument("--cuda", action="store_true", default=False)
@@ -26,22 +26,19 @@ USE_CUDA = user_args.cuda
 DEF_OUTPUT_PATH = user_args.out_dir
 BATCH_SIZE = int(user_args.batch_size)
 
-## ------------------------------ DEVICE ------------------------------
-
+# ------------------------------ DEVICE ------------------------------
 PROVIDER = ["CPUExecutionProvider"]
-if USE_CUDA:
-    if "CUDAExecutionProvider" in onnxruntime.get_available_providers():
-        PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"]
-        print(">>> Running on CUDA")
-    else:
-        USE_CUDA = False
-        print(">>> CUDA not available, running on CPU")
+if USE_CUDA and "CUDAExecutionProvider" in onnxruntime.get_available_providers():
+    PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    print(">>> Running on CUDA")
+else:
+    USE_CUDA = False
+    print(">>> Running on CPU")
 
 device = "cuda" if USE_CUDA else "cpu"
 EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
 
-## ------------------------------ LOAD MODELS ------------------------------
-
+# ------------------------------ MODELS ------------------------------
 FACE_ANALYSER = insightface.app.FaceAnalysis(name="buffalo_l", providers=PROVIDER)
 FACE_ANALYSER.prepare(ctx_id=0, det_size=(640, 640), det_thresh=0.6)
 
@@ -51,64 +48,79 @@ FACE_SWAPPER = Inswapper(
     providers=PROVIDER,
 )
 
-## ------------------------------ PROCESS ------------------------------
+# ------------------------------ ENHANCERS ------------------------------
+ENHANCER_CHOICES = ["NONE"] + get_available_enhancer_names()
+# ví dụ: ["NONE", "CodeFormer", "GFPGAN", "REAL-ESRGAN 2x", "REAL-ESRGAN 4x", "REAL-ESRGAN 8x"]
 
-def swap_faces(image_path, male_source_path, female_source_path, face_enhancer_name="NONE"):
+# ------------------------------ PROCESS ------------------------------
+def swap_faces(image_path, male_source_path, female_source_path, enhancer_name="NONE"):
     start_time = time.time()
 
-    # Load target
-    target = cv2.imread(image_path)
+    # Load target & sources
+    target_bgr = cv2.imread(image_path)
+    if target_bgr is None:
+        raise ValueError("Không đọc được ảnh đích.")
+
+    src_male_img = cv2.imread(male_source_path)
+    src_female_img = cv2.imread(female_source_path)
+    if src_male_img is None or src_female_img is None:
+        raise ValueError("Không đọc được ảnh nguồn (nam/nữ).")
 
-    # Load source male/female
-    analysed_source_male = analyse_face(cv2.imread(male_source_path), FACE_ANALYSER)
-    analysed_source_female = analyse_face(cv2.imread(female_source_path), FACE_ANALYSER)
+    analysed_source_male = analyse_face(src_male_img, FACE_ANALYSER)
+    analysed_source_female = analyse_face(src_female_img, FACE_ANALYSER)
 
-    # Analyse target
-    analysed_faces = FACE_ANALYSER.get(target)
+    # Phân tích các khuôn mặt trong ảnh đích
+    analysed_faces = FACE_ANALYSER.get(target_bgr)
 
     preds, matrs = [], []
     for analysed_face in tqdm(analysed_faces, desc="Swapping faces"):
-        if analysed_face["gender"] == 1:  # male
-            src = analysed_source_male
-        else:  # female
-            src = analysed_source_female
-
-        batch_pred, batch_matr = FACE_SWAPPER.get([target], [analysed_face], [src])
+        # gender: 1 = male, 0 = female (theo insightface)
+        src_face = analysed_source_male if analysed_face.get("gender", 1) == 1 else analysed_source_female
+        batch_pred, batch_matr = FACE_SWAPPER.get([target_bgr], [analysed_face], [src_face])
         preds.extend(batch_pred)
         matrs.extend(batch_matr)
         EMPTY_CACHE()
 
-    # Paste back
+    # Ghép lại
     for p, m in zip(preds, matrs):
-        target = paste_to_whole(p, target, m, blend_method="laplacian")
-
-    # Enhance (optional)
-    if face_enhancer_name != "NONE":
-        model, runner = load_face_enhancer_model(face_enhancer_name, device=device)
-        target = runner(target, model)
-
-    elapsed = time.time() - start_time
-    print(f"✔ Done in {elapsed:.2f} sec")
-    return target[:, :, ::-1]  # BGR->RGB for display
-
-
-## ------------------------------ GRADIO UI ------------------------------
-
+        target_bgr = paste_to_whole(
+            foreground=p,
+            background=target_bgr,
+            matrix=m,
+            mask=None,
+            crop_mask=(0, 0, 0, 0),
+            blur_amount=0.1,
+            erode_amount=0.15,
+            blend_method="laplacian"  # tự nhiên hơn
+        )
+
+    # Enhance (nếu chọn)
+    if enhancer_name != "NONE":
+        try:
+            model, runner = load_face_enhancer_model(name=enhancer_name, device=device)
+            target_bgr = runner(target_bgr, model)
+        except AssertionError as e:
+            print(f"[Enhancer] {e}. Trả về ảnh không enhance.")
+        except Exception as e:
+            print(f"[Enhancer] Lỗi khi chạy {enhancer_name}: {e}. Trả về ảnh không enhance.")
+
+    print(f"✔ Hoàn tất trong {time.time() - start_time:.2f}s")
+    return target_bgr[:, :, ::-1]  # BGR -> RGB để hiển thị
+
+# ------------------------------ UI ------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧑➡👩 Face Swapper (Male+Female sources)")
+    gr.Markdown("## 🧑‍🦱➡👩 Face Swapper (2 nguồn nam/nữ) + Enhancer (CodeFormer / GFPGAN / Real-ESRGAN)")
 
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(label="Target Image", type="filepath")
-            male_input = gr.Image(label="Source Male", type="filepath")
-            female_input = gr.Image(label="Source Female", type="filepath")
-            enhancer = gr.Dropdown(
-                ["NONE"] + cv2_interpolations, label="Face Enhancer", value="NONE"
-            )
+            image_input = gr.Image(label="Ảnh đích (Target Image)", type="filepath")
+            male_input = gr.Image(label="Ảnh nguồn Nam", type="filepath")
+            female_input = gr.Image(label="Ảnh nguồn Nữ", type="filepath")
+            enhancer = gr.Dropdown(ENHANCER_CHOICES, label="Face Enhancer", value="NONE")
             run_btn = gr.Button("✨ Swap")
 
         with gr.Column():
-            output_image = gr.Image(label="Output")
+            output_image = gr.Image(label="Kết quả")
 
     run_btn.click(
         fn=swap_faces,