train.py

"""
Project: Prostate Cancer Classification using Vision Transformers (ViT) presented at the XXXVIII Conference on Graphics, Patterns and Images- Workshop on Digital and Computational Pathology (WDCPath).

Authors:Betania Caroline Silva de Albuquerque
        Federal University of Sao Paulo
        Email: betania.caroline@unifesp.br
        
        Leandro Alves Neves
        Sao Paulo State University
        Email: nevesleandro@gmail.com
         
        
        Hanna Beatriz Couto Monteiro Fernandes de Castro
        Federal University of Sao Paulo
        Email: h.castro01@unifesp.br
        
        Marcelo Zanchetta do Nascimento
        Federal University of Uberlandia
        Email: marcelozanchetta@gmail.com
        
        Thaına Aparecida Azevedo Tosta
        Federal University of Sao Paulo
        Email: tosta.thaina@gmail.com
        
Description: Pre-trained ViT google/vit-large-patch16-224 training, feature extraction,
             classification and ensemble using ViT classification, SVM, SVM-polynomial, K-NN, RF, LR and NB.
"""

#%%Import Libraries
#Standard Library Imports
import os
import time
import psutil

# Third-Party Library Imports (Numerical & Data)
import numpy as np
import pandas as pd
import torch
from scipy.special import softmax
from scipy.stats import mode

#Machine Learning & Computer Vision (Scikit-Learn)
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    roc_auc_score, 
    accuracy_score, 
    f1_score, 
    confusion_matrix, 
    ConfusionMatrixDisplay
)

# Deep Learning & Transformers (Hugging Face)
from torch.utils.data import DataLoader
from transformers import (
    logging,
    ViTModel,
    ViTImageProcessor,
    ViTForImageClassification,
    TrainingArguments,
    Trainer
)

#Visualization Imports
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator

# Local Module Imports
import dataset

#%% Defining Hyperparameters

#diretory to save the plots
diretorio = './saved_plots/' 

batch_size = 16
epoch = 50
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model_name = "google/vit-large-patch16-224"

train_args = TrainingArguments(
    output_dir="output_dir",
    save_total_limit=2,
    report_to=None,
    save_strategy="epoch",
    eval_strategy="epoch",
    learning_rate=2e-4,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=epoch,
    weight_decay=0.01, 
    load_best_model_at_end=True,
    logging_dir="logs",
    remove_unused_columns=False,
    fp16=True,
    logging_steps=32, 
    dataloader_num_workers=4
)


#%%Functions
# Metrics

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=1)
    accuracy = accuracy_score(labels, predictions)
    return {"accuracy": accuracy} 


def extract_features(data_dl, model_extractor, device):
    all_features = []
    all_labels = []
    with torch.no_grad():
        for batch in data_dl:
            pixel_values = batch["pixel_values"].to(device)
            labels = batch["labels"].to(device)
            outputs = model_extractor(pixel_values)
            # CLS token 
            features = outputs.last_hidden_state[:, 0, :].cpu().numpy()
            all_features.append(features)
            all_labels.append(labels.cpu().numpy())
    return np.concatenate(all_features, axis=0), np.concatenate(all_labels, axis=0)

def evaluate_classifier(name, model, features, labels):
    print(f"\n--- Training and evaluating: {name} ---")
    train_features_to_use = train_features_scaled if name in ["SVM", "Logistic Regression", "Polynomial"] else train_features
    
    start_time = time.time()
    model.fit(train_features_to_use, train_labels)
    training_time = time.time() - start_time
    
    eval_features_to_use = features
    if hasattr(model, "predict_proba"):
        probabilities = model.predict_proba(eval_features_to_use)
        predicted_classes = np.argmax(probabilities, axis=1)
        auc_val = roc_auc_score(labels, probabilities[:, 1])
    else:
        predicted_classes = model.predict(eval_features_to_use)
        auc_val = "N/A"

    acc = accuracy_score(labels, predicted_classes)
    tn, fp, fn, tp = confusion_matrix(labels, predicted_classes).ravel()
    sens = tp / (tp + fn) if (tp + fn) != 0 else 0.0
    spec = tn / (tn + fp) if (tn + fp) != 0 else 0.0
    f1_val = f1_score(labels, predicted_classes)

    print(f"AUC: {auc_val}")
    print(f"Accuracy: {acc:.4f}")
    return {
        "name": name, "model": model, "predictions": predicted_classes,
        "probabilities": probabilities[:, 1] if hasattr(model, "predict_proba") else None,
        "auc": auc_val, "accuracy": acc, "sensitivity": sens, "specificity": spec, "f1": f1_val,
        "training_time": training_time
    }

#%% Check free RAM. For faster training we uploaded the dataset to the RAM memory. If you have no RAM memory, you may need to change that.

