README.md

metadata

title: AI Wig Try-On System
emoji: 💇‍♀️
colorFrom: pink
colorTo: blue
sdk: gradio
sdk_version: 5.35.0
app_file: app.py
pinned: false

🎭 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
• 🖼️ 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
├── 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

pip install -r requirements.txt

2. Kiểm tra setup

python setup.py

3. Chạy ứng dụng

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 workflow mới
5. Kết quả: Xem kết quả cuối cùng với tóc giả được ghép tự nhiên

🔄 Workflow Mới

1. Scale Image Source: Điều chỉnh kích thước ảnh source sao cho khuôn mặt bằng với khuôn mặt background (sử dụng MediaPipe)
2. Bald Head: Thực hiện baldhead image background
3. Face Overlay: Lấy khuôn mặt từ background(baldhead) có trán và overlay lên khuôn mặt source đã được scale
4. Hair Extraction: Trích xuất tóc + mặt + trán từ kết quả bước 3 và overlay lên background(baldhead) gốc và làm mịn vùng overlay lưu vào variable temp sau đó get_hair từ source overlay lên variable temp với vị trí trùng với vị trí tóc của temp_result
5. Face Swapping: Thực hiện face swap giữa background(baldhead) và kết quả tạm để có kết quả cuối cùng

🤖 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

📊 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

# 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:

# 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:

# 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