NN.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, TensorDataset
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import KFold, StratifiedKFold
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
import matplotlib.pyplot as plt
import seaborn as sns
import time

# Import UCI ML Repository package for dataset fetching
try:
    from ucimlrepo import fetch_ucirepo
except ImportError:
    print("The ucimlrepo package is not installed. Please install it using:")
    print("pip install ucimlrepo")
    exit(1)

# Set random seed for reproducibility
torch.manual_seed(42)
np.random.seed(42)

# Define a PyTorch Dataset for the landmine data
class LandmineDataset(Dataset):
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels
        
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self, idx):
        x = self.features[idx]
        y = self.labels[idx]
        return torch.FloatTensor(x), torch.LongTensor([y])[0]  # Return long tensor for CrossEntropyLoss

# Define a neural network for multi-class landmine detection
class MultiClassLandmineDetector(nn.Module):
    def __init__(self, input_size, num_classes, hidden_size1=64, hidden_size2=32, hidden_size3=16):
        super(MultiClassLandmineDetector, self).__init__()
        
        # Define the layers
        self.fc1 = nn.Linear(input_size, hidden_size1)
        self.bn1 = nn.BatchNorm1d(hidden_size1)  # Batch normalization for better training
        self.dropout1 = nn.Dropout(0.2)  # Dropout for regularization
        
        self.fc2 = nn.Linear(hidden_size1, hidden_size2)
        self.bn2 = nn.BatchNorm1d(hidden_size2)
        self.dropout2 = nn.Dropout(0.2)
        
        self.fc3 = nn.Linear(hidden_size2, hidden_size3)
        self.bn3 = nn.BatchNorm1d(hidden_size3)
        
        # Output layer for multi-class classification
        self.fc4 = nn.Linear(hidden_size3, num_classes)
        
    def forward(self, x):
        # Forward pass through the network
        x = self.fc1(x)
        x = nn.functional.relu(x)
        if x.size(0) > 1:  # Batch normalization requires batch size > 1
            x = self.bn1(x)
        x = self.dropout1(x)
        
        x = self.fc2(x)
        x = nn.functional.relu(x)
        if x.size(0) > 1:
            x = self.bn2(x)
        x = self.dropout2(x)
        
        x = self.fc3(x)
        x = nn.functional.relu(x)
        if x.size(0) > 1:
            x = self.bn3(x)
        
        # No activation function here - CrossEntropyLoss will apply softmax
        x = self.fc4(x)
        return x

# Function to train the model for one fold
def train_model(model, train_loader, criterion, optimizer, num_epochs=100, verbose=True):
    # Lists to store metrics
    train_losses = []
    train_accuracies = []
    
    # Training loop
    for epoch in range(num_epochs):
        # Training phase
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        for inputs, labels in train_loader:
            # Zero the parameter gradients
            optimizer.zero_grad()
            
            # Forward pass
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            # Backward pass and optimize
            loss.backward()
            optimizer.step()
            
            # Track statistics
            train_loss += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs, 1)
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum().item()
        
        # Calculate average loss and accuracy for the epoch
        train_loss = train_loss / len(train_loader.dataset)
        train_accuracy = train_correct / train_total
        
        # Save the metrics
        train_losses.append(train_loss)
        train_accuracies.append(train_accuracy)
        
        # Print progress for certain epochs if verbose
        if verbose and (epoch + 1) % 25 == 0:
            print(f'Epoch {epoch+1}/{num_epochs}:')
            print(f'Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}')
    
    return train_losses, train_accuracies

# Function to evaluate the model
def evaluate_model(model, test_loader):
    model.eval()
    test_loss = 0.0
    all_preds = []
    all_labels = []
    
    criterion = nn.CrossEntropyLoss()
    
    with torch.no_grad():
        for inputs, labels in test_loader:
            # Forward pass
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            # Calculate loss
            test_loss += loss.item() * inputs.size(0)
            
            # Get predictions
            _, predicted = torch.max(outputs, 1)
            
            # Store predictions and labels
            all_preds.extend(predicted.numpy())
            all_labels.extend(labels.numpy())
    
    # Calculate average test loss
    test_loss = test_loss / len(test_loader.dataset)
    
    # Convert to numpy arrays
    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)
    
    # Calculate metrics
    accuracy = accuracy_score(all_labels, all_preds)
    
    return test_loss, accuracy, all_preds, all_labels

