Spaces:
Runtime error
Runtime error
File size: 3,930 Bytes
07f3e72 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
"""
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()
|