Update app.py

app.py

@@ -1,182 +1,173 @@
-
-import gradio as gr
-from overlay import overlay_source
-from detect_face import predict, NUM_CLASSES
-from swapface import swap_face_now
-import os
-from pathlib import Path
-
-BASE_DIR = Path(__file__).parent  # thư mục chứa app.py
-FOLDER = BASE_DIR / "example_wigs"
-
-# --- Hàm load ảnh từ folder ---
-def load_images_from_folder(folder_path: str) -> list[str]:
-    """
-    Trả về list[str] chứa tất cả các hình (jpg, png, gif, bmp) trong folder_path.
-    """
-    supported = {'.jpg', '.jpeg', '.png', '.gif', '.bmp'}
-    if not os.path.isdir(folder_path):
-        print(f"Cảnh báo: '{folder_path}' không phải folder hợp lệ.")
-        return []
-    files = [
-        os.path.join(folder_path, fn)
-        for fn in os.listdir(folder_path)
-        if os.path.splitext(fn)[1].lower() in supported
-    ]
-    if not files:
-        print(f"Không tìm thấy hình trong: {folder_path}")
-    return files
-
-def on_gallery_select(evt: gr.SelectData):
-    """
-    Khi click thumbnail: trả về
-      1) filepath để nạp vào Image Source
-      2) tên file (basename) để hiển thị trong Textbox
-    """
-    val = evt.value
-
-    # --- logic trích filepath y như cũ ---
-    if isinstance(val, dict):
-        img = val.get("image")
-        if isinstance(img, str):
-            filepath = img
-        elif isinstance(img, dict):
-            filepath = img.get("path") or img.get("url")
-        else:
-            filepath = next(
-                (v for v in val.values() if isinstance(v, str) and os.path.isfile(v)),
-                None
-            )
-    elif isinstance(val, str):
-        filepath = val
-    else:
-        raise ValueError(f"Kiểu không hỗ trợ: {type(val)}")
-
-    filename = os.path.basename(filepath) if filepath else ""
-    return filepath, filename
-
-# --- Hàm xác định folder dựa trên phân lớp ---
-def infer_folder(image) -> str:
-    cls = predict(image)["predicted_class"]
-    folder = str(FOLDER / cls)
-    return folder
-
-# --- Hàm gộp: phân loại + load ảnh ---
-def handle_bg_change(image):
-    """
-    Khi thay đổi background:
-    1. Phân loại khuôn mặt
-    2. Load ảnh từ folder tương ứng
-    """
-    if image is None:
-        return "", []
-    
-    try:
-        folder = infer_folder(image)
-        images = load_images_from_folder(folder)
-        return folder, images
-    except Exception as e:
-        print(f"Lỗi xử lý ảnh: {e}")
-        return "", []
-
-# --- Hàm swap face ---
-def swap_face_wrapper(background_img, result_img):
-    """
-    Wrapper function cho swap face giữa background và result image
-    """
-    if background_img is None or result_img is None:
-        return None
-    
-    try:
-        # Swap face từ background vào result image
-        swapped = swap_face_now(background_img, result_img, do_enhance=True)
-        return swapped
-    except Exception as e:
-        print(f"Lỗi swap face: {e}")
-        return result_img  # Trả về ảnh gốc nếu có lỗi
-
-# --- Hàm gộp overlay + swap face ---
-def combined_hair_and_face(background_img, source_img):
-    """
-    Hàm gộp: chạy overlay trước, sau đó swap face
-    """
-    if background_img is None or source_img is None:
-        return None
-    
-    try:
-        # Bước 1: Chạy overlay (ghép tóc)
-        hair_result = overlay_source(background_img, source_img)
-        
-        # Bước 2: Chạy swap face từ background lên kết quả overlay
-        final_result = swap_face_wrapper(background_img, hair_result)
-        
-        return final_result
-    except Exception as e:
-        print(f"Lỗi trong quá trình gộp hair + face: {e}")
-        return None
-
-# --- Xây dựng giao diện Gradio ---
-def build_demo():
-    with gr.Blocks(title="Hair Try-On & Face Swap", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""
-        # 🎯 Hair Try-On & Face Swap Application
-        """)
-        with gr.Row():
-            bg = gr.Image(type="pil", label="Background", height=500)
-            src = gr.Image(type="pil", label="Source", height=500, interactive=False)
-            out = gr.Image(label="Result", height=500, interactive=False)
-
-        folder_path_box = gr.Textbox(label="Folder path", visible=False)
-
-        
-        with gr.Row():
-            src_name_box = gr.Textbox(
-                        label="Wigs Name",
-                        interactive=False,
-                        show_copy_button=True ,     # tuỳ chọn – tiện copy đường dẫn
-                        scale = 1
-                        )
-            gallery = gr.Gallery(
-                            label="Recommend For You",
-                            height=300,
-                            value=[],
-                            type="filepath",
-                            interactive=False,
-                            columns=5,
-                            object_fit="cover",
-                            allow_preview=True,
-                            scale = 8
-                            )
-            with gr.Column(scale=1):
-                combined_btn = gr.Button("🔄✨ Run Hair + Face Swap", variant="primary")
-                btn = gr.Button("🔄 Run Hair Only", variant="secondary")
-                swap_btn = gr.Button("👤 Swap Face Only", variant="secondary")
- 
-
-
-        # Chạy gộp hair + face swap
-        combined_btn.click(fn=combined_hair_and_face, inputs=[bg, src], outputs=[out])
-        
-        # Chạy ghép tóc
-        btn.click(fn=overlay_source, inputs=[bg, src], outputs=[out])
-        
-        # Chạy swap face
-        swap_btn.click(fn=swap_face_wrapper, inputs=[bg, out], outputs=[out])
-        
-        # Khi đổi ảnh background, tự động phân loại và load ảnh gợi ý
-        bg.change(
-            fn=handle_bg_change,
-            inputs=[bg],
-            outputs=[folder_path_box, gallery],
-            show_progress=True
-        )
-        # Khi chọn ảnh trong gallery, cập nhật vào khung Source
-        gallery.select(
-            fn=on_gallery_select,
-            outputs=[src, src_name_box]
-        )
-
-    return demo
-
-if __name__ == "__main__":
-    build_demo().launch()
+
+import gradio as gr
+from overlay import overlay_source
+from detect_face import predict, NUM_CLASSES
+from swapface import swap_face_now
+import os
+from pathlib import Path
+
+BASE_DIR = Path(__file__).parent  # thư mục chứa app.py
+FOLDER = BASE_DIR / "example_wigs"
+
+# --- Hàm load ảnh từ folder ---
+def load_images_from_folder(folder_path: str) -> list[str]:
+    """
+    Trả về list[str] chứa tất cả các hình (jpg, png, gif, bmp) trong folder_path.
+    """
+    supported = {'.jpg', '.jpeg', '.png', '.gif', '.bmp'}
+    if not os.path.isdir(folder_path):
+        print(f"Cảnh báo: '{folder_path}' không phải folder hợp lệ.")
+        return []
+    files = [
+        os.path.join(folder_path, fn)
+        for fn in os.listdir(folder_path)
+        if os.path.splitext(fn)[1].lower() in supported
+    ]
+    if not files:
+        print(f"Không tìm thấy hình trong: {folder_path}")
+    return files
+
+def on_gallery_select(evt: gr.SelectData):
+    """
+    Khi click thumbnail: trả về
+      1) filepath để nạp vào Image Source
+      2) tên file (basename) để hiển thị trong Textbox
+    """
+    val = evt.value
+
+    # --- logic trích filepath y như cũ ---
+    if isinstance(val, dict):
+        img = val.get("image")
+        if isinstance(img, str):
+            filepath = img
+        elif isinstance(img, dict):
+            filepath = img.get("path") or img.get("url")
+        else:
+            filepath = next(
+                (v for v in val.values() if isinstance(v, str) and os.path.isfile(v)),
+                None
+            )
+    elif isinstance(val, str):
+        filepath = val
+    else:
+        raise ValueError(f"Kiểu không hỗ trợ: {type(val)}")
+
+    filename = os.path.basename(filepath) if filepath else ""
+    return filepath, filename
+
+# --- Hàm xác định folder dựa trên phân lớp ---
+def infer_folder(image) -> str:
+    cls = predict(image)["predicted_class"]
+    folder = str(FOLDER / cls)
+    return folder
+
+# --- Hàm gộp: phân loại + load ảnh ---
+def handle_bg_change(image):
+    """
+    Khi thay đổi background:
+    1. Phân loại khuôn mặt
+    2. Load ảnh từ folder tương ứng
+    """
+    if image is None:
+        return "", []
+    
+    try:
+        folder = infer_folder(image)
+        images = load_images_from_folder(folder)
+        return folder, images
+    except Exception as e:
+        print(f"Lỗi xử lý ảnh: {e}")
+        return "", []
+
+# --- Hàm swap face ---
+def swap_face_wrapper(background_img, result_img):
+    """
+    Wrapper function cho swap face giữa background và result image
+    """
+    if background_img is None or result_img is None:
+        return None
+    
+    try:
+        # Swap face từ background vào result image
+        swapped = swap_face_now(background_img, result_img, do_enhance=True)
+        return swapped
+    except Exception as e:
+        print(f"Lỗi swap face: {e}")
+        return result_img  # Trả về ảnh gốc nếu có lỗi
+
+# --- Hàm gộp overlay + swap face ---
+def combined_hair_and_face(background_img, source_img):
+    """
+    Hàm gộp: chạy overlay trước, sau đó swap face
+    """
+    if background_img is None or source_img is None:
+        return None
+    
+    try:
+        # Bước 1: Chạy overlay (ghép tóc)
+        hair_result = overlay_source(background_img, source_img)
+        
+        # Bước 2: Chạy swap face từ background lên kết quả overlay
+        final_result = swap_face_wrapper(background_img, hair_result)
+        
+        return final_result
+    except Exception as e:
+        print(f"Lỗi trong quá trình gộp hair + face: {e}")
+        return None
+
+# --- Xây dựng giao diện Gradio ---
+def build_demo():
+    with gr.Blocks(title="Hair Try-On & Face Swap", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🎯 Hair Try-On & Face Swap Application
+        """)
+        with gr.Row():
+            bg = gr.Image(type="pil", label="Background", height=500)
+            src = gr.Image(type="pil", label="Source", height=500, interactive=False)
+            out = gr.Image(label="Result", height=500, interactive=False)
+
+        folder_path_box = gr.Textbox(label="Folder path", visible=False)
+
+        
+        with gr.Row():
+            src_name_box = gr.Textbox(
+                        label="Wigs Name",
+                        interactive=False,
+                        show_copy_button=True ,     # tuỳ chọn – tiện copy đường dẫn
+                        scale = 1
+                        )
+            gallery = gr.Gallery(
+                            label="Recommend For You",
+                            height=300,
+                            value=[],
+                            type="filepath",
+                            interactive=False,
+                            columns=5,
+                            object_fit="cover",
+                            allow_preview=True,
+                            scale = 8
+                            )
+            with gr.Column(scale=1):
+                combined_btn = gr.Button("🔄✨ Run Hair + Face Swap", variant="primary")
+
+        # Chạy gộp hair + face swap
+        combined_btn.click(fn=combined_hair_and_face, inputs=[bg, src], outputs=[out])
+        
+        
+        # Khi đổi ảnh background, tự động phân loại và load ảnh gợi ý
+        bg.change(
+            fn=handle_bg_change,
+            inputs=[bg],
+            outputs=[folder_path_box, gallery],
+            show_progress=True
+        )
+        # Khi chọn ảnh trong gallery, cập nhật vào khung Source
+        gallery.select(
+            fn=on_gallery_select,
+            outputs=[src, src_name_box]
+        )
+
+    return demo
+
+if __name__ == "__main__":
+    build_demo().launch()