Upload 60 files

README.md

@@ -1,6 +1,6 @@
 ---
-title: Ghep Image
-emoji: 📉
+title: AI Wig Try-On System
+emoji: 💇‍♀️
 colorFrom: pink
 colorTo: blue
 sdk: gradio
@@ -9,4 +9,159 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🎭 AI Wig Try-On System
+
+Hệ thống thử tóc giả thông minh sử dụng AI - Tự động phân loại khuôn mặt và gợi ý tóc giả phù hợp.
+
+## ✨ Tính năng chính
+
+- 🤖 **Phân loại khuôn mặt tự động**: AI nhận diện 5 dạng mặt (Heart, Oblong, Oval, Round, Square)
+- 🎯 **Gợi ý tóc giả thông minh**: Tự động hiển thị tóc phù hợp với khuôn mặt
+- 🎨 **Face Swap chất lượng cao**: Sử dụng công nghệ roop để face swapping tự nhiên
+- 🌈 **Color Transfer tự động**: Điều chỉnh màu da từ vùng trán để hài hòa
+- 🖼️ **Giao diện thân thiện**: Web interface đơn giản với Gradio
+
+## 🏗️ Kiến trúc hệ thống
+
+```
+📁 Project Structure
+├── app.py                 # Main Gradio interface
+├── detect_face.py         # Face shape classification (EfficientNet-B4)
+├── overlay.py             # Hair overlay processing
+├── swapface.py            # Face swapping using roop
+├── apply_color_transfer.py # Color matching from forehead region
+├── baldhead.py            # Hair removal using GAN
+├── segmentation.py        # Hair/face segmentation
+├── setup.py               # Project setup checker
+└── example_wigs/          # Wig samples organized by face shape
+    ├── Heart/
+    ├── Oblong/
+    ├── Oval/
+    ├── Round/
+    └── Square/
+```
+
+## 🚀 Cài đặt và chạy
+
+### 1. Cài đặt dependencies
+```bash
+pip install -r requirements.txt
+```
+
+### 2. Kiểm tra setup
+```bash
+python setup.py
+```
+
+### 3. Chạy ứng dụng
+```bash
+python app.py
+```
+
+## 🎯 Cách sử dụng
+
+1. **Upload ảnh khuôn mặt** vào khung "Background"
+2. **Tự động phân loại**: Hệ thống sẽ phân tích khuôn mặt và hiển thị tóc giả gợi ý
+3. **Chọn tóc giả**: Click vào gallery để chọn kiểu tóc mong muốn
+4. **Xử lý**: Nhấn nút "🔄 Run" để thực hiện face swap
+5. **Kết quả**: Xem kết quả cuối cùng với màu da được điều chỉnh tự nhiên
+
+## 🤖 AI Models được sử dụng
+
+- **Face Classification**: EfficientNet-B4 custom trained
+- **Hair Segmentation**: SegFormer from Hugging Face
+- **Face Detection**: MediaPipe + RetinaFace
+- **Hair Removal**: Custom GAN model
+- **Face Swapping**: roop library
+- **Color Transfer**: LAB color space histogram matching
+
+## 🎨 Color Transfer Technology
+
+Hệ thống tự động:
+- 🎯 **Detect vùng trán** bằng MediaPipe
+- 🌈 **Lấy màu reference** từ vùng trán (30% phần trên khuôn mặt)
+- 🔄 **Apply color matching** với opacity 90%
+- ✨ **Kết quả tự nhiên** với màu da hài hòa
+
+## 📊 Supported Face Shapes
+
+| Face Shape | Description | Wig Recommendations |
+|------------|-------------|-------------------|
+| **Heart** | Trán rộng, cằm nhọn | Tóc bob, lob, layers |
+| **Oblong** | Mặt dài, tỷ lệ cao | Tóc có volume, fringe |
+| **Oval** | Cân đối, lý tưởng | Phù hợp mọi kiểu tóc |
+| **Round** | Mặt tròn, má đầy | Tóc dài, layers |
+| **Square** | Góc cạnh, hàm vuông | Tóc xoăn, soft waves |
+
+## 🛠️ Technical Requirements
+
+- Python 3.8+
+- CUDA-compatible GPU (recommended)
+- RAM: 8GB+
+- Storage: 5GB+ for models
+
+## 📋 Key Dependencies
+
+- `gradio` - Web interface
+- `torch` + `torchvision` - Deep learning
+- `transformers` - Hugging Face models  
+- `mediapipe` - Face detection
+- `opencv-python` - Image processing
+- `insightface` - Face analysis
+- `tensorflow` - GAN models
+
+## 🏃‍♂️ Quick Start
+
+```bash
+# Clone and setup
+git clone <repo-url>
+cd be_rejection
+
+# Install requirements
+pip install -r requirements.txt
+
+# Check setup
+python setup.py
+
+# Run app
+python app.py
+```
+
+## 🎭 Adding New Wigs
+
+Để thêm tóc giả mới:
+1. Đặt ảnh tóc giả vào folder tương ứng trong `example_wigs/`
+2. Format: PNG/JPG với background trong suốt (khuyến nghị)
+3. Kích thước: Tối thiểu 512x512px
+4. Chất lượng: Ảnh rõ nét, góc chụp thẳng
+
+## 🔧 Troubleshooting
+
+**Lỗi model không load:**
+```bash
+# Kiểm tra kết nối internet và Hugging Face Hub
+python -c "from transformers import pipeline; print('HF Hub OK')"
+```
+
+**Lỗi CUDA:**
+```bash
+# Chuyển sang CPU mode
+export CUDA_VISIBLE_DEVICES=""
+```
+
+**Lỗi memory:**
+- Giảm batch size
+- Sử dụng GPU có RAM lớn hơn
+- Giảm kích thước ảnh input
+
+## 📝 License
+
+Dự ��n này được phát triển cho mục đích giáo dục và nghiên cứu.
+
+## 🤝 Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+
+---
+
+*Powered by AI • Made with ❤️ in Vietnam*

