5_predict.py

"""
Step 5: Inference Tool (Enhanced)
Command-line tool for making predictions on EEG .edf files 
using the trained model from Step 4.

Enhanced features:
- Single file prediction
- Batch processing of multiple files
- Drag-and-drop support
- CSV output for batch results

Usage:
    # Single file (CNN model)
    python 5_predict.py --model_path models/eeg_classify_cnn_v1.h5 --file_path path/to/file.edf
    
    # Multiple files (LSTM model)
    python 5_predict.py --model_path models/eeg_classify_lstm_v1.h5 --file_path file1.edf file2.edf file3.edf
    
    # Directory of files (CNN model)
    python 5_predict.py --model_path models/eeg_classify_cnn_v1.h5 --batch_dir path/to/edf_files/
    
    # Drag and drop (Windows, LSTM model)
    python 5_predict.py --model_path models/eeg_classify_lstm_v1.h5 file1.edf file2.edf
"""

import tensorflow as tf
import numpy as np
import mne
import argparse
import os
import pickle
import glob
import csv
from pathlib import Path
import sys
from datetime import datetime

# Segmentation parameters (must match those used in Step 2)
SEGMENT_LENGTH = 1000
OVERLAP = 500

def create_segments(data_array, segment_length=SEGMENT_LENGTH, overlap=OVERLAP):
    """
    Create overlapping segments from a data array.
    This function must use the EXACT same parameters as Step 2.
    
    Args:
        data_array (np.ndarray): Input data of shape (timesteps, channels)
        segment_length (int): Length of each segment
        overlap (int): Number of overlapping time steps
        
    Returns:
        np.ndarray: Array of segments with shape (num_segments, segment_length, channels)
    """
    segments = []
    step_size = segment_length - overlap
    
    # Calculate number of segments
    num_timesteps = data_array.shape[0]
    
    # Create segments with overlap
    start_idx = 0
    while start_idx + segment_length <= num_timesteps:
        segment = data_array[start_idx:start_idx + segment_length]
        segments.append(segment)
        start_idx += step_size
    
    return np.array(segments)

def load_and_preprocess_edf(file_path, scaler):
    """
    Load an EDF file and preprocess it for prediction.
    
    Args:
        file_path (str): Path to the EDF file
        scaler: Fitted StandardScaler object
        
    Returns:
        np.ndarray: Preprocessed segments ready for prediction
    """
    try:
        # Load EDF file with MNE
        raw = mne.io.read_raw_edf(file_path, preload=True, verbose=False)
        
        # Extract numerical data and transpose to (timesteps, channels) format
        data = raw.get_data().T
        
        print(f"Loaded EDF file: {os.path.basename(file_path)}")
        print(f"  Original shape: {data.shape} (timesteps, channels)")
        print(f"  Duration: {data.shape[0] / raw.info['sfreq']:.2f} seconds")
        print(f"  Sampling rate: {raw.info['sfreq']} Hz")
        print(f"  Number of channels: {data.shape[1]}")
        
        # Create segments
        segments = create_segments(data)
        print(f"  Created {len(segments)} segments")
        
        if len(segments) == 0:
            raise ValueError(f"No segments could be created. File too short? "
                           f"Minimum length needed: {SEGMENT_LENGTH} timesteps")
        
        # Normalize using the same scaler as training
        original_shape = segments.shape
        segments_reshaped = segments.reshape(-1, segments.shape[-1])
        segments_scaled = scaler.transform(segments_reshaped)
        segments_scaled = segments_scaled.reshape(original_shape)
        
        return segments_scaled.astype(np.float32)
        
    except Exception as e:
        raise RuntimeError(f"Error processing EDF file '{file_path}': {str(e)}")