# Function to plot training history
def plot_training_history(losses, accuracies, title_prefix=""):
    # Create figure with two subplots
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # If we have multiple folds, plot the mean and std dev across folds
    if isinstance(losses[0], list) or isinstance(losses[0], np.ndarray):
        # Convert to numpy arrays
        losses = np.array(losses)
        accuracies = np.array(accuracies)
        
        # Calculate mean and std dev
        mean_losses = np.mean(losses, axis=0)
        std_losses = np.std(losses, axis=0)
        mean_accuracies = np.mean(accuracies, axis=0)
        std_accuracies = np.std(accuracies, axis=0)
        
        # Plot means
        epochs = np.arange(1, len(mean_losses) + 1)
        ax1.plot(epochs, mean_losses)
        ax1.fill_between(epochs, mean_losses - std_losses, mean_losses + std_losses, alpha=0.3)
        ax2.plot(epochs, mean_accuracies)
        ax2.fill_between(epochs, mean_accuracies - std_accuracies, mean_accuracies + std_accuracies, alpha=0.3)
        
        ax1.set_title(f'{title_prefix}Mean Loss across Folds')
        ax2.set_title(f'{title_prefix}Mean Accuracy across Folds')
    else:
        # Plot single fold history
        ax1.plot(losses)
        ax2.plot(accuracies)
        
        ax1.set_title(f'{title_prefix}Loss over Epochs')
        ax2.set_title(f'{title_prefix}Accuracy over Epochs')
    
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Loss')
    
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Accuracy')
    
    # Adjust layout and save
    plt.tight_layout()
    plt.savefig(f'{title_prefix.lower().replace(" ", "_")}training_history.png')
    plt.close()

# Function to plot confusion matrix
def plot_confusion_matrix(conf_matrix, class_names, title_prefix=""):
    plt.figure(figsize=(10, 8))
    sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues',
                xticklabels=class_names, yticklabels=class_names)
    plt.title(f'{title_prefix}Confusion Matrix')
    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')
    plt.tight_layout()
    plt.savefig(f'{title_prefix.lower().replace(" ", "_")}confusion_matrix.png')
    plt.close()

