setup_models.py

"""
Model Setup Script for KYC POC

This script downloads and sets up the required ML models:
1. AuraFace - Face recognition model from HuggingFace
2. Silent-Face-Anti-Spoofing - Liveness detection models from GitHub

Run this script before starting the application:
    python setup_models.py
"""

import os
import sys
import shutil
import subprocess
from pathlib import Path


def setup_auraface():
    """Download AuraFace model from HuggingFace."""
    print("=" * 50)
    print("Setting up AuraFace model...")
    print("=" * 50)

    try:
        from huggingface_hub import snapshot_download

        # InsightFace expects models at {root}/models/{name}/
        # So we download to insightface_models/models/auraface/
        model_dir = Path("insightface_models/models/auraface")
        model_dir.mkdir(parents=True, exist_ok=True)

        print("Downloading AuraFace-v1 from HuggingFace...")
        snapshot_download(
            repo_id="fal/AuraFace-v1",
            local_dir=str(model_dir),
            local_dir_use_symlinks=False
        )
        print(f"AuraFace model downloaded to: {model_dir}")
        return True

    except ImportError:
        print("ERROR: huggingface_hub not installed. Run: pip install huggingface-hub")
        return False
    except Exception as e:
        print(f"ERROR downloading AuraFace: {e}")
        return False


def setup_silent_face_anti_spoofing():
    """Clone Silent-Face-Anti-Spoofing repository and copy models."""
    print("\n" + "=" * 50)
    print("Setting up Silent-Face-Anti-Spoofing...")
    print("=" * 50)

    repo_dir = Path("Silent-Face-Anti-Spoofing")
    models_dir = Path("models/anti_spoof")

    # Clone repository if not exists or is empty
    if not repo_dir.exists() or not any(repo_dir.iterdir()):
        print("Cloning Silent-Face-Anti-Spoofing repository...")
        # Remove empty directory if it exists
        if repo_dir.exists():
            shutil.rmtree(repo_dir)
        try:
            result = subprocess.run(
                ["git", "clone", "--depth", "1", "https://github.com/minivision-ai/Silent-Face-Anti-Spoofing.git"],
                capture_output=True,
                text=True
            )
            if result.returncode != 0:
                print(f"ERROR cloning repository: {result.stderr}")
                return False
            print("Repository cloned successfully.")
        except FileNotFoundError:
            print("ERROR: git not found. Please install git and try again.")
            return False
    else:
        print("Repository already exists, skipping clone.")

    # Copy model files
    models_dir.mkdir(parents=True, exist_ok=True)

    # Copy anti_spoof_models
    src_anti_spoof = repo_dir / "resources" / "anti_spoof_models"
    dst_anti_spoof = models_dir / "anti_spoof_models"

    if src_anti_spoof.exists():
        if dst_anti_spoof.exists():
            shutil.rmtree(dst_anti_spoof)
        shutil.copytree(src_anti_spoof, dst_anti_spoof)
        print(f"Copied anti_spoof_models to: {dst_anti_spoof}")
    else:
        print(f"WARNING: {src_anti_spoof} not found")

    # Copy detection_model
    src_detection = repo_dir / "resources" / "detection_model"
    dst_detection = models_dir / "detection_model"

    if src_detection.exists():
        if dst_detection.exists():
            shutil.rmtree(dst_detection)
        shutil.copytree(src_detection, dst_detection)
        print(f"Copied detection_model to: {dst_detection}")
    else:
        print(f"WARNING: {src_detection} not found")

    return True


def verify_models():
    """Verify all required model files exist."""
    print("\n" + "=" * 50)
    print("Verifying model files...")
    print("=" * 50)

    required_files = [
        # AuraFace models (InsightFace expects {root}/models/{name}/)
        "insightface_models/models/auraface",
        # Anti-spoofing models
        "models/anti_spoof/anti_spoof_models",
        "models/anti_spoof/detection_model",
    ]

    all_exist = True
    for file_path in required_files:
        path = Path(file_path)
        exists = path.exists()
        status = "OK" if exists else "MISSING"
        print(f"  [{status}] {file_path}")
        if not exists:
            all_exist = False

    return all_exist


def main():
    """Main setup function."""
    print("\n" + "#" * 60)
    print("# KYC POC - Model Setup")
    print("#" * 60 + "\n")

    # Change to script directory
    script_dir = Path(__file__).parent
    os.chdir(script_dir)

    success = True

    # Setup AuraFace
    if not setup_auraface():
        success = False

    # Setup Silent-Face-Anti-Spoofing
    if not setup_silent_face_anti_spoofing():
        success = False

    # Verify all models
    if not verify_models():
        success = False

    print("\n" + "#" * 60)
    if success:
        print("# Setup completed successfully!")
        print("# You can now run the application with: uvicorn app.main:app --reload")
    else:
        print("# Setup completed with errors. Please check the messages above.")
    print("#" * 60 + "\n")

    return 0 if success else 1


if __name__ == "__main__":
    sys.exit(main())