mem = psutil.virtual_memory()

print(f"Total: {mem.total / 1e9:.2f} GB")
print(f"Available: {mem.available / 1e9:.2f} GB")
print(f"In use: {mem.used / 1e9:.2f} GB")
print(f"Used percentage: {mem.percent}%")

#%%

logging.set_verbosity_info()

os.environ["WANDB_DISABLED"] = "true" # To disable notifications for wandb
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" # Avoid shared library error

#%% Load data into RAM
train_ds = dataset.InMemoryTensorDataset(
    caminho_csv="C:/Users/lipai/Datasets/Cancer_de_Prostata1/Dados_Originais/train/train_com_aug.csv",
    pasta_tensores="C:/Users/lipai/Datasets/Cancer_de_Prostata1/Dados_Normalizados_BKSVD/train_processed"
)

val_ds = dataset.InMemoryTensorDataset(
    caminho_csv="C:/Users/lipai/Datasets/Cancer_de_Prostata1/Dados_Originais/val/val.csv",
    pasta_tensores="C:/Users/lipai/Datasets/Cancer_de_Prostata1/Dados_Normalizados_BKSVD/val_processed"
)

test_ds = dataset.CustomTensorDataset(
    caminho_csv="C:/Users/lipai/Datasets/Cancer_de_Prostata1/Dados_Originais/test/test.csv",
    pasta_tensores="C:/Users/lipai/Datasets/Cancer_de_Prostata1/Dados_Normalizados_BKSVD/test_processed"
)

#%% Loading data in batches

train_dl = DataLoader(train_ds, collate_fn=dataset.collate_fn, batch_size=batch_size, num_workers=4)
test_dl = DataLoader(test_ds, collate_fn=dataset.collate_fn, batch_size=batch_size, num_workers=4)
val_dl = DataLoader(val_ds, collate_fn=dataset.collate_fn, batch_size=batch_size, num_workers=4)


#%% Define labels
id2label = {0: "benign", 1: "malignant"}
label2id = {"benign": 0, "malignant": 1}

#%% Importing pre-trained model
processor = ViTImageProcessor.from_pretrained(model_name)

# Setting ignore_mismatched_sizes to True allows the model to adjust its layers to accommodate size differences without throwing an error.
model = ViTForImageClassification.from_pretrained(
    model_name, id2label=id2label, label2id=label2id, ignore_mismatched_sizes=True
)

#%% Training
trainer = Trainer(
    model,
    train_args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    data_collator=dataset.collate_fn,
    processing_class=processor,
    compute_metrics=compute_metrics,
)
trainer.train(resume_from_checkpoint=False)

best_model_path= trainer.state.best_model_checkpoint

if best_model_path:
    print(f"O melhor modelo foi salvo em: {best_model_path}")
    model_path = best_model_path 
else:
    model_path = "output_dir" 
    print("Aviso: Melhor checkpoint não encontrado, usando diretório padrão.")
#%% Evaluation on validation set

# evaluate the method
metrics = trainer.evaluate()
print(metrics)

# get history
history = trainer.state.log_history

# Initialize lists for metrics per epoch
eval_losses_per_epoch = []
eval_accuracies_per_epoch = []
train_losses_per_epoch = []
epochs_list = [] # To store actual epoch numbers

# Filter history to collect data from each epoch
for log_entry in history:

    if "eval_loss" in log_entry and "eval_accuracy" in log_entry:
        eval_losses_per_epoch.append(log_entry["eval_loss"])
        eval_accuracies_per_epoch.append(log_entry["eval_accuracy"])
        epochs_list.append(log_entry["epoch"])

#collect train_loss per epoch ---
processed_train_losses = {}
for log_entry in history:
    if "loss" in log_entry and "epoch" in log_entry and log_entry.get("eval_loss") is None:
        epoch_num = int(round(log_entry["epoch"]))
        if epoch_num not in processed_train_losses:
            processed_train_losses[epoch_num] = []
        processed_train_losses[epoch_num].append(log_entry["loss"])

# Calculate mean training loss for each epoch
final_train_losses = []
final_train_epochs = []
for epoch_num in sorted(processed_train_losses.keys()):
    if epoch_num > 0: 
        final_train_epochs.append(epoch_num)
        final_train_losses.append(np.mean(processed_train_losses[epoch_num]))


if not os.path.exists(diretorio):
    os.makedirs(diretorio)
    print(f"Directory '{diretorio}' created.")
else:
    print(f"Directory '{diretorio}' already exists.")

#%% Plot the graphs
if not epochs_list:
    print("No evaluation data (eval_loss, eval_accuracy) found. Check compute_metrics and training.")
elif not final_train_epochs:
    print("No training data (train_loss) found. Check logging_strategy.")
