setup.py

"""
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()