main.py

#!/usr/bin/env python3
"""
RecognAIze Main Script
=====================

Command-line interface for the RecognAIze face recognition system.
This script provides end-to-end functionality for face detection, 
identification, and report generation.

Usage:
    python main.py [command] [options]

Commands:
    inference    - Run inference on test images
    train        - Process training data and build prototypes
    analyze      - Analyze dataset statistics
    test         - Run system tests
"""

import argparse
import logging
import sys
import os
from pathlib import Path
import json
from datetime import datetime

# Add src to path for imports
sys.path.insert(0, str(Path(__file__).parent / 'src'))

from src.face_recognition.inference.test_inference import TestInferencePipeline
from src.face_recognition.inference.pipeline import FaceRecognitionPipeline
from src.face_recognition.inference.embedding_generator import EmbeddingGenerator


# Configure logging
def setup_logging(level=logging.INFO, log_file=None):
    """Set up logging configuration."""
    log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    
    handlers = [logging.StreamHandler(sys.stdout)]
    if log_file:
        handlers.append(logging.FileHandler(log_file))
    
    logging.basicConfig(
        level=level,
        format=log_format,
        handlers=handlers
    )


def run_inference(args):
    """Run inference on test images."""
    print("=" * 60)
    print("RecognAIze - Running Inference Pipeline")
    print("=" * 60)
    
    # Check dataset path
    if not os.path.exists(args.dataset):
        print(f"Error: Dataset path '{args.dataset}' does not exist!")
        return False
    
    try:
        # Initialize inference pipeline
        print(f"Initializing inference pipeline...")
        print(f"Dataset: {args.dataset}")
        print(f"Similarity threshold: {args.threshold}")
        print(f"Face confidence: {args.confidence}")
        print(f"Output directory: {args.output}")
        
        inference = TestInferencePipeline(
            dataset_root=args.dataset,
            output_dir=args.output,
            similarity_threshold=args.threshold
        )
        
        # Run complete inference
        print("\nRunning complete inference pipeline...")
        output_files = inference.run_complete_inference(
            min_confidence=args.confidence,
            max_images=args.max_images,
            save_debug=args.debug
        )
        
        # Display results
        print("\n" + "=" * 60)
        print("INFERENCE COMPLETE!")
        print("=" * 60)
        
        print(f"\nOutput files generated:")
        for name, path in output_files.items():
            print(f"  {name.upper()}: {path}")
        
        # Show quick stats
        stats = inference.inference_stats
        total = stats['total_images']
        detected = stats['images_with_faces']
        identified = stats['faces_identified']
        
        print(f"\nQuick Statistics:")
        print(f"  Total images processed: {total}")
        print(f"  Images with faces: {detected} ({100*detected/total:.1f}%)")
        print(f"  Total faces found: {stats['total_faces_found']}")
        print(f"  Employees identified: {identified}")
        print(f"  Unknown faces: {stats['faces_unknown']}")
        
        # Show prediction summary
        summary = inference.get_employee_prediction_summary()
        print(f"\nPrediction Summary:")
        for emp_id, count in list(summary.items())[:10]:  # Show top 10
            print(f"  {emp_id}: {count} images")
        if len(summary) > 10:
            print(f"  ... and {len(summary) - 10} more employees")
        
        return True
        
    except Exception as e:
        print(f"Error during inference: {e}")
        if args.debug:
            import traceback
            traceback.print_exc()
        return False