app.py

@@ -1,182 +1,182 @@
-
-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
-from PIL import Image
-
-
-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)}")
-
- # Lấy tên file không có phần mở rộng
-    filename = os.path.splitext(os.path.basename(filepath))[0] 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 "", []
-
-# --- Pipeline đầy đủ: ghép tóc + swap face cuối ---
-def complete_pipeline(background: Image.Image, source: Image.Image):
-    if background is None or source is None:
-        return None
-
-    try:
-        # Import color transfer function
-        from apply_color_transfer import apply_color_transfer_to_output
-        
-        # Bước 1: overlay tóc
-        intermediate_result = overlay_source(background, source)
-        if intermediate_result is None:
-            print("Overlay failed.")
-            return None
-
-        # Bước 2: swap face - truyền trực tiếp PIL Images
-        final_result = swap_face_now(
-            background.convert("RGB"),
-            intermediate_result.convert("RGB")
-        )
-
-        if final_result is None:
-            print("Swap failed.")
-            result_img = intermediate_result  # fallback
-        else:
-            print("Swap success.")
-            result_img = final_result  # đây đã là ảnh PIL.Image
-
-        # Bước 3: Apply color transfer để match màu da
-        result_with_color_transfer = apply_color_transfer_to_output(
-            swapped_img=result_img,
-            original_bg=background,
-            opacity=0.9
-        )
-
-        return result_with_color_transfer
-
-    except Exception as e:
-        print(f"Lỗi trong complete_pipeline: {e}")
-        return None
-        
-# --- Xây dựng giao diện Gradio ---
-def build_demo():
-    with gr.Blocks(title="Xử lý hai hình ảnh", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("Upload Background & Source, click **Run** to try on wigs.")
-
-        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
-                            )
-            btn = gr.Button("🔄 Run", variant="primary",scale = 1)  
- 
-
-
-        # Chạy pipeline đầy đủ: ghép tóc + swap face cuối
-        btn.click(fn=complete_pipeline, 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
-        )
-        # Nút tải lại ảnh thủ công (backup)
-        # 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
+from PIL import Image
+
+
+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)}")
+
+ # Lấy tên file không có phần mở rộng
+    filename = os.path.splitext(os.path.basename(filepath))[0] 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 "", []
+
+# --- Pipeline đầy đủ: ghép tóc + swap face cuối ---
+def complete_pipeline(background: Image.Image, source: Image.Image):
+    if background is None or source is None:
+        return None
+
+    try:
+        # Import color transfer function
+        from apply_color_transfer import apply_color_transfer_to_output
+        
+        # Bước 1: overlay tóc
+        intermediate_result = overlay_source(background, source)
+        if intermediate_result is None:
+            print("Overlay failed.")
+            return None
+
+        # Bước 2: swap face - truyền trực tiếp PIL Images
+        final_result = swap_face_now(
+            background.convert("RGB"),
+            intermediate_result.convert("RGB")
+        )
+
+        if final_result is None:
+            print("Swap failed.")
+            result_img = intermediate_result  # fallback
+        else:
+            print("Swap success.")
+            result_img = final_result  # đây đã là ảnh PIL.Image
+
+        # Bước 3: Apply color transfer để match màu da
+        result_with_color_transfer = apply_color_transfer_to_output(
+            swapped_img=result_img,
+            original_bg=background,
+            opacity=0.9
+        )
+
+        return result_with_color_transfer
+
+    except Exception as e:
+        print(f"Lỗi trong complete_pipeline: {e}")
+        return None
+        
+# --- Xây dựng giao diện Gradio ---
+def build_demo():
+    with gr.Blocks(title="Xử lý hai hình ảnh", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("Upload Background & Source, click **Run** to try on wigs.")
+
+        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
+                            )
+            btn = gr.Button("🔄 Run", variant="primary",scale = 1)  
+ 
+
+
+        # Chạy pipeline đầy đủ: ghép tóc + swap face cuối
+        btn.click(fn=complete_pipeline, 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
+        )
+        # Nút tải lại ảnh thủ công (backup)
+        # 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()

