src/train.py

import os
import sys
import pandas as pd
import numpy as np
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
from tensorflow.keras.metrics import AUC
import matplotlib.pyplot as plt


# --- BEGIN ROBUST GPU FIX ---
# We must do this before any other TensorFlow operations
import tensorflow as tf
print("Applying robust GPU configuration...")
try:
    # Get all GPUs that TensorFlow can see
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        for gpu in gpus:
            # Set memory growth to True for each GPU
            tf.config.experimental.set_memory_growth(gpu, True)
            print(f"  > Enabled memory growth for: {gpu.name}")
    else:
        print("  > No GPUs found by TensorFlow. Will run on CPU.")
except Exception as e:
    print(f"  > Error applying GPU configuration: {e}")
# --- END ROBUST GPU FIX ---


# 1. --- Import from our project files ---
try:
    import config
    from model import build_baseline_model
except ImportError:
    print("Error: Could not import config.py or model.py.")
    print("Make sure they are in the 'src/' directory.")
    sys.exit(1)
    
# Suppress TensorFlow logs
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 

def create_balanced_dataframe():
    """
    Scans the training directory and creates a balanced DataFrame
    of 'real' and 'fake' image paths. This is our undersampling step.
    """
    print("Creating balanced training dataframe...")
    
    # Get lists of all real and fake training images
    real_paths = [os.path.join(config.TRAIN_REAL_DIR, f) for f in os.listdir(config.TRAIN_REAL_DIR) if f.endswith('.jpg')]
    fake_paths = [os.path.join(config.TRAIN_FAKE_DIR, f) for f in os.listdir(config.TRAIN_FAKE_DIR) if f.endswith('.jpg')]
    
    # Create DataFrames
    df_real = pd.DataFrame({'filepath': real_paths, 'label': 'real'})
    df_fake = pd.DataFrame({'filepath': fake_paths, 'label': 'fake'})
    
    # --- This is the key undersampling step ---
    # We sample the 'fake' DataFrame to have the same number of
    # images as the 'real' DataFrame.
    df_fake_sampled = df_fake.sample(n=len(df_real), random_state=42)
    
    # Combine and shuffle
    df_train_balanced = pd.concat([df_real, df_fake_sampled]).sample(frac=1, random_state=42).reset_index(drop=True)
    
    print(f"Balanced training set created: {len(df_train_balanced)} total images")
    print(f"  Real: {len(df_real)} images")
    print(f"  Fake: {len(df_fake_sampled)} images")
    
    return df_train_balanced

def create_generators(train_df):
    """
    Creates the Keras Data Generators for training and validation.
    """
    print("Creating Data Generators...")
    
    # --- Training Generator with Data Augmentation ---
    # Data augmentation creates "new" versions of our images on-the-fly
    # (flipped, rotated, etc.) to make our model more robust.
    train_datagen = ImageDataGenerator(
        rescale=1./255,          # Normalize pixel values
        rotation_range=20,
        width_shift_range=0.1,
        height_shift_range=0.1,
        horizontal_flip=True,
        fill_mode='nearest',
        validation_split=0.2     # We'll use 20% of our training data for validation
    )
    
    # --- Test/Validation Generator (No Augmentation) ---
    # We *never* augment our validation or test data.
    test_datagen = ImageDataGenerator(rescale=1./255)
    
    # --- Create the generators from our DataFrames ---
    
    # Training Generator (from the balanced DataFrame)
    train_generator = train_datagen.flow_from_dataframe(
        dataframe=train_df,
        x_col='filepath',
        y_col='label',
        target_size=(config.TARGET_IMAGE_SIZE, config.TARGET_IMAGE_SIZE),
        batch_size=config.BATCH_SIZE,
        class_mode='binary',
        subset='training',
        shuffle=True
    )
    
    # Validation Generator (also from the balanced DataFrame)
    validation_generator = train_datagen.flow_from_dataframe(
        dataframe=train_df,
        x_col='filepath',
        y_col='label',
        target_size=(config.TARGET_IMAGE_SIZE, config.TARGET_IMAGE_SIZE),
        batch_size=config.BATCH_SIZE,
        class_mode='binary',
        subset='validation',
        shuffle=False  # No need to shuffle validation data
    )
    
    # Test Generator (from the *unbalanced* test directory)
    # This is our real-world test.
    test_generator = test_datagen.flow_from_directory(
        directory=config.TEST_DIR,
        target_size=(config.TARGET_IMAGE_SIZE, config.TARGET_IMAGE_SIZE),
        batch_size=config.BATCH_SIZE,
        class_mode='binary',
        shuffle=False
    )
    
    return train_generator, validation_generator, test_generator