def run_training(args):
    """Process training data and build employee prototypes."""
    print("=" * 60)
    print("RecognAIze - Training Pipeline")
    print("=" * 60)
    
    try:
        # Initialize embedding generator
        print(f"Initializing training pipeline...")
        print(f"Dataset: {args.dataset}")
        
        embedding_gen = EmbeddingGenerator(
            dataset_root=args.dataset
        )
        
        # Generate reference embeddings
        print("\nGenerating reference embeddings...")
        embedding_gen.generate_reference_embeddings(
            min_confidence=args.confidence
        )
        
        # Build employee prototypes
        print("Building employee prototypes...")
        embedding_gen.build_employee_prototypes(
            method="mean",
            min_faces=1
        )
        
        # Initialize similarity matcher for threshold tuning
        print("\nTuning similarity threshold on training data...")
        from src.face_recognition.inference.similarity_matcher import SimilarityMatcher
        
        matcher = SimilarityMatcher(
            embedding_generator=embedding_gen,
            similarity_threshold=0.5  # Initial guess
        )
        
        # Evaluate on training data to find optimal threshold
        print("Evaluating on training data to optimize threshold...")
        evaluation_results = matcher.evaluate_on_training_data(max_samples=200)
        
        if evaluation_results:
            print(f"Training evaluation results:")
            print(f"  Overall accuracy: {evaluation_results['overall_accuracy']:.3f}")
            print(f"  Known accuracy: {evaluation_results['known_accuracy']:.3f}")
            print(f"  Unknown accuracy: {evaluation_results['unknown_accuracy']:.3f}")
            print(f"  Current threshold: {evaluation_results['similarity_threshold']:.3f}")
        
        # Generate training embeddings for threshold tuning
        print("Generating training embeddings for threshold optimization...")
        training_embeddings = embedding_gen.generate_training_embeddings(
            include_unknown=True,
            max_images=300,
            save_embeddings=False
        )
        
        if training_embeddings:
            # Prepare validation data
            val_embeddings = []
            val_labels = []
            for image_id, data in training_embeddings.items():
                val_embeddings.append(data['embedding'])
                val_labels.append(data['employee_id'])
            
            import numpy as np
            val_embeddings = np.array(val_embeddings)
            
            # Tune threshold
            print("Optimizing similarity threshold...")
            best_threshold, best_accuracy = matcher.tune_threshold(
                val_embeddings, 
                val_labels,
                threshold_range=(0.4, 0.95),
                num_steps=30
            )
            
            print(f"\nOptimal threshold found: {best_threshold:.3f}")
            print(f"Best validation accuracy: {best_accuracy:.3f}")
            
            # Save optimal threshold
            config_path = "optimal_config.json"
            config = {
                "optimal_similarity_threshold": float(best_threshold),
                "validation_accuracy": float(best_accuracy),
                "tuned_on": datetime.now().isoformat()
            }
            
            with open(config_path, 'w') as f:
                json.dump(config, f, indent=2)
            print(f"Optimal configuration saved to: {config_path}")
        
        # Show statistics
        num_employees = len(embedding_gen.employee_prototypes)
        ref_stats = embedding_gen.get_embedding_statistics()
        
        print("\n" + "=" * 60)
        print("TRAINING COMPLETE!")
        print("=" * 60)
        
        print(f"\nResults:")
        print(f"  Employee prototypes created: {num_employees}")
        print(f"  Total reference faces processed: {ref_stats.get('total_reference_faces', 0)}")
        print(f"  Average faces per employee: {ref_stats.get('avg_faces_per_employee', 0):.1f}")
        
        print(f"\nPrototype files saved:")
        print(f"  Reference embeddings: embeddings/reference_embeddings.pkl")
        print(f"  Employee prototypes: embeddings/employee_prototypes.pkl")
        
        if training_embeddings:
            print(f"\nRecommended similarity threshold: {best_threshold:.3f}")
            print(f"Use this threshold for inference: --threshold {best_threshold:.3f}")
        
        return True
        
    except Exception as e:
        print(f"Error during training: {e}")
        if args.debug:
            import traceback
            traceback.print_exc()
        return False


def run_analysis(args):
    """Analyze dataset statistics."""
    print("=" * 60)
    print("RecognAIze - Dataset Analysis")
    print("=" * 60)
    
    try:
        # Run data exploration
        import subprocess
        result = subprocess.run([
            sys.executable, 'data_exploration.py'
        ], capture_output=True, text=True)
        
        if result.returncode == 0:
            print(result.stdout)
            return True
        else:
            print(f"Error running data exploration: {result.stderr}")
            return False
            
    except Exception as e:
        print(f"Error during analysis: {e}")
        return False


def run_tests(args):
    """Run system tests."""
    print("=" * 60)
    print("RecognAIze - Running Tests")
    print("=" * 60)
    
    try:
        # Run test suite
        import subprocess
        
        if args.module:
            # Run specific test module
            cmd = [sys.executable, 'tests/run_tests.py', args.module]
        else:
            # Run all tests
            cmd = [sys.executable, 'tests/run_tests.py']
        
        result = subprocess.run(cmd, capture_output=False)
        return result.returncode == 0
        
    except Exception as e:
        print(f"Error running tests: {e}")
        return False