apply_color_transfer.py

@@ -1,16 +1,79 @@
 import cv2
 import numpy as np
 from PIL import Image
+import mediapipe as mp
 
+# MediaPipe Face Detection
+mp_face_detection = mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5)
 
-def apply_color_transfer_to_output(swapped_img, original_bg, opacity=0.7):
+
+def get_forehead_region(image):
+    """
+    Detect forehead region using MediaPipe face detection.
+    Returns the highest point of the face (forehead area).
+    
+    Args:
+        image: PIL Image or numpy array
+    
+    Returns:
+        tuple: (x, y, width, height) of forehead region or None if no face detected
+    """
+    try:
+        # Convert to numpy if PIL Image
+        if isinstance(image, Image.Image):
+            img_array = np.array(image.convert('RGB'))
+        else:
+            img_array = image
+        
+        # Detect face
+        results = mp_face_detection.process(img_array)
+        
+        if not results.detections:
+            return None
+        
+        # Get first detected face
+        detection = results.detections[0]
+        bbox = detection.location_data.relative_bounding_box
+        
+        h, w = img_array.shape[:2]
+        
+        # Convert relative coordinates to absolute
+        x = int(bbox.xmin * w)
+        y = int(bbox.ymin * h)
+        width = int(bbox.width * w)
+        height = int(bbox.height * h)
+        
+        # Define forehead region (upper 30% of face)
+        forehead_height = int(height * 0.3)
+        forehead_y = max(0, y - int(forehead_height * 0.2))  # Slightly above face bbox
+        
+        # Ensure forehead region is within image bounds
+        forehead_x = max(0, x + int(width * 0.2))  # Start from 20% into face width
+        forehead_width = int(width * 0.6)  # Use 60% of face width
+        
+        # Clamp to image boundaries
+        forehead_x = max(0, min(forehead_x, w - 1))
+        forehead_y = max(0, min(forehead_y, h - 1))
+        forehead_width = max(1, min(forehead_width, w - forehead_x))
+        forehead_height = max(1, min(forehead_height, h - forehead_y))
+        
+        return (forehead_x, forehead_y, forehead_width, forehead_height)
+    
+    except Exception as e:
+        print(f"Forehead detection failed: {e}")
+        return None
+
+
+def apply_color_transfer_to_output(swapped_img, original_bg, opacity=0.7, color_region=None):
     """
     Apply color transfer to make the swapped face match the original background lighting.
+    Automatically uses forehead region as color reference by default.
     
     Args:
         swapped_img: PIL Image - The image after face swapping
         original_bg: PIL Image - The original background image for color reference
         opacity: float - Opacity of color transfer (0.0 = no effect, 1.0 = full effect)
+        color_region: tuple or None - (x, y, width, height) for specific color sampling region
     
     Returns:
         PIL Image - The color-corrected result
@@ -23,8 +86,26 @@ def apply_color_transfer_to_output(swapped_img, original_bg, opacity=0.7):
         swapped_np = np.array(swapped_img.convert('RGB'))
         original_np = np.array(original_bg.convert('RGB'))
         
-        # Apply color transfer using histogram matching
-        color_corrected = histogram_matching(swapped_np, original_np)
+        # If no specific region provided, automatically detect forehead
+        if color_region is None:
+            color_region = get_forehead_region(original_bg)
+            if color_region:
+                print(f"Auto-detected forehead region: {color_region}")
+        
+        # If specific region is provided or auto-detected, use it for color reference
+        if color_region is not None:
+            x, y, w, h = color_region
+            # Extract color reference from specific region
+            reference_region = original_np[y:y+h, x:x+w]
+            if reference_region.size > 0:
+                # Apply color transfer using the selected region
+                color_corrected = histogram_matching_with_region(swapped_np, original_np, reference_region)
+            else:
+                # Fallback to full image if region is invalid
+                color_corrected = histogram_matching(swapped_np, original_np)
+        else:
+            # Apply color transfer using histogram matching on full image
+            color_corrected = histogram_matching(swapped_np, original_np)
         
         # Blend original and color-corrected image based on opacity
         result = blend_images(swapped_np, color_corrected, opacity)
@@ -90,6 +171,40 @@ def histogram_matching(source, reference):
     return result
 
 
+def histogram_matching_with_region(source, reference_full, reference_region):
+    """
+    Perform histogram matching using a specific region as color reference.
+    
+    Args:
+        source: numpy array - Source image to be color corrected
+        reference_full: numpy array - Full reference image (not used directly)
+        reference_region: numpy array - Specific region to extract color from
+    
+    Returns:
+        numpy array - Color corrected image
+    """
+    # Convert to LAB color space
+    source_lab = cv2.cvtColor(source, cv2.COLOR_RGB2LAB)
+    region_lab = cv2.cvtColor(reference_region, cv2.COLOR_RGB2LAB)
+    
+    # Split channels
+    source_l, source_a, source_b = cv2.split(source_lab)
+    region_l, region_a, region_b = cv2.split(region_lab)
+    
+    # Apply histogram matching using the region's histogram
+    matched_l = match_histograms(source_l, region_l.flatten().reshape(-1, 1).squeeze())
+    matched_a = match_histograms(source_a, region_a.flatten().reshape(-1, 1).squeeze())
+    matched_b = match_histograms(source_b, region_b.flatten().reshape(-1, 1).squeeze())
+    
+    # Merge channels back
+    matched_lab = cv2.merge([matched_l, matched_a, matched_b])
+    
+    # Convert back to RGB
+    result = cv2.cvtColor(matched_lab, cv2.COLOR_LAB2RGB)
+    
+    return result
+
+
 def match_histograms(source, reference):
     """
     Match histogram of source image to reference image.