def make_prediction(model, segments):
    """
    Make predictions on EEG segments and aggregate the results.
    
    Args:
        model: Trained Keras model
        segments (np.ndarray): Preprocessed segments
        
    Returns:
        tuple: (prediction_label, confidence_score, individual_predictions)
    """
    # Get predictions for all segments
    predictions = model.predict(segments, verbose=0)
    
    # predictions is array of shape (num_segments, 1) with probabilities
    probabilities = predictions.flatten()
    
    # Aggregate predictions (average probability)
    avg_probability = np.mean(probabilities)
    
    # Final prediction (0 = healthy, 1 = schizophrenia)
    prediction_label = int(avg_probability > 0.5)
    
    # Confidence score
    if prediction_label == 1:
        confidence = avg_probability  # Confidence in schizophrenia prediction
    else:
        confidence = 1 - avg_probability  # Confidence in healthy prediction
    
    return prediction_label, confidence, probabilities

def process_single_file(file_path, model, scaler, verbose=False):
    """
    Process a single EDF file and return prediction results.
    
    Args:
        file_path (str): Path to the EDF file
        model: Trained Keras model
        scaler: Fitted StandardScaler object
        verbose (bool): Show detailed analysis
        
    Returns:
        dict: Prediction results with metadata
    """
    results = {
        'file_path': file_path,
        'file_name': os.path.basename(file_path),
        'success': False,
        'error_message': None,
        'prediction_label': None,
        'prediction_class': None,
        'confidence': None,
        'num_segments': None,
        'processing_time': None
    }
    
    start_time = datetime.now()
    
    try:
        # Load and preprocess the EDF file
        segments = load_and_preprocess_edf(file_path, scaler)
        
        # Make prediction
        prediction_label, confidence, probabilities = make_prediction(model, segments)
        
        # Store results
        results.update({
            'success': True,
            'prediction_label': prediction_label,
            'prediction_class': "Schizophrenia" if prediction_label == 1 else "Healthy",
            'confidence': confidence,
            'num_segments': len(segments),
            'processing_time': (datetime.now() - start_time).total_seconds()
        })
        
        # Display results
        if verbose:
            print_detailed_results(prediction_label, confidence, probabilities, file_path)
        else:
            class_name = "Schizophrenia" if prediction_label == 1 else "Healthy"
            print(f"📁 {os.path.basename(file_path)}: {class_name} (Confidence: {confidence * 100:.1f}%)")
        
    except Exception as e:
        results['error_message'] = str(e)
        print(f"❌ Error processing {os.path.basename(file_path)}: {str(e)}")
    
    return results

def find_edf_files(paths):
    """
    Find all EDF files from various input types.
    
    Args:
        paths (list): List of file paths, directory paths, or glob patterns
        
    Returns:
        list: List of valid EDF file paths
    """
    edf_files = []
    
    for path in paths:
        path = Path(path)
        
        if path.is_file() and path.suffix.lower() == '.edf':
            # Single EDF file
            edf_files.append(str(path))
        elif path.is_dir():
            # Directory - find all EDF files recursively
            pattern = path / "**" / "*.edf"
            edf_files.extend(glob.glob(str(pattern), recursive=True))
        elif '*' in str(path) or '?' in str(path):
            # Glob pattern
            edf_files.extend(glob.glob(str(path)))
        else:
            # Check if it's an EDF file without extension or with different case
            if path.exists():
                edf_files.append(str(path))
    
    # Remove duplicates and sort
    edf_files = sorted(list(set(edf_files)))
    
    # Validate files exist and are EDF files
    valid_files = []
    for file_path in edf_files:
        if os.path.exists(file_path):
            try:
                # Quick check if it's a valid EDF file
                raw = mne.io.read_raw_edf(file_path, preload=False, verbose=False)
                valid_files.append(file_path)
            except:
                print(f"⚠️ Skipping invalid EDF file: {os.path.basename(file_path)}")
    
    return valid_files

