setup.py

"""
Setup script for ALPR system
Verifies installation and downloads required models
"""

import sys
import subprocess
import platform
from pathlib import Path


def check_python_version():
    """Check Python version"""
    print("Checking Python version...")
    version = sys.version_info

    if version.major != 3 or version.minor < 10:
        print(f"❌ Python 3.10+ required. Current: {version.major}.{version.minor}")
        return False

    print(f"✓ Python {version.major}.{version.minor}.{version.micro}")
    return True


def check_cuda():
    """Check CUDA availability"""
    print("\nChecking CUDA...")

    try:
        import torch

        if torch.cuda.is_available():
            print(f"✓ CUDA available: {torch.cuda.get_device_name(0)}")
            print(f"  CUDA Version: {torch.version.cuda}")
            print(f"  Device Count: {torch.cuda.device_count()}")
            return True
        else:
            print("⚠ CUDA not available - will use CPU (slower)")
            return False

    except ImportError:
        print("❌ PyTorch not installed")
        return False


def check_dependencies():
    """Check if all dependencies are installed"""
    print("\nChecking dependencies...")

    required_packages = [
        'cv2',
        'torch',
        'ultralytics',
        'easyocr',
        'yaml',
        'loguru'
    ]

    missing = []

    for package in required_packages:
        try:
            if package == 'cv2':
                import cv2
            elif package == 'torch':
                import torch
            elif package == 'ultralytics':
                import ultralytics
            elif package == 'easyocr':
                import easyocr
            elif package == 'yaml':
                import yaml
            elif package == 'loguru':
                import loguru

            print(f"✓ {package}")

        except ImportError:
            print(f"❌ {package} not found")
            missing.append(package)

    if missing:
        print(f"\n❌ Missing packages: {', '.join(missing)}")
        print("Run: pip install -r requirements.txt")
        return False

    return True


def download_models():
    """Download required YOLO models"""
    print("\nDownloading models...")

    try:
        from ultralytics import YOLO

        # Download YOLOv8 nano model (small and fast)
        print("Downloading YOLOv8n...")
        model = YOLO('yolov8n.pt')
        print("✓ YOLOv8n downloaded")

        return True

    except Exception as e:
        print(f"❌ Error downloading models: {e}")
        return False


def create_directories():
    """Create necessary directories"""
    print("\nCreating directories...")

    dirs = ['logs', 'detections', 'models']

    for dir_name in dirs:
        dir_path = Path(dir_name)
        dir_path.mkdir(exist_ok=True)
        print(f"✓ {dir_name}/")

    return True


def test_camera():
    """Test camera access"""
    print("\nTesting camera...")

    try:
        import cv2

        cap = cv2.VideoCapture(0)

        if not cap.isOpened():
            print("⚠ Could not open default camera (index 0)")
            print("  You may need to configure a different camera source")
            return False

        ret, frame = cap.read()
        cap.release()

        if ret:
            print(f"✓ Camera working: {frame.shape[1]}x{frame.shape[0]}")
            return True
        else:
            print("⚠ Camera opened but could not read frame")
            return False

    except Exception as e:
        print(f"❌ Error testing camera: {e}")
        return False


def main():
    """Main setup function"""
    print("=" * 60)
    print("ALPR System - Setup and Verification")
    print("=" * 60)

    results = {
        'Python Version': check_python_version(),
        'CUDA': check_cuda(),
        'Dependencies': check_dependencies(),
        'Directories': create_directories(),
        'Models': download_models(),
        'Camera': test_camera()
    }

    print("\n" + "=" * 60)
    print("Setup Summary")
    print("=" * 60)

    for name, status in results.items():
        symbol = "✓" if status else "❌"
        print(f"{symbol} {name}")

    all_good = all(results.values())

    print("\n" + "=" * 60)

    if all_good:
        print("✓ Setup complete! You can now run: python main.py")
    else:
        print("⚠ Some checks failed. Please review and fix issues above.")

    print("=" * 60)

    return 0 if all_good else 1


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