baldhead.py

@@ -213,17 +213,36 @@ def build_generator():
 
 HF_REPO_ID = "VanNguyen1214/baldhead"
 HF_FILENAME = "model_G_5_170.hdf5"
-HF_TOKEN = os.environ["HUGGINGFACEHUB_API_TOKEN"]
 
 def load_generator_from_hub():
     """
     Download the .hdf5 weights from HF Hub into cache,
     rebuild the generator, then load weights.
     """
-    local_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_FILENAME,token=HF_TOKEN)
-    gen = build_generator()
-    gen.load_weights(local_path)
-    return gen
+    try:
+        # Try with token first if available
+        token = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+        local_path = hf_hub_download(
+            repo_id=HF_REPO_ID, 
+            filename=HF_FILENAME,
+            token=token,
+            local_dir="models/baldhead"
+        )
+        gen = build_generator()
+        gen.load_weights(local_path)
+        return gen
+    except Exception as e:
+        print(f"[WARNING] Could not download from HF Hub: {e}")
+        # Try to load from local cache if exists
+        local_cache_path = f"models/baldhead/{HF_FILENAME}"
+        if os.path.exists(local_cache_path):
+            print(f"[INFO] Loading from local cache: {local_cache_path}")
+            gen = build_generator()
+            gen.load_weights(local_cache_path)
+            return gen
+        else:
+            print(f"[ERROR] No local cache found at {local_cache_path}")
+            return None
 
 # Load once at startup
 try:

overlay.py