else:
    #Training and Validation Loss
    plt.figure(figsize=(7, 7))
    ax1 = plt.gca()
    ax1.plot(final_train_epochs, final_train_losses, label="Train Loss", marker='o', color='blue')
    ax1.plot(epochs_list, eval_losses_per_epoch, label="Validation Loss", marker='o', color='red')
    ax1.set_xlabel("Epochs")
    ax1.set_ylabel("Loss")
    ax1.set_title("Training and Validation Loss Over the Epochs", y=1.05)
    ax1.legend()
    ax1.grid(True)
    ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
    plt.savefig(os.path.join(diretorio, 'training_validation_loss.png'))
    plt.show()

    #Validation Accuracy and Loss
    plt.figure(figsize=(7, 7))
    ax2 = plt.gca()
    ax2.plot(epochs_list, eval_accuracies_per_epoch, label="Validation Accuracy", linestyle='--', marker='o', color='green')
    ax2.plot(epochs_list, eval_losses_per_epoch, label="Validation Loss", marker='o', color='red')
    ax2.set_xlabel("Epochs")
    ax2.set_ylabel("Accuracy and Loss")
    ax2.set_title("Validation Accuracy and Loss Over the Epochs", y=1.05)
    ax2.legend()
    ax2.grid(True)
    ax2.xaxis.set_major_locator(MaxNLocator(integer=True))
    plt.savefig(os.path.join(diretorio, 'validation_accuracy_loss.png'))
    plt.show()

    #Validation Accuracy
    plt.figure(figsize=(7, 7))
    ax3 = plt.gca()
    ax3.plot(epochs_list, eval_accuracies_per_epoch, label="Validation Accuracy", linestyle='--', marker='o', color='green')
    ax3.set_xlabel("Epochs")
    ax3.set_ylabel("Accuracy")
    ax3.set_title("Validation Accuracy Over the Epochs", y=1.05)
    ax3.legend()
    ax3.grid(True)
    ax3.xaxis.set_major_locator(MaxNLocator(integer=True))
    plt.savefig(os.path.join(diretorio, 'validation_accuracy.png'))
    plt.show()

#%% Predict on test set
outputs = trainer.predict(test_ds)
print(outputs.metrics)

#%% Confusion Matrix

y_true = outputs.label_ids
y_pred = outputs.predictions.argmax(1)

labels = ["benign", "malignant"]
cm = confusion_matrix(y_true, y_pred)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=labels)
disp.plot(xticks_rotation=45)
plt.savefig("confusion_matrix.png", bbox_inches='tight')

#%% Calculate Metrics
# Perform predictions on test set
predictions = outputs
logits = predictions.predictions
labels = predictions.label_ids

# Apply softmax to get probabilities
probabilities = softmax(logits, axis=1)
predicted_classes = np.argmax(probabilities, axis=1)

# Calculate metrics
auc = roc_auc_score(labels, probabilities[:, 1])
accuracy = accuracy_score(labels, predicted_classes)
tn, fp, fn, tp = confusion_matrix(labels, predicted_classes).ravel()
sensitivity = tp / (tp + fn)
specificity = tn / (tn + fp)
f1 = f1_score(labels, predicted_classes)

print("Metrics on test set:")
print(f"AUC: {auc:.4f}")
print(f"Accuracy: {accuracy:.4f}")
print(f"Sensitivity (Recall): {sensitivity:.4f}")
print(f"Specificity: {specificity:.4f}")
print(f"F1-Score: {f1:.4f}")

#%% Feature Extraction

try:
    feature_extractor_model = ViTModel.from_pretrained(model_path)
    print("ViTModel successfully loaded from checkpoint.")
except Exception as e:
    print(f"Error loading ViTModel directly: {e}")
    from transformers import ViTForImageClassification
    full_model = ViTForImageClassification.from_pretrained(
        model_path, num_labels=2, id2label=id2label, label2id=label2id
    )
    feature_extractor_model = full_model.vit
    print("ViTForImageClassification loaded and ViT base extracted.")

feature_extractor_model.eval()
feature_extractor_model.to(device)

train_features, train_labels = extract_features(train_dl, feature_extractor_model, device)
test_features, test_labels = extract_features(test_dl, feature_extractor_model, device)
val_features, val_labels = extract_features(val_dl, feature_extractor_model, device)

print("\nFeature extraction completed.")

#%% Classifiers evaluation
scaler = StandardScaler()
train_features_scaled = scaler.fit_transform(train_features)
test_features_scaled = scaler.transform(test_features)
val_features_scaled = scaler.transform(val_features)