def save_batch_results(results, output_path):
    """
    Save batch processing results to CSV file.
    
    Args:
        results (list): List of prediction result dictionaries
        output_path (str): Path to save CSV file
    """
    if not results:
        return
    
    fieldnames = ['file_name', 'file_path', 'success', 'prediction_class', 
                  'confidence', 'num_segments', 'processing_time', 'error_message']
    
    with open(output_path, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        
        for result in results:
            # Convert confidence to percentage for CSV
            if result['confidence'] is not None:
                result['confidence'] = f"{result['confidence'] * 100:.2f}%"
            writer.writerow(result)
    
    print(f"📊 Batch results saved to: {output_path}")

def print_detailed_results(prediction_label, confidence, probabilities, file_path):
    """
    Print detailed prediction results.
    
    Args:
        prediction_label (int): Final prediction (0 or 1)
        confidence (float): Confidence score
        probabilities (np.ndarray): Individual segment predictions
        file_path (str): Path to the input file
    """
    print("\n" + "=" * 60)
    print("PREDICTION RESULTS")
    print("=" * 60)
    
    # Main prediction
    class_name = "Schizophrenia" if prediction_label == 1 else "Healthy"
    print(f"File: {os.path.basename(file_path)}")
    print(f"Prediction: {class_name}")
    print(f"Confidence: {confidence * 100:.1f}%")
    
    # Detailed statistics
    print("\nDetailed Analysis:")
    print(f"  Number of segments analyzed: {len(probabilities)}")
    print(f"  Average probability (schizophrenia): {np.mean(probabilities):.4f}")
    print(f"  Standard deviation: {np.std(probabilities):.4f}")
    print(f"  Min probability: {np.min(probabilities):.4f}")
    print(f"  Max probability: {np.max(probabilities):.4f}")
    
    # Segment-wise consensus
    segment_predictions = (probabilities > 0.5).astype(int)
    schizophrenia_segments = np.sum(segment_predictions)
    healthy_segments = len(segment_predictions) - schizophrenia_segments
    
    print(f"  Segments predicting schizophrenia: {schizophrenia_segments} ({schizophrenia_segments/len(probabilities)*100:.1f}%)")
    print(f"  Segments predicting healthy: {healthy_segments} ({healthy_segments/len(probabilities)*100:.1f}%)")
    
    # Confidence level interpretation
    print(f"\nConfidence Level: ", end="")
    if confidence >= 0.9:
        print("Very High")
    elif confidence >= 0.8:
        print("High")
    elif confidence >= 0.7:
        print("Moderate")
    elif confidence >= 0.6:
        print("Low")
    else:
        print("Very Low")
    
    print("=" * 60)

def main():
    """Main prediction process with batch support."""
    parser = argparse.ArgumentParser(
        description="EEG Classification Prediction Tool - Enhanced with Batch Processing",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Single file
  python 5_predict.py --model_path models/eeg_classify_cnn_v1.h5 --files data/test.edf
  
  # Multiple files
  python 5_predict.py --model_path models/eeg_classify_cnn_v1.h5 --files file1.edf file2.edf file3.edf
  
  # Directory of files
  python 5_predict.py --model_path models/eeg_classify_cnn_v1.h5 --batch_dir data/edf_files/
  
  # Drag and drop (pass files as arguments)
  python 5_predict.py --model_path models/eeg_classify_cnn_v1.h5 file1.edf file2.edf
  
  # With verbose output and CSV export
  python 5_predict.py --model_path models/eeg_classify_cnn_v1.h5 --files *.edf --verbose --csv_output results.csv
        """
    )
    
    parser.add_argument(
        "--model_path", 
        required=True,
        help="Path to the trained model file (.h5)"
    )
    
    parser.add_argument(
        "--files", 
        nargs="*",
        help="One or more EDF files to analyze"
    )
    
    parser.add_argument(
        "--batch_dir",
        help="Directory containing EDF files to process (recursive search)"
    )
    
    parser.add_argument(
        "--scaler_path",
        default="models/scaler_cnn.pkl",
        help="Path to the fitted scaler file (default: models/scaler_cnn.pkl)"
    )
    
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Show detailed prediction analysis"
    )
    
    parser.add_argument(
        "--csv_output",
        help="Save batch results to CSV file"
    )
    
    # Parse known args to handle drag-and-drop files as positional arguments
    args, unknown_args = parser.parse_known_args()
    
    # Combine files from --files and unknown args (drag-and-drop)
    all_files = []
    if args.files:
        all_files.extend(args.files)
    if unknown_args:
        all_files.extend(unknown_args)
    
    # Add files from batch directory
    if args.batch_dir:
        all_files.append(args.batch_dir)
    
    # If no files specified, show help
    if not all_files:
        parser.print_help()
        print(f"\nError: No input files specified.")
        print("Use --files, --batch_dir, or drag-and-drop files onto the script.")
        return 1
    
    print("EEG Classification Prediction Tool (Enhanced)")
    print("=" * 60)
    
    # Validate input files
    if not os.path.exists(args.model_path):
        print(f"Error: Model file '{args.model_path}' does not exist.")
        return 1
    
    if not os.path.exists(args.scaler_path):
        print(f"Error: Scaler file '{args.scaler_path}' does not exist.")
        return 1
    
    try:
        # Find all EDF files
        print("🔍 Finding EDF files...")
        edf_files = find_edf_files(all_files)
        
        if not edf_files:
            print("❌ No valid EDF files found.")
            return 1
        
        print(f"📁 Found {len(edf_files)} EDF file(s) to process")
        
        # Load the trained model
        print(f"🤖 Loading model: {args.model_path}")
        model = tf.keras.models.load_model(args.model_path)
        print("  ✓ Model loaded successfully")
        
        # Load the scaler
        print(f"📊 Loading scaler: {args.scaler_path}")
        with open(args.scaler_path, 'rb') as f:
            scaler = pickle.load(f)
        print("  ✓ Scaler loaded successfully")
        
        # Process files
        print(f"\n🔄 Processing {len(edf_files)} file(s)...")
        print("-" * 60)
        
        all_results = []
        successful_predictions = 0
        
        for i, file_path in enumerate(edf_files, 1):
            print(f"\n[{i}/{len(edf_files)}] Processing: {os.path.basename(file_path)}")
            
            # Process single file
            result = process_single_file(file_path, model, scaler, args.verbose)
            all_results.append(result)
            
            if result['success']:
                successful_predictions += 1
        
        # Summary
        print("\n" + "=" * 60)
        print("BATCH PROCESSING SUMMARY")
        print("=" * 60)
        print(f"Total files processed: {len(edf_files)}")
        print(f"Successful predictions: {successful_predictions}")
        print(f"Failed predictions: {len(edf_files) - successful_predictions}")
        
        if successful_predictions > 0:
            # Calculate overall statistics
            successful_results = [r for r in all_results if r['success']]
            healthy_count = sum(1 for r in successful_results if r['prediction_label'] == 0)
            schizo_count = sum(1 for r in successful_results if r['prediction_label'] == 1)
            avg_confidence = np.mean([r['confidence'] for r in successful_results])
            total_time = sum([r['processing_time'] for r in successful_results])
            
            print(f"\nPrediction Distribution:")
            print(f"  Healthy: {healthy_count} ({healthy_count/successful_predictions*100:.1f}%)")
            print(f"  Schizophrenia: {schizo_count} ({schizo_count/successful_predictions*100:.1f}%)")
            print(f"  Average confidence: {avg_confidence*100:.1f}%")
            print(f"  Total processing time: {total_time:.2f} seconds")
            print(f"  Average time per file: {total_time/successful_predictions:.2f} seconds")
        
        # Save CSV results if requested
        if args.csv_output:
            save_batch_results(all_results, args.csv_output)
        
        print(f"\n✅ Batch processing completed!")
        
        return 0
        
    except Exception as e:
        print(f"❌ Error during prediction: {str(e)}")
        return 1

if __name__ == "__main__":
    # Suppress TensorFlow warnings for cleaner output
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    
    # Suppress MNE warnings for cleaner output
    mne.set_log_level('ERROR')
    
    exit_code = main()
    exit(exit_code)