def plot_history(history, save_path):
    """
    Plots the training history (accuracy and loss) and saves it to a file.
    """
    print(f"Saving training history plot to {save_path}...")
    
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10))

    # Plot training & validation accuracy values
    ax1.plot(history.history['accuracy'])
    ax1.plot(history.history['val_accuracy'])
    ax1.set_title('Model Accuracy')
    ax1.set_ylabel('Accuracy')
    ax1.set_xlabel('Epoch')
    ax1.legend(['Train', 'Validation'], loc='upper left')

    # Plot training & validation loss values
    ax2.plot(history.history['loss'])
    ax2.plot(history.history['val_loss'])
    ax2.set_title('Model Loss')
    ax2.set_ylabel('Loss')
    ax2.set_xlabel('Epoch')
    ax2.legend(['Train', 'Validation'], loc='upper left')
    
    plt.tight_layout()
    plt.savefig(save_path)
    print("History plot saved.")


def main():
    """
    Main training function.
    """
    print("--- Phase 2: Starting Baseline Model Training ---")
    
    # 1. Handle class imbalance
    train_df = create_balanced_dataframe()
    
    # 2. Create data generators
    train_gen, val_gen, test_gen = create_generators(train_df)
    
    # 3. Build the model
    print("Building model...")
    model = build_baseline_model(config.TARGET_IMAGE_SIZE)
    
    # 4. Compile the model
    # We use AUC (Area Under the Curve) as our main metric.
    # It's much better than accuracy for imbalanced test sets.
    model.compile(
        optimizer=Adam(learning_rate=config.LEARNING_RATE),
        loss='binary_crossentropy',
        metrics=['accuracy', AUC(name='auc')]
    )
    
    model.summary()
    
    # 5. Define Callbacks
    # This will save the *best* model based on validation AUC
    checkpoint_path = os.path.join(config.MODEL_DIR, "baseline_model.h5")
    model_checkpoint = ModelCheckpoint(
        filepath=checkpoint_path,
        save_best_only=True,
        monitor='val_auc',
        mode='max',
        verbose=1
    )
    
    # This will stop training early if it stops improving
    early_stopping = EarlyStopping(
        monitor='val_auc',
        mode='max',
        patience=5,  # Stop after 5 epochs of no improvement
        verbose=1,
        restore_best_weights=True
    )
    
    # 6. Start Training
    print("Starting model training...")
    history = model.fit(
        train_gen,
        steps_per_epoch=train_gen.n // config.BATCH_SIZE,
        validation_data=val_gen,
        validation_steps=val_gen.n // config.BATCH_SIZE,
        epochs=config.EPOCHS,
        callbacks=[model_checkpoint, early_stopping]
    )
    
    print("Training complete.")
    
    # 7. Evaluate on the (imbalanced) Test Set
    print("Evaluating model on the test set...")
    results = model.evaluate(test_gen, steps=test_gen.n // config.BATCH_SIZE)
    
    print("\n--- Test Set Evaluation ---")
    print(f"Test Loss:     {results[0]:.4f}")
    print(f"Test Accuracy: {results[1]:.4f}")
    print(f"Test AUC:      {results[2]:.4f}")
    
    # 8. Save history plot
    plot_path = os.path.join(config.RESULTS_DIR, "baseline_training_history.png")
    plot_history(history, plot_path)
    
    print("\n--- Baseline Model Training Finished ---")
    print(f"Best model saved to: {checkpoint_path}")

if __name__ == "__main__":
    
    # Create models/ and results/ directories if they don't exist
    os.makedirs(config.MODEL_DIR, exist_ok=True)
    os.makedirs(config.RESULTS_DIR, exist_ok=True)
    
    main()