# Main function to run the pipeline with k-fold cross-validation
def main():
    # Step 1: Fetch the UCI landmine dataset
    print("Fetching the UCI landmine dataset (ID: 763)...")
    try:
        # Fetch the dataset using ucimlrepo
        land_mines = fetch_ucirepo(id=763)
        
        # Extract features and targets
        X = land_mines.data.features
        y = land_mines.data.targets
        
        print("UCI landmine dataset successfully fetched.")
        print(f"Dataset shape: {X.shape}")
        
        # Check if dataset is valid
        if X is None or y is None or X.shape[0] == 0:
            raise ValueError("The fetched dataset is empty or invalid.")
            
    except Exception as e:
        print(f"Error fetching UCI dataset: {e}")
        print("Creating a synthetic multi-class dataset for demonstration instead...")
        
        # Create a synthetic multi-class dataset
        n_samples = 1000
        n_features = 10
        n_classes = 4  # For example: 4 different types of mines
        
        # Generate random features
        X = np.random.randn(n_samples, n_features)
        
        # Generate multi-class targets
        y = np.random.randint(0, n_classes, size=n_samples)
    
    # Convert y to numpy array if it's a DataFrame or Series
    if hasattr(y, 'values'):
        y = y.values
    
    # Make sure y is flattened
    y = y.flatten() if hasattr(y, 'ndim') and y.ndim > 1 else y
    
    # Encode class labels to integers starting from 0
    label_encoder = LabelEncoder()
    y_encoded = label_encoder.fit_transform(y)
    
    # Get class names (original class labels)
    class_names = label_encoder.classes_
    num_classes = len(class_names)
    
    # Print dataset information
    print(f"Dataset shape: {X.shape}")
    print(f"Number of classes: {num_classes}")
    print(f"Class labels: {class_names}")
    print(f"Class distribution: {np.bincount(y_encoded)}")
    
    # Set up k-fold cross-validation
    n_folds = 10
    kfold = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=42)
    
    # Standardize all features at once
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    # Lists to store fold results
    fold_train_losses = []
    fold_train_accuracies = []
    fold_test_losses = []
    fold_test_accuracies = []
    all_conf_matrices = []
    
    # Overall predictions and true labels for final evaluation
    all_predictions = []
    all_true_labels = []
    
    # Track start time
    start_time = time.time()
    
    # Perform k-fold cross-validation
    for fold, (train_idx, test_idx) in enumerate(kfold.split(X_scaled, y_encoded)):
        print(f"\n{'='*20} Fold {fold+1}/{n_folds} {'='*20}")
        
        # Split data
        X_train, X_test = X_scaled[train_idx], X_scaled[test_idx]
        y_train, y_test = y_encoded[train_idx], y_encoded[test_idx]
        
        # Create DataLoaders
        train_dataset = LandmineDataset(X_train, y_train)
        test_dataset = LandmineDataset(X_test, y_test)
        
        batch_size = 32
        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        test_loader = DataLoader(test_dataset, batch_size=batch_size)
        
        # Initialize model
        input_size = X.shape[1]  # Number of features
        model = MultiClassLandmineDetector(input_size=input_size, num_classes=num_classes)
        
        # Define loss function and optimizer
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=0.001)
        
        # Train the model
        print(f"Training model for fold {fold+1}...")
        train_losses, train_accuracies = train_model(
            model, train_loader, criterion, optimizer, num_epochs=500, verbose=False
        )
        
        # Evaluate on test set
        print(f"Evaluating model for fold {fold+1}...")
        test_loss, test_accuracy, fold_preds, fold_labels = evaluate_model(model, test_loader)
        
        # Store fold results
        fold_train_losses.append(train_losses)
        fold_train_accuracies.append(train_accuracies)
        fold_test_losses.append(test_loss)
        fold_test_accuracies.append(test_accuracy)
        
        # Store predictions and true labels for this fold
        all_predictions.extend(fold_preds)
        all_true_labels.extend(fold_labels)
        
        # Calculate confusion matrix for this fold
        fold_conf_matrix = confusion_matrix(fold_labels, fold_preds)
        all_conf_matrices.append(fold_conf_matrix)
        
        # Print fold results
        print(f"Fold {fold+1} Results:")
        print(f"  Train Loss (final epoch): {train_losses[-1]:.4f}")
        print(f"  Train Accuracy (final epoch): {train_accuracies[-1]:.4f}")
        print(f"  Test Loss: {test_loss:.4f}")
        print(f"  Test Accuracy: {test_accuracy:.4f}")
        
        # Optional: Save the model for each fold
        torch.save(model.state_dict(), f'landmine_detector_fold_{fold+1}.pth')
    
    # Calculate execution time
    execution_time = time.time() - start_time
    
    # Calculate average losses and accuracies across folds
    avg_train_loss = np.mean([losses[-1] for losses in fold_train_losses])
    avg_train_acc = np.mean([accs[-1] for accs in fold_train_accuracies])
    avg_test_loss = np.mean(fold_test_losses)
    avg_test_acc = np.mean(fold_test_accuracies)
    std_test_acc = np.std(fold_test_accuracies)
    
    # Calculate overall confusion matrix and classification report
    overall_conf_matrix = confusion_matrix(all_true_labels, all_predictions)
    report = classification_report(all_true_labels, all_predictions, 
                                  target_names=[str(name) for name in class_names])
    
    # Print overall results
    print("\n" + "="*50)
    print(f"10-Fold Cross-Validation Results:")
    print(f"Total execution time: {execution_time:.2f} seconds")
    print(f"Average Training Loss: {avg_train_loss:.4f}")
    print(f"Average Training Accuracy: {avg_train_acc:.4f}")
    print(f"Average Test Loss: {avg_test_loss:.4f}")
    print(f"Average Test Accuracy: {avg_test_acc:.4f} ± {std_test_acc:.4f}")
    print("\nClassification Report:")
    print(report)
    
    # Plot training history (average across folds)
    plot_training_history(fold_train_losses, fold_train_accuracies, title_prefix="10-Fold CV ")
    
    # Plot overall confusion matrix
    plot_confusion_matrix(overall_conf_matrix, 
                         [str(name) for name in class_names], 
                         title_prefix="10-Fold CV ")
    
    # Print accuracy for each fold
    print("\nTest Accuracy for each fold:")
    for fold, acc in enumerate(fold_test_accuracies):
        print(f"Fold {fold+1}: {acc:.4f}")
    
    print("\nModel training and evaluation complete!")

if __name__ == "__main__":
    main()