def create_config_file(args):
    """Create a configuration file with current settings."""
    config = {
        'dataset_path': args.dataset,
        'output_directory': args.output,
        'similarity_threshold': args.threshold,
        'face_confidence': args.confidence,
        'max_images': args.max_images,
        'debug_mode': args.debug,
        'created': datetime.now().isoformat()
    }
    
    config_path = 'recognaize_config.json'
    with open(config_path, 'w') as f:
        json.dump(config, f, indent=2)
    
    print(f"Configuration saved to: {config_path}")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="RecognAIze Face Recognition System",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run inference on test images
  python main.py inference --dataset ./dataset

  # Run training to build prototypes
  python main.py train --dataset ./dataset

  # Analyze dataset statistics
  python main.py analyze --dataset ./dataset

  # Run system tests
  python main.py test

  # Run with custom parameters
  python main.py inference --dataset ./dataset --threshold 0.5 --confidence 0.3
        """
    )
    
    # Add subcommands
    subparsers = parser.add_subparsers(dest='command', required=True, help='Available commands')
    
    # Inference command
    inference_parser = subparsers.add_parser(
        'inference', 
        help='Run inference on test images'
    )
    inference_parser.add_argument(
        '--dataset', 
        type=str, 
        default='identity-employees-in-surveillance-cctv/dataset',
        help='Path to dataset directory'
    )
    inference_parser.add_argument(
        '--output', 
        type=str, 
        default='results',
        help='Output directory for results'
    )
    inference_parser.add_argument(
        '--threshold', 
        type=float, 
        default=0.85,  # Much higher default threshold
        help='Similarity threshold for employee identification'
    )
    inference_parser.add_argument(
        '--confidence', 
        type=float, 
        default=0.3,  # Lower face detection threshold
        help='Minimum confidence for face detection'
    )
    inference_parser.add_argument(
        '--max-images', 
        type=int, 
        default=None,
        help='Maximum number of images to process (for testing)'
    )
    inference_parser.add_argument(
        '--debug', 
        action='store_true',
        help='Enable debug mode with detailed output'
    )
    
    # Training command
    train_parser = subparsers.add_parser(
        'train', 
        help='Process training data and build prototypes'
    )
    train_parser.add_argument(
        '--dataset', 
        type=str, 
        default='identity-employees-in-surveillance-cctv/dataset',
        help='Path to dataset directory'
    )
    train_parser.add_argument(
        '--confidence', 
        type=float, 
        default=0.2,  # Lower threshold for reference images
        help='Minimum confidence for face detection in reference images'
    )
    train_parser.add_argument(
        '--debug', 
        action='store_true',
        help='Enable debug mode'
    )
    
    # Analysis command
    analysis_parser = subparsers.add_parser(
        'analyze', 
        help='Analyze dataset statistics'
    )
    analysis_parser.add_argument(
        '--dataset', 
        type=str, 
        default='identity-employees-in-surveillance-cctv/dataset',
        help='Path to dataset directory'
    )
    
    # Test command
    test_parser = subparsers.add_parser(
        'test', 
        help='Run system tests'
    )
    test_parser.add_argument(
        '--module', 
        type=str, 
        default=None,
        help='Specific test module to run (e.g., face_detection)'
    )
    test_parser.add_argument(
        '--debug', 
        action='store_true',
        help='Enable debug mode'
    )
    
    # Global arguments
    parser.add_argument(
        '--log-level', 
        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'],
        default='INFO',
        help='Set logging level'
    )
    parser.add_argument(
        '--log-file', 
        type=str, 
        default=None,
        help='Log to file (in addition to console)'
    )
    parser.add_argument(
        '--save-config', 
        action='store_true',
        help='Save current configuration to file'
    )
    
    # Parse arguments
    args = parser.parse_args()
    
    # Set up logging
    log_level = getattr(logging, args.log_level)
    setup_logging(level=log_level, log_file=args.log_file)
    
    # Show welcome message
    print("RecognAIze Face Recognition System")
    print(f"Command: {args.command}")
    print(f"Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print()
    
    # Save configuration if requested
    if args.save_config and hasattr(args, 'dataset'):
        create_config_file(args)
    
    # Route to appropriate command
    success = False
    
    if args.command == 'inference':
        success = run_inference(args)
    elif args.command == 'train':
        success = run_training(args)
    elif args.command == 'analyze':
        success = run_analysis(args)
    elif args.command == 'test':
        success = run_tests(args)
    else:
        parser.print_help()
        return 1
    
    # Exit with appropriate code
    if success:
        print(f"\n✅ Command '{args.command}' completed successfully!")
        return 0
    else:
        print(f"\n❌ Command '{args.command}' failed!")
        return 1


if __name__ == '__main__':
    exit_code = main()
    sys.exit(exit_code)