classifiers = {
    "SVM": SVC(probability=True, random_state=42),
    "K-Nearest Neighbors": KNeighborsClassifier(n_neighbors=5),
    "Random Forest": RandomForestClassifier(random_state=42),
    "Logistic Regression": LogisticRegression(random_state=42, solver='liblinear'),
    "Naive Bayes": GaussianNB(),
    "Polynomial": SVC(kernel='poly', probability=True, random_state=42)
}

results = []
trained_models = {}
for name, model_obj in classifiers.items():
    res = evaluate_classifier(name, model_obj, test_features_scaled if name in ["SVM", "Logistic Regression", "Polynomial", "K-Nearest Neighbors"] else test_features, test_labels)
    results.append(res)
    trained_models[name] = res["model"]

#%% Majority Vote Ensemble
all_individual_probabilities = []
all_individual_predictions = []

for name, model_obj in trained_models.items():
    feat = test_features_scaled if name in ["SVM", "Logistic Regression", "Polynomial", "K-Nearest Neighbors"] else test_features
    if hasattr(model_obj, "predict_proba"):
        all_individual_probabilities.append(model_obj.predict_proba(feat)[:, 1])
    else:
        all_individual_probabilities.append(None)
    all_individual_predictions.append(model_obj.predict(feat))

if all_individual_predictions:
    valid_probs = [p for p in all_individual_probabilities if p is not None]
    ensemble_auc = roc_auc_score(test_labels, np.mean(valid_probs, axis=0)) if valid_probs else "N/A"
    
    stacked_preds = np.stack(all_individual_predictions, axis=1)
    majority_vote, _ = mode(stacked_preds, axis=1, keepdims=True)
    ensemble_predicted_classes = majority_vote.flatten()

    ensemble_accuracy = accuracy_score(test_labels, ensemble_predicted_classes)
    tn, fp, fn, tp = confusion_matrix(test_labels, ensemble_predicted_classes).ravel()
    ensemble_sensitivity = tp / (tp + fn) if (tp + fn) != 0 else 0.0
    ensemble_specificity = tn / (tn + fp) if (tn + fp) != 0 else 0.0
    ensemble_f1 = f1_score(test_labels, ensemble_predicted_classes)

    print(f"Ensemble Accuracy: {ensemble_accuracy:.4f}")

#%% Fine-tuned ViT Metrics Calculation
print("\nCalculating Fine-Tuned ViT Metrics")
model_from_checkpoint = ViTForImageClassification.from_pretrained(model_path, num_labels=2, id2label=id2label, label2id=label2id)
model_from_checkpoint.eval()
model_from_checkpoint.to(device)

all_logits = []
all_labels_test = []
with torch.no_grad():
    for batch in test_dl:
        pix = batch["pixel_values"].to(device)
        lab = batch["labels"].to(device)
        all_logits.append(model_from_checkpoint(pix).logits.cpu().numpy())
        all_labels_test.append(lab.cpu().numpy())

final_logits = np.concatenate(all_logits, axis=0)
final_labels = np.concatenate(all_labels_test, axis=0)
probs = softmax(final_logits, axis=1)
pred_classes = np.argmax(probs, axis=1)

vit_metrics = {
    "auc": roc_auc_score(final_labels, probs[:, 1]),
    "accuracy": accuracy_score(final_labels, pred_classes),
    "f1": f1_score(final_labels, pred_classes)
}

#%% Consolidated Results and Plots
all_methods_metrics = {"ViT (Fine-tuned)": vit_metrics}
for r in results:
    all_methods_metrics[r["name"]] = {"auc": r["auc"], "accuracy": r["accuracy"], "f1": r["f1"]}
all_methods_metrics["Ensemble"] = {"auc": ensemble_auc, "accuracy": ensemble_accuracy, "f1": ensemble_f1}

metrics_df = pd.DataFrame.from_dict(all_methods_metrics, orient='index')
print(metrics_df)

output_graphs_dir = "comparative_plots"
os.makedirs(output_graphs_dir, exist_ok=True)

# Generate Bar Plots
for metric in ["auc", "accuracy", "f1"]:
    plt.figure(figsize=(10, 6))
    metrics_df[metric].plot(kind='bar', color='skyblue')
    plt.title(f"Comparison of {metric.upper()}")
    plt.ylabel(metric.capitalize())
    plt.tight_layout()
    plt.savefig(os.path.join(output_graphs_dir, f"compare_{metric}.png"))
    plt.show()

#%% Save to CSV
vit_training_params = {
    "model_name": model_name, "epochs": epoch, "lr": 2e-4, "batch_size": batch_size
}
all_data_rows = [{"Type": "Parameters", **vit_training_params}]
for method, m_vals in all_methods_metrics.items():
    all_data_rows.append({"Type": "Metrics", "Method": method, **m_vals})

pd.DataFrame(all_data_rows).to_csv("experiment_results_complete.csv", index=False)
print("Results saved to experiment_results_complete.csv")