@@ -1,93 +1,93 @@
-import numpy as np
-from PIL import Image
-import mediapipe as mp
-
-from baldhead import inference       # cạo tóc background
-from segmentation import extract_hair_face_full_forehead
-from swapface import swap_face_now 
-from apply_color_transfer import apply_color_transfer_to_output
-
-# MediaPipe Face Detection
-mp_fd = mp.solutions.face_detection.FaceDetection(model_selection=1,
-                                                  min_detection_confidence=0.5)
-
-def get_face_bbox(img: Image.Image) -> tuple[int,int,int,int] | None:
-    arr = np.array(img.convert("RGB"))
-    res = mp_fd.process(arr)
-    if not res.detections:
-        return None
-    d = res.detections[0].location_data.relative_bounding_box
-    h, w = arr.shape[:2]
-    x1 = int(d.xmin * w)
-    y1 = int(d.ymin * h)
-    x2 = x1 + int(d.width * w)
-    y2 = y1 + int(d.height * h)
-    return x1, y1, x2, y2
-
-def compute_scale(w_bg, h_bg, w_src, h_src) -> float:
-    return ((w_bg / w_src) + (h_bg / h_src)) / 2
-
-def compute_offset(bbox_bg, bbox_src, scale) -> tuple[int,int]:
-    x1, y1, x2, y2 = bbox_bg
-    bg_cx = x1 + (x2 - x1)//2
-    bg_cy = y1 + (y2 - y1)//2
-    sx1, sy1, sx2, sy2 = bbox_src
-    src_cx = int((sx1 + (sx2 - sx1)//2) * scale)
-    src_cy = int((sy1 + (sy2 - sy1)//2) * scale)
-    return bg_cx - src_cx, bg_cy - src_cy
-
-def paste_with_alpha(bg: np.ndarray, src: np.ndarray, offset: tuple[int,int]) -> Image.Image:
-    res = bg.copy()
-    x, y = offset
-    h, w = src.shape[:2]
-    x1, y1 = max(x,0), max(y,0)
-    x2 = min(x+w, bg.shape[1])
-    y2 = min(y+h, bg.shape[0])
-    if x1>=x2 or y1>=y2:
-        return Image.fromarray(res)
-    cs = src[y1-y:y2-y, x1-x:x2-x]
-    cd = res[y1:y2, x1:x2]
-    mask = cs[...,3] > 0
-    if cd.shape[2] == 3:
-        cd[mask] = cs[mask][...,:3]
-    else:
-        cd[mask] = cs[mask]
-    res[y1:y2, x1:x2] = cd
-    return Image.fromarray(res)
-
-def overlay_source(background: Image.Image, source: Image.Image):
-    # 1) detect bboxes
-    bbox_bg  = get_face_bbox(background)
-    bbox_src = get_face_bbox(source)
-    if bbox_bg is None:
-        return None, "❌ No face in background."
-    if bbox_src is None:
-        return None, "❌ No face in source."
-
-    # 2) compute scale & resize source
-    w_bg, h_bg = bbox_bg[2]-bbox_bg[0], bbox_bg[3]-bbox_bg[1]
-    w_src, h_src = bbox_src[2]-bbox_src[0], bbox_src[3]-bbox_src[1]
-    scale = compute_scale(w_bg, h_bg, w_src, h_src)
-    src_scaled = source.resize(
-        (int(source.width*scale), int(source.height*scale)),
-        Image.Resampling.LANCZOS
-    )
-
-    # 3) compute offset
-    offset = compute_offset(bbox_bg, bbox_src, scale)
-
-    # 4) baldhead background
-    bg_bald = inference(background)
-
-    # 5) extract hair-only from source
-    full_head = extract_hair_face_full_forehead(src_scaled)
-
-    # 6) paste onto bald background
-    src_change = paste_with_alpha(
-        np.array(bg_bald.convert("RGBA")),
-        np.array(full_head),
-        offset
-    )
-    
-    # Return the intermediate result for face swapping in complete_pipeline
-    return src_change.convert("RGB")
+import numpy as np
+from PIL import Image
+import mediapipe as mp
+
+from baldhead import inference       # cạo tóc background
+from segmentation import extract_hair_face_full_forehead
+from swapface import swap_face_now 
+from apply_color_transfer import apply_color_transfer_to_output
+
+# MediaPipe Face Detection
+mp_fd = mp.solutions.face_detection.FaceDetection(model_selection=1,
+                                                  min_detection_confidence=0.5)
+
+def get_face_bbox(img: Image.Image) -> tuple[int,int,int,int] | None:
+    arr = np.array(img.convert("RGB"))
+    res = mp_fd.process(arr)
+    if not res.detections:
+        return None
+    d = res.detections[0].location_data.relative_bounding_box
+    h, w = arr.shape[:2]
+    x1 = int(d.xmin * w)
+    y1 = int(d.ymin * h)
+    x2 = x1 + int(d.width * w)
+    y2 = y1 + int(d.height * h)
+    return x1, y1, x2, y2
+
+def compute_scale(w_bg, h_bg, w_src, h_src) -> float:
+    return ((w_bg / w_src) + (h_bg / h_src)) / 2
+
+def compute_offset(bbox_bg, bbox_src, scale) -> tuple[int,int]:
+    x1, y1, x2, y2 = bbox_bg
+    bg_cx = x1 + (x2 - x1)//2
+    bg_cy = y1 + (y2 - y1)//2
+    sx1, sy1, sx2, sy2 = bbox_src
+    src_cx = int((sx1 + (sx2 - sx1)//2) * scale)
+    src_cy = int((sy1 + (sy2 - sy1)//2) * scale)
+    return bg_cx - src_cx, bg_cy - src_cy
+
+def paste_with_alpha(bg: np.ndarray, src: np.ndarray, offset: tuple[int,int]) -> Image.Image:
+    res = bg.copy()
+    x, y = offset
+    h, w = src.shape[:2]
+    x1, y1 = max(x,0), max(y,0)
+    x2 = min(x+w, bg.shape[1])
+    y2 = min(y+h, bg.shape[0])
+    if x1>=x2 or y1>=y2:
+        return Image.fromarray(res)
+    cs = src[y1-y:y2-y, x1-x:x2-x]
+    cd = res[y1:y2, x1:x2]
+    mask = cs[...,3] > 0
+    if cd.shape[2] == 3:
+        cd[mask] = cs[mask][...,:3]
+    else:
+        cd[mask] = cs[mask]
+    res[y1:y2, x1:x2] = cd
+    return Image.fromarray(res)
+
+def overlay_source(background: Image.Image, source: Image.Image):
+    # 1) detect bboxes
+    bbox_bg  = get_face_bbox(background)
+    bbox_src = get_face_bbox(source)
+    if bbox_bg is None:
+        return None, "❌ No face in background."
+    if bbox_src is None:
+        return None, "❌ No face in source."
+
+    # 2) compute scale & resize source
+    w_bg, h_bg = bbox_bg[2]-bbox_bg[0], bbox_bg[3]-bbox_bg[1]
+    w_src, h_src = bbox_src[2]-bbox_src[0], bbox_src[3]-bbox_src[1]
+    scale = compute_scale(w_bg, h_bg, w_src, h_src)
+    src_scaled = source.resize(
+        (int(source.width*scale), int(source.height*scale)),
+        Image.Resampling.LANCZOS
+    )
+
+    # 3) compute offset
+    offset = compute_offset(bbox_bg, bbox_src, scale)
+
+    # 4) baldhead background
+    bg_bald = inference(background)
+
+    # 5) extract hair-only from source
+    full_head = extract_hair_face_full_forehead(src_scaled)
+
+    # 6) paste onto bald background
+    src_change = paste_with_alpha(
+        np.array(bg_bald.convert("RGBA")),
+        np.array(full_head),
+        offset
+    )
+    
+    # Return the intermediate result for face swapping in complete_pipeline
+    return src_change.convert("RGB")

requirements.txt

@@ -1,35 +1,35 @@
---extra-index-url https://download.pytorch.org/whl/cu118 # Dòng này có vẻ là comment hoặc cấu hình cho pip, không phải là một gói
-# spaces # Dòng này không rõ ràng là một gói, có thể là ghi chú. Nếu không phải gói, hãy xóa đi.
-huggingface_hub>=0.20.3
-numpy==1.23.5
-transformers==4.30.0
-opencv-python-headless==4.7.0.72
-onnx==1.14.0
-insightface==0.7.3
-psutil==5.9.5
-tk==0.1.0 # Lưu ý: tk thường được bao gồm trong bản cài đặt Python chuẩn, không phải lúc nào cũng cần cài qua pip.
-customtkinter==5.1.3
-pillow==9.5.0
-torch==2.0.1+cu118; sys_platform != 'darwin'
-torch==2.0.1; sys_platform == 'darwin'
-torchvision==0.15.2+cu118; sys_platform != 'darwin'
-torchvision==0.15.2; sys_platform == 'darwin'
-# onnxruntime==1.15.0; # Bỏ comment cho dòng này nếu bạn muốn cố định phiên bản cho mọi OS
-# sys_platform == 'darwin' and platform_machine != 'arm64' # Comment
-onnxruntime-silicon==1.13.1; sys_platform == 'darwin' and platform_machine == 'arm64'
-onnxruntime-gpu==1.15.0; sys_platform != 'darwin' # Nên giữ lại dòng này cho non-darwin GPU
-onnxruntime==1.15.0; sys_platform == 'darwin' and platform_machine != 'arm64' # Thêm lại dòng onnxruntime cho Mac Intel
-tensorflow==2.12.0
-# sys_platform != 'darwin' # Comment
-opennsfw2==0.10.2
-# protobuf==4.23.2 # Thay thế dòng này
-protobuf==4.25.3  # *** THAY ĐỔI QUAN TRỌNG ***
-tqdm==4.65.0
-gfpgan==1.3.8
-# torch # Dòng này không cần thiết vì torch đã được định nghĩa ở trên với phiên bản cụ thể.
-
-# Thêm các thư viện mới cần thiết cho app.py đã cập nhật
-scikit-image>=0.19 # Hoặc một phiên bản cụ thể hơn nếu bạn muốn, ví dụ: scikit-image==0.19.3
-mediapipe==0.10.14 # *** THÊM MỚI HOẶC CẬP NHẬT *** (Phiên bản này yêu cầu protobuf >=4.25.3)
-git+https://github.com/keras-team/keras-contrib.git
+--extra-index-url https://download.pytorch.org/whl/cu118 # Dòng này có vẻ là comment hoặc cấu hình cho pip, không phải là một gói
+# spaces # Dòng này không rõ ràng là một gói, có thể là ghi chú. Nếu không phải gói, hãy xóa đi.
+huggingface_hub>=0.20.3
+numpy==1.23.5
+transformers==4.30.0
+opencv-python-headless==4.7.0.72
+onnx==1.14.0
+insightface==0.7.3
+psutil==5.9.5
+tk==0.1.0 # Lưu ý: tk thường được bao gồm trong bản cài đặt Python chuẩn, không phải lúc nào cũng cần cài qua pip.
+customtkinter==5.1.3
+pillow==9.5.0
+torch==2.0.1+cu118; sys_platform != 'darwin'
+torch==2.0.1; sys_platform == 'darwin'
+torchvision==0.15.2+cu118; sys_platform != 'darwin'
+torchvision==0.15.2; sys_platform == 'darwin'
+# onnxruntime==1.15.0; # Bỏ comment cho dòng này nếu bạn muốn cố định phiên bản cho mọi OS
+# sys_platform == 'darwin' and platform_machine != 'arm64' # Comment
+onnxruntime-silicon==1.13.1; sys_platform == 'darwin' and platform_machine == 'arm64'
+onnxruntime-gpu==1.15.0; sys_platform != 'darwin' # Nên giữ lại dòng này cho non-darwin GPU
+onnxruntime==1.15.0; sys_platform == 'darwin' and platform_machine != 'arm64' # Thêm lại dòng onnxruntime cho Mac Intel
+tensorflow==2.12.0
+# sys_platform != 'darwin' # Comment
+opennsfw2==0.10.2
+# protobuf==4.23.2 # Thay thế dòng này
+protobuf==4.25.3  # *** THAY ĐỔI QUAN TRỌNG ***
+tqdm==4.65.0
+gfpgan==1.3.8
+# torch # Dòng này không cần thiết vì torch đã được định nghĩa ở trên với phiên bản cụ thể.
+
+# Thêm các thư viện mới cần thiết cho app.py đã cập nhật
+scikit-image>=0.19 # Hoặc một phiên bản cụ thể hơn nếu bạn muốn, ví dụ: scikit-image==0.19.3
+mediapipe==0.10.14 # *** THÊM MỚI HOẶC CẬP NHẬT *** (Phiên bản này yêu cầu protobuf >=4.25.3)
+git+https://github.com/keras-team/keras-contrib.git
 retina-face==0.0.13

setup.py

@@ -0,0 +1,131 @@
+"""
+Setup script for FaceSwap AI Wig Try-On Project
+Ensures all models and dependencies are properly configured
+"""
+
+import os
+import sys
+from pathlib import Path
+
+def check_models():
+    """Check if all required models are accessible"""
+    print("🔍 Checking AI models...")
+    
+    models_status = {}
+    
+    # Check face detection model
+    try:
+        from detect_face import predict, _MODEL
+        if _MODEL is not None:
+            models_status['face_detection'] = "✅ Ready"
+        else:
+            models_status['face_detection'] = "❌ Failed to load"
+    except Exception as e:
+        models_status['face_detection'] = f"❌ Error: {e}"
+    
+    # Check segmentation model
+    try:
+        from segmentation import processor, model
+        models_status['segmentation'] = "✅ Ready"
+    except Exception as e:
+        models_status['segmentation'] = f"❌ Error: {e}"
+    
+    # Check baldhead model
+    try:
+        from baldhead import GENERATOR
+        if GENERATOR is not None:
+            models_status['baldhead'] = "✅ Ready"
+        else:
+            models_status['baldhead'] = "⚠️ Model not loaded (may download on first use)"
+    except Exception as e:
+        models_status['baldhead'] = f"❌ Error: {e}"
+    
+    # Check roop face swap
+    try:
+        import roop.core
+        models_status['face_swap'] = "✅ Ready"
+    except Exception as e:
+        models_status['face_swap'] = f"❌ Error: {e}"
+    
+    return models_status
+
+def check_directories():
+    """Ensure required directories exist"""
+    print("📁 Checking directories...")
+    
+    required_dirs = [
+        "example_wigs/Heart",
+        "example_wigs/Oblong", 
+        "example_wigs/Oval",
+        "example_wigs/Round",
+        "example_wigs/Square",
+        "models",
+        "temp"
+    ]
+    
+    for dir_path in required_dirs:
+        Path(dir_path).mkdir(parents=True, exist_ok=True)
+        print(f"   ✅ {dir_path}")
+
+def check_wig_samples():
+    """Check if wig samples exist"""
+    print("🎭 Checking wig samples...")
+    
+    wig_count = {}
+    for shape in ["Heart", "Oblong", "Oval", "Round", "Square"]:
+        shape_dir = Path(f"example_wigs/{shape}")
+        if shape_dir.exists():
+            wigs = list(shape_dir.glob("*.png")) + list(shape_dir.glob("*.jpg"))
+            wig_count[shape] = len(wigs)
+            print(f"   {shape}: {len(wigs)} wigs")
+        else:
+            wig_count[shape] = 0
+            print(f"   {shape}: 0 wigs ❌")
+    
+    return wig_count
+
+def main():
+    """Main setup check"""
+    print("🚀 FaceSwap AI Wig Try-On - Setup Check")
+    print("=" * 50)
+    
+    # Check directories
+    check_directories()
+    print()
+    
+    # Check wig samples  
+    wig_count = check_wig_samples()
+    print()
+    
+    # Check models
+    models_status = check_models()
+    print()
+    
+    # Summary
+    print("📊 SETUP SUMMARY:")
+    print("-" * 30)
+    
+    print("🤖 AI Models:")
+    for model_name, status in models_status.items():
+        print(f"   {model_name}: {status}")
+    
+    print(f"\n🎭 Wig Samples: {sum(wig_count.values())} total")
+    for shape, count in wig_count.items():
+        print(f"   {shape}: {count}")
+    
+    # Check if ready to run
+    critical_models = ['face_detection', 'segmentation', 'face_swap']
+    ready = all("✅" in models_status.get(model, "") for model in critical_models)
+    
+    if ready and sum(wig_count.values()) > 0:
+        print("\n🎉 Project is ready to run!")
+        print("   Run: python app.py")
+    else:
+        print("\n⚠️  Some components need attention:")
+        if not ready:
+            print("   - AI models need fixing")
+        if sum(wig_count.values()) == 0:
+            print("   - Add wig samples to example_wigs/ folders")
+
+if __name__ == "__main__":
+    main()

swapface.py

@@ -1,53 +1,53 @@
-import os
-from PIL import Image
-import numpy as np
-import roop.globals
-from roop.core import start, decode_execution_providers, suggest_max_memory, suggest_execution_threads
-from roop.processors.frame.core import get_frame_processors_modules
-
-def to_pil(img):
-    """Convert np.ndarray to PIL.Image if needed."""
-    if isinstance(img, np.ndarray):
-        return Image.fromarray(img)
-    return img
-
-def swap_face_now(source_img, target_img, do_enhance=True): 
-    TEMP_DIR = "temp/"
-    os.makedirs(TEMP_DIR, exist_ok=True)
-
-    # ✅ Ép kiểu an toàn
-    source_img = to_pil(source_img)
-    target_img = to_pil(target_img)
-
-    source_path = os.path.join(TEMP_DIR, "input.jpg")
-    target_path = os.path.join(TEMP_DIR, "target.jpg")
-    output_path = os.path.join(TEMP_DIR, "output.jpg")
-
-    source_img.save(source_path)
-    target_img.save(target_path)
-
-    # Roop config
-    roop.globals.source_path = source_path
-    roop.globals.target_path = target_path
-    roop.globals.output_path = output_path
-    roop.globals.frame_processors = ["face_swapper", "face_enhancer"] if do_enhance else ["face_swapper"]
-    roop.globals.headless = True
-    roop.globals.many_faces = False
-    roop.globals.max_memory = suggest_max_memory()
-    roop.globals.execution_providers = decode_execution_providers(["cuda"])
-    roop.globals.execution_threads = suggest_execution_threads()
-    roop.globals.keep_fps = False
-    roop.globals.keep_audio = False
-    roop.globals.keep_frames = False
-    roop.globals.video_encoder = None
-
-    for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
-        if not frame_processor.pre_check():
-            return None
-
-    start()
-
-    if not os.path.exists(output_path):
-        return None
-
+import os
+from PIL import Image
+import numpy as np
+import roop.globals
+from roop.core import start, decode_execution_providers, suggest_max_memory, suggest_execution_threads
+from roop.processors.frame.core import get_frame_processors_modules
+
+def to_pil(img):
+    """Convert np.ndarray to PIL.Image if needed."""
+    if isinstance(img, np.ndarray):
+        return Image.fromarray(img)
+    return img
+
+def swap_face_now(source_img, target_img, do_enhance=True): 
+    TEMP_DIR = "temp/"
+    os.makedirs(TEMP_DIR, exist_ok=True)
+
+    # ✅ Ép kiểu an toàn
+    source_img = to_pil(source_img)
+    target_img = to_pil(target_img)
+
+    source_path = os.path.join(TEMP_DIR, "input.jpg")
+    target_path = os.path.join(TEMP_DIR, "target.jpg")
+    output_path = os.path.join(TEMP_DIR, "output.jpg")
+
+    source_img.save(source_path)
+    target_img.save(target_path)
+
+    # Roop config
+    roop.globals.source_path = source_path
+    roop.globals.target_path = target_path
+    roop.globals.output_path = output_path
+    roop.globals.frame_processors = ["face_swapper", "face_enhancer"] if do_enhance else ["face_swapper"]
+    roop.globals.headless = True
+    roop.globals.many_faces = False
+    roop.globals.max_memory = suggest_max_memory()
+    roop.globals.execution_providers = decode_execution_providers(["cuda"])
+    roop.globals.execution_threads = suggest_execution_threads()
+    roop.globals.keep_fps = False
+    roop.globals.keep_audio = False
+    roop.globals.keep_frames = False
+    roop.globals.video_encoder = None
+
+    for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
+        if not frame_processor.pre_check():
+            return None
+
+    start()
+
+    if not os.path.exists(output_path):
+        return None
+
     return Image.open(output_path)