diff --git a/recognition/Alzheimers_Classifier_s4693608/README.md b/recognition/Alzheimers_Classifier_s4693608/README.md new file mode 100644 index 000000000..83fcdc64f --- /dev/null +++ b/recognition/Alzheimers_Classifier_s4693608/README.md @@ -0,0 +1,111 @@ +# Alzheimer’s Disease MRI Classifier — COMP3710 Task 8 + +This project implements a deep learning classifier to distinguish between Alzheimer’s Disease (AD) and Normal Control (NC) using 2D MRI brain slices from the ADNI dataset. +The model is based on ConvNeXt, trained using PyTorch, and makes predictions on the slice-level and then aggregates through these predictions and averages them to make patient-level predictions. +The model was trained and tested on UQ's rangpur, achieving a final patient prediction accuracy of 80.22%. + +Developed for COMP3710 — Pattern Recognition and Analysis, University of Queensland by Christian Vever - s4693608. + +## Dataset Structure + +/home/groups/comp3710/ADNI/ +├── AD_NC/ +│ ├── train/ +│ │ ├── AD/ +│ │ │ ├── `_.jpeg` +│ │ └── NC/ +│ │ ├── `_.jpeg` +│ ├── test/ +│ │ ├── AD/ +│ │ └── NC/ +└── meta_data_with_label.json + +All images are greyscale JPEGs of size (256 x 240), and are 2D slices of MRI brain scans form the ADNI dataset. +Filenames follow `_.jpeg` naming convention. + +## Requirements + +### System Requirments + +* Python 3.10 + +* Parallel processing using GPU (optional but recommended) + +### Dependency Requirements + +* torch +* torchvision +* timm +* pillow +* numpy + +## Training Model Parameters + +* Model: ConvNeXt (pretrained on ImageNet) +* Epochs: 30 +* Batch Size: 32 (optimal tested number of batches) +* Learning Rate: 1e-4 +* Scheduler: OneCycleLR (allows for increased learnin rate over first few epochs before decreasing) +* Dropout: Enabled in classifier head (set to 0.5; prevents overfitting) +* Loss: Weighted CrossEntropyLoss (weigth = [2.0, 1.0] to focus more on AD cases) +* Augmentation: Random rotation +/- 10 degrees and horizontal flip (to force the classifier to focus more on general features) +* Normalisation: mean = 0.1159, std = 0.2199 (the calculated mean and std of the dataset) + +The script saves the best checkpoint to `best_model.pth`. + +## Running Scripts and Outputs + +Scripts were run on rangpur using sbatcb. +sbatch runner for train.py: +``` +#!/bin/bash +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=1 +#SBATCH --gres=gpu:1 +#SBATCH --partition=a100 +#SBATCH --job-name=train +#SBATCH --time=04:00:00 +#SBATCH -o train.out +#SBATCH -e train.err + +conda activate torch +python train.py +``` + +The train.py script outputs (per epoch): +* The current epoch number (out of total epochs) +* The training loss for this epoch +* The validation loss for this epoch +* The validation accuracy for this epoch +* The current learning rate +Example output: +``` +Epoch 7/30 | Train Loss: 0.0914 | Val Loss: 0.6345 | Val Acc: 0.7551 | LR: 0.000447 +``` +Once all epochs have been run, the script outputs the best validation accuracy. + +sbatch runner for predict.py: +``` +#!/bin/bash +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=1 +#SBATCH --gres=gpu:a100 +#SBATCH --job-name=predict +#SBATCH -o predict.out +#SBATCH -e predict.err + +conda activate torch +python predict.py +``` + +The predict.py script outputs: +* The patiend id +* The predicted case (AD or NC) +* The true case (AD or NC) +* The patient prediction accuracy across all slices +Example output: +``` +Patient 389298: Predicted AD (99.99%) | True: AD +``` +Once all patient id's have been tested, the script outputs the overall patient prediction accuracy. \ No newline at end of file diff --git a/recognition/Alzheimers_Classifier_s4693608/dataset.py b/recognition/Alzheimers_Classifier_s4693608/dataset.py new file mode 100644 index 000000000..f68a6256a --- /dev/null +++ b/recognition/Alzheimers_Classifier_s4693608/dataset.py @@ -0,0 +1,85 @@ +import os +import torch +from torch.utils.data import Dataset, DataLoader +from PIL import Image +import torchvision.transforms as transforms + +class ADNIDataset(Dataset): + """ + A PyTorch Dataset for loading 2D MRI image slices from the ADNI dataset, + organized into Alzheimer's disease (AD) and normal control (NC) categories. + + Args: + root_dir (str): Path to the AD_NC directory containing "train" and "test". + split (str): Which dataset split to use, "train" or "test". + transform (callable, optional): Optional transform to be applied on an image. + + Attributes: + samples (list): List of tuples (image_path, label) for all samples + in the specified split. + """ + + def __init__(self, root_dir, split="train", transform=None): + self.root_dir = os.path.join(root_dir, split) # path automatically goes to train folder + self.transform = transform + self.samples = [] # (image_path, label) + + # get all (image_path, label) pairs + for label_name, label in [("AD", 1), ("NC", 0)]: + class_dir = os.path.join(self.root_dir, label_name) + for fname in os.listdir(class_dir): + self.samples.append((os.path.join(class_dir, fname), label)) + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx): + img_path, label = self.samples[idx] + image = Image.open(img_path) + + if self.transform: + image = self.transform(image) + + return image, torch.tensor(label, dtype=torch.long) + +def get_dataloaders(root_dir, batch_size=16): + """ + Create PyTorch DataLoaders for the ADNI dataset. + + This function initializes ADNIDataset instances for the training and + testing splits, applies preprocessing transforms (resize, tensor conversion, + normalisation), and returns DataLoaders for batched access. + + Args: + root_dir (str): Path to the AD_NC directory containing 'train' and 'test'. + batch_size (int, optional): Number of samples per batch. Default is 16. + + Returns: + tuple: + - train_loader (DataLoader): DataLoader for the training set, + with shuffling enabled. + - test_loader (DataLoader): DataLoader for the test set, + with shuffling disabled. + """ + train_transform = transforms.Compose([ + transforms.Resize((224, 224)), + transforms.RandomHorizontalFlip(p=0.5), + transforms.RandomRotation(degrees=10), + transforms.RandomResizedCrop(224, scale=(0.8, 1.0)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.1159], std=[0.2199]) + ]) + + test_transform = transforms.Compose([ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.1159], std=[0.2199]) + ]) + + train_dataset = ADNIDataset(os.path.join(root_dir, "AD_NC"), split="train", transform=train_transform) + test_dataset = ADNIDataset(os.path.join(root_dir, "AD_NC"), split="test", transform=test_transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + return train_loader, test_loader \ No newline at end of file diff --git a/recognition/Alzheimers_Classifier_s4693608/modules.py b/recognition/Alzheimers_Classifier_s4693608/modules.py new file mode 100644 index 000000000..017a90055 --- /dev/null +++ b/recognition/Alzheimers_Classifier_s4693608/modules.py @@ -0,0 +1,45 @@ +import torch +import torch.nn as nn +import timm + +class AlzheimersClassifier(nn.Module): + """ + ConvNeXt-based classifier for Alzheimer's (AD) vs Normal Control (NC). + + This class wraps a pretrained ConvNeXt model from the `timm` library and modifies it to: + - Accept grayscale input images (1 channel) by duplicating them into 3 channels, + since ConvNeXt expects RGB input. + - Replace the final classification head with a dropout rate and linear layer outputting 2 classes + (Alzheimer's disease = 1, Normal Control = 0). + + Attributes: + model : timm.models.ConvNeXt + The ConvNeXt backbone model with a modified classifier head. + + Methods: + forward(x: torch.Tensor) -> torch.Tensor + Performs a forward pass through the network. + Takes grayscale input of shape [B, 1, H, W], duplicates channels, + and outputs logits of shape [B, 2]. + """ + def __init__(self, model_name="convnext_base", num_classes=2, pretrained=True, dropout=0.5): + super().__init__() + # Load pretrained ConvNeXt backbone + self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=0) + + # Ensure fixed size feature vector + self.pool = nn.AdaptiveAvgPool2d(1) + + # Replace classifier with dropout + linear + self.dropout = nn.Dropout(dropout) + self.fc = nn.Linear(self.model.num_features, num_classes) + + def forward(self, x): + # duplicate x channels ([B, 1, H, W] -> [B, 3, H, W]) + x = x.repeat(1, 3, 1, 1) + x = self.model.forward_features(x) + x = self.pool(x) + x = torch.flatten(x, 1) + x = self.dropout(x) + x = self.fc(x) + return x \ No newline at end of file diff --git a/recognition/Alzheimers_Classifier_s4693608/predict.py b/recognition/Alzheimers_Classifier_s4693608/predict.py new file mode 100644 index 000000000..957cc2c5b --- /dev/null +++ b/recognition/Alzheimers_Classifier_s4693608/predict.py @@ -0,0 +1,111 @@ +import os +import torch +from torchvision import transforms +from PIL import Image +from modules import AlzheimersClassifier +from collections import defaultdict +import numpy as np + +def load_model(model_path="best_model.pth"): + """ + This function initializes an instance of the AlzheimersClassifier model, + loads the trained weights from the specified checkpoint file, and sets + the model to evaluation mode on the appropriate device (CPU or GPU). + + Args: + model_path (str, optional): Path to the saved model checkpoint (.pth file). + Defaults to the value of "best_model.pth". + + Returns: + AlzheimersClassifier: + A ConvNeXt-based PyTorch model ready for inference. + """ + model = AlzheimersClassifier() + model.load_state_dict(torch.load(model_path, map_location="cuda")) + model.to("cuda") + model.eval() + return model + +def predict_slice(image_path, model): + """ + The function loads a grayscale MRI slice, applies the same preprocessing + transformations used during training (resize, tensor conversion, normalization), + and performs a forward pass through the trained model to obtain the predicted + probability of AD/NC case. + + Args: + image_path (str): Path to the input image (e.g., .jpeg slice). + model (torch.nn.Module): Trained AlzheimersClassifier model instance. + + Returns: + probs (float): probability of AD/NC case. + """ + image = Image.open(image_path) + + # Apply transform to image + transform = transforms.Compose([ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.1159], std=[0.2199]) + ]) + image = transform(image).unsqueeze(0).to("cuda") + + with torch.no_grad(): + outputs = model(image) + probs = torch.softmax(outputs, dim=1)[0].cpu().numpy() + + return probs + +def aggregate_patient_predictions(patient_probs): + """ + Given a list of [NC_prob, AD_prob] arrays for one patient, + average them and return predicted label + confidence. + + Args: + patient_probs (list): probabilities of AD case for all slices of one patient. + + Returns: + tuple: + - label (str): Predicted class label value (1 for "AD" or 0 for "NC"). + - confidence (float): Model confidence for the predicted class, + between 0.0 and 1.0. + - mean_probs (float): mean probability of AD case across aggregated slices + for a patient. + """ + mean_probs = np.mean(patient_probs, axis=0) + label = np.argmax(mean_probs) + confidence = mean_probs[label] + + return label, confidence, mean_probs + +if __name__ == "__main__": + model = load_model() + root_dir = "/home/groups/comp3710/ADNI/AD_NC/test" + class_names = ["NC", "AD"] + + # Group all slices by patient id + patient_slices = defaultdict(list) + for cls in ["AD", "NC"]: + folder = os.path.join(root_dir, cls) + for fname in os.listdir(folder): + patient_id = fname.split('_')[0] + patient_slices[patient_id].append(os.path.join(folder, fname)) + + # Predict each slice and aggregate + patient_results = {} + for pid, slice_paths in patient_slices.items(): + slice_probs = [predict_slice(p, model) for p in slice_paths] + label, conf, mean_probs = aggregate_patient_predictions(slice_probs) + patient_results[pid] = (label, conf, mean_probs) + + # Evaluate patient accuracy + correct, total = 0, 0 + for pid, (label, conf, probs) in patient_results.items(): + true_label = 1 if any("AD/" in p for p in patient_slices[pid]) else 0 + total += 1 + correct += int(label == true_label) + + print(f"Patient {pid}: Predicted {class_names[label]} ({conf * 100:.2f}%)" + f" | True: {class_names[true_label]}") + + print(f"\nPatient Prediction Accuracy: {100 * correct / total:.2f}%") \ No newline at end of file diff --git a/recognition/Alzheimers_Classifier_s4693608/train.py b/recognition/Alzheimers_Classifier_s4693608/train.py new file mode 100644 index 000000000..5a1162252 --- /dev/null +++ b/recognition/Alzheimers_Classifier_s4693608/train.py @@ -0,0 +1,115 @@ +import os +import torch +import torch.nn as nn +import torch.optim as optim +from torch.optim.lr_scheduler import OneCycleLR +from dataset import get_dataloaders +from modules import AlzheimersClassifier + +def train_model(root_dir, epochs=10, batch_size=16, lr=1e-4, device='cuda'): + """ + Train and evaluate a deep learning model for Alzheimer's disease classification. + + This function performs supervised training of a model using the given training + data, evaluates it on the test data at each epoch, and reports progress. + + Args: + root_dir (str): Path to the ADNI dataset root directory. + epochs (int, optional): Number of training epochs. Default is 10. + batch_size (int, optional): Number of samples per training batch. Default is 16. + lr (float, optional): Learning rate for the optimiser. Default is 1e-4. + device (str, optional): Device to run training on ('cuda' or 'cpu'). Default is "cuda". + + Returns: + None + + Attributes: + - Saves the best model's state_dict to "best_model.pth" in the current directory. + - Generates and saves two plots: + - "loss_curve.png" for training and validation loss + - "val_acc_curve.png" for validation accuracy + - Prints training/validation progress and best validation accuracy to the console. + """ + train_loader, test_loader = get_dataloaders(root_dir, batch_size=batch_size) + + # Define model, loss, optimiser + model = AlzheimersClassifier().to(device) + criterion = nn.CrossEntropyLoss(weight=torch.tensor([2.0, 1.0]).to(device)) + optimiser = optim.AdamW(model.parameters(), lr=lr) + + # scheduler + scheduler = OneCycleLR( + optimiser, + max_lr=lr * 5, + steps_per_epoch=len(train_loader), + epochs=epochs, + pct_start=0.3, + anneal_strategy='cos', + div_factor=10, + final_div_factor=1e4 + ) + + train_losses, val_losses, val_accs = [], [], [] + best_acc = 0.0 + + for epoch in range(epochs): + # Training + model.train() + running_loss = 0.0 + for images, labels in train_loader: + images, labels = images.to(device), labels.to(device) + + optimiser.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimiser.step() + scheduler.step() + + running_loss += loss.item() + + train_loss = running_loss / len(train_loader) + train_losses.append(train_loss) + + # Validation + model.eval() + correct, total, val_loss = 0, 0, 0.0 + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + loss = criterion(outputs, labels) + val_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + + val_acc = correct / total + val_loss /= len(test_loader) + + val_losses.append(val_loss) + val_accs.append(val_acc) + + print(f"Epoch {epoch+1}/{epochs} | " + f"Train Loss: {train_loss:.4f} | " + f"Val Loss: {val_loss:.4f} | " + f"Val Acc: {val_acc:.4f} | " + f"LR: {current_lr:.6f}") + + # Update scheduler + scheduler.step(val_acc) + current_lr = optimiser.param_groups[0]['lr'] + print(f"Learning rate: {current_lr:.6f}") + + # Save the best model + if val_acc > best_acc: + best_acc = val_acc + torch.save(model.state_dict(), "best_model.pth") + print("New best model saved!") + + print(f"Best Validation Accuracy: {best_acc:.4f}") + +if __name__ == "__main__": + root_dir = "/home/groups/comp3710/ADNI" + train_model(root_dir, epochs=30, batch_size=32, lr=1e-4) \ No newline at end of file