<jupyter_start><jupyter_text>Using PyTorch and MONAI for MRI Alzheimer's Classification<jupyter_code>import logging
import os
import sys
from tqdm.notebook import tqdm

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

import monai
from monai.data import ImageDataset
from monai.transforms import AddChannel, Compose, SpatialCrop, Resize, ScaleIntensity, EnsureType, RandRotate90

from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix, precision_score, recall_score, accuracy_score, balanced_accuracy_score, roc_auc_score

monai.config.print_config()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
from sklearn.utils import shuffle

def generate_filenames(labels=['control', 'dementia']):
    pairs = []
    
    for label in labels:
        label_dir = os.path.join('/home/jack/Code/GitHub/Polygence/Data/OASIS/mri_data', label)
        mr_ids = os.listdir(label_dir)
        mr_ids.sort()
        
        for mr_id in tqdm(mr_ids, desc=label):
            scans = os.path.join(label_dir, mr_id)
            img_file = [file for file in os.listdir(scans) if "nii.gz" in file]
            img_path = os.path.join(scans, img_file[0])
            i = 1 if label == 'dementia' else 0
            pairs.append([img_path, i])
    
    m = len(pairs)
    filenames = []
    labels = np.zeros(m, dtype=np.int64)
    
    idx = 0
    for filename, label in pairs:
        filenames.append(filename)
        labels[idx] = label
        idx += 1
    
    filenames = np.array(filenames)
    filenames, labels = shuffle(filenames, labels)
    return filenames, labels
        

X_filenames, y_labels = generate_filenames()
# There are two labels: 0 represents control patients, and 1 represents dementia patients
for idx in range(len(X_filenames)):
    if idx < 10:
        print(y_labels[idx], X_filenames[idx])
    dementia = "dementia" in X_filenames[idx]
    label = y_labels[idx]
    if int(dementia) != label:
        print("Bad: ", X_filenames[idx], label)
X_train, X_test, y_train, y_test = train_test_split(
    X_filenames, y_labels, test_size=.25, shuffle=True)
print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)
from collections import Counter
counter_train = Counter(y_train)
counter_test = Counter(y_test)
print(counter_train, (counter_train[0] / (counter_train[0] + counter_train[1]),
                     (counter_train[1] / (counter_train[0] + counter_train[1]))))
print(counter_test, (counter_test[0] / (counter_test[0] + counter_test[1]),
                     (counter_test[1] / (counter_test[0] + counter_test[1]))))
# Validate the train and val sets
for idx in range(len(X_train)):
    dementia = "dementia" in X_train[idx]
    label = y_train[idx]
    if int(dementia) != label:
        print("Bad: ", X_train, y_train)

for idx in range(len(X_test)):
    dementia = "dementia" in X_test[idx]
    label = y_test[idx]
    if int(dementia) != label:
        print("Bad: ", X_test, y_test)
# define transforms
train_transforms = Compose([ScaleIntensity(), AddChannel(), Resize((128,128,128)), SpatialCrop(roi_start=(15, 15, 0), roi_end=(110,110,110)), RandRotate90(), EnsureType()])
val_transforms = Compose([ScaleIntensity(), AddChannel(), Resize((128,128,128)), SpatialCrop(roi_start=(15, 15, 0), roi_end=(110,110,110)), EnsureType()])
# Check the dataset
check_ds = ImageDataset(image_files=X_filenames, labels=y_labels, transform=train_transforms)
check_loader = DataLoader(check_ds, batch_size=5, num_workers=12, pin_memory=torch.cuda.is_available())
im, label = monai.utils.misc.first(check_loader)
print(type(im), im.shape, label)
fig, ax = plt.subplots(1, 3, figsize=(15,5))
ax[0].imshow(np.rot90(im[0, 0, :, :, int(110/2)], 3))
ax[1].imshow(np.rot90(im[0, 0, :, int(95/2), :]))
ax[2].imshow(im[0, 0, int(95/2), :, :])
from IPython import display as dspl
import ipywidgets as widgets

# graphs the MRI image
def show_graph(img, height, width, depth, color, patient):
    # Plots brain scans in all three dimensions
    fig, ax = plt.subplots(1, 3, figsize=(15,7))
    pat_types = ['control', 'dementia']
    print(pat_types[label[patient]])
    print(img.shape)

    ax[0].set_title("height")
    ax[1].set_title("width")
    ax[2].set_title("depth")

    colormap = 'gray' if color == 'Gray' else 'viridis'
    ax[0].imshow(img[patient, 0, height, :, :], cmap=colormap)
    ax[1].imshow(np.rot90(img[patient, 0, :, width, :]), cmap=colormap)
    ax[2].imshow(np.rot90(img[patient, 0, :, :, depth], 3), cmap=colormap)
    

def disp(pat_type, image_data):
    num_pats, _, height, width, depth  = image_data.shape
    co = widgets.ToggleButtons(options=['Gray', 'Color'], description='Color Map')
    hi = widgets.IntSlider(value=height/2, max=height-1, continous_updates=False, description='View 1')
    wi = widgets.IntSlider(value=width/2, max=width-1, continous_updates=False, description='View 2')
    de = widgets.IntSlider(value=depth/2, max=depth-1, continous_updates=False, description='View 3')
    pat = widgets.Dropdown(options=range(num_pats), value=0, description='Patient Number')
    ui = widgets.HBox([hi, wi, de, co, pat])
    
    display_output = widgets.interactive_output(show_graph, {'img':widgets.fixed(image_data),
                                                             'height':hi,
                                                             'width':wi,
                                                             'depth':de,
                                                             'color':co,
                                                             'patient':pat})
    dspl.display(ui, display_output)

disp(label[0], im)
# create a training data loader
train_ds = ImageDataset(image_files=X_train, labels=y_train, transform=train_transforms)
train_loader = DataLoader(train_ds, batch_size=5, shuffle=True, num_workers=12, pin_memory=torch.cuda.is_available())

# create a validation data loader
val_ds = ImageDataset(image_files=X_test, labels=y_test, transform=val_transforms)
val_loader = DataLoader(val_ds, batch_size=5, num_workers=12, pin_memory=torch.cuda.is_available())
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
# Create class weights
(unique, counts) = np.unique(y_labels, return_counts=True)
frequencies = np.asarray((unique, counts)).T
print(frequencies)
print(f'{frequencies[0,1]/(np.sum(frequencies[:,1]))*100:.2f}% of the data is control')

# 0 is control, 1 is dementia
control_cnt = frequencies[0,1]
dementia_cnt = frequencies[1,1]
total_cnt = control_cnt + dementia_cnt
print(control_cnt, dementia_cnt, total_cnt)

weight_for_0 = (1 / control_cnt) * (total_cnt / 2.0)
weight_for_1 = (1 / dementia_cnt) * (total_cnt / 2.0)
class_weight = {0: weight_for_0, 1: weight_for_1}
print(class_weight)

class_weights = torch.FloatTensor([weight_for_0, weight_for_1]).to(device)
print(class_weights)
model = monai.networks.nets.DenseNet121(spatial_dims=3, in_channels=1, out_channels=2).to(device)
loss_function = torch.nn.CrossEntropyLoss(weight=class_weights)
optimizer = torch.optim.Adam(model.parameters(), 1e-5)
!rm -rf runs
val_interval = 1
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
metric_values = list()
writer = SummaryWriter()
n_epochs = 50
# start a typical PyTorch training
for epoch in tqdm(range(n_epochs), desc='Epoch'):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{n_epochs}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step += 1
        inputs, labels = batch_data[0].to(device), batch_data[1].to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        
        epoch_loss += loss.item()
        epoch_len = len(train_ds) // train_loader.batch_size
        print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}", end='\r')
        writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)
        
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

    if (epoch + 1) % val_interval == 0:
        model.eval()
        with torch.no_grad():
            y_pred = torch.tensor([], dtype=torch.float32, device=device)
            y = torch.tensor([], dtype=torch.long, device=device)

            iteration = 0
            leng = len(val_loader)
            for val_data in val_loader:
                print(f'Reading validation set: {iteration}/{leng}', end='\r')
                iteration += 1

                val_images, val_labels = val_data[0].to(device), val_data[1].to(device)
                y_pred = torch.cat([y_pred, model(val_images)], dim=0)
                y = torch.cat([y, val_labels], dim=0)

            predictions = y_pred.argmax(dim=1).cpu()
            actual = y.cpu()

            # Accuracy
            acc = accuracy_score(actual, predictions)
            balanced_acc = balanced_accuracy_score(actual, predictions)
            # Check if accuracy is better
            if balanced_acc > best_metric:
                best_metric = balanced_acc
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), "best_metric_model_classification3d_dict.pth")
                print(f"Saved new best metric model: best metric of {best_metric}")
            writer.add_scalar("val_accuracy", acc, epoch + 1)
            writer.add_scalar("val_balanced_accuracy", balanced_acc, epoch + 1)

            # ROC-AUC
            auc_result = roc_auc_score(actual, predictions)

            # Precision
            precision = 0
            if 1 in predictions:
                precision = precision_score(actual, predictions)
            writer.add_scalar("val_precision", precision, epoch + 1)

            # Recall
            recall = recall_score(actual, predictions)
            writer.add_scalar("val_recall", recall, epoch + 1)
            
            print(f"epoch: {epoch+1}, accuracy: {acc:.4f}, balanced_acc: {balanced_acc:.4f}, AUC: {auc_result:.4f}, precision: {precision:.4f}, recall: {recall:.4f}")
            print(f"best epoch: {best_metric_epoch}, best accuracy: {best_metric}")

            # Confusion matrix
            cm = confusion_matrix(actual, predictions)

            fig, ax = plt.subplots()
            sns.heatmap(cm, ax=ax, square=True, annot=True, fmt='d', cbar=False, cmap='Blues')
            ax.set_xlabel('predicted label')
            ax.set_ylabel('true label', rotation=0)
            ax.set_xticklabels(['Control', 'Dementia'], rotation=0)
            ax.set_yticklabels(['Control', 'Dementia'], rotation=0)
            ax.set_title(f'Validation set confusion matrix for epoch {epoch}')
            plt.show()

print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}")
writer.close()
eval_model = monai.networks.nets.DenseNet121(spatial_dims=3, in_channels=1, out_channels=2).to(device)
eval_model.load_state_dict(torch.load("/home/jack/Code/GitHub/Polygence/CNN/best_metric_model_classification3d_dict.pth"))
eval_model.eval()
def evaluate(device, model, loader):
    model.eval()
    with torch.no_grad():
        y_pred = torch.tensor([], dtype=torch.float32, device=device)
        y = torch.tensor([], dtype=torch.long, device=device)

        iteration = 0
        leng = len(loader)
        for val_data in loader:
            print(f'Reading set: {iteration}/{leng}', end='\r')
            iteration += 1

            val_images, val_labels = val_data[0].to(device), val_data[1].to(device)
            y_pred = torch.cat([y_pred, model(val_images)], dim=0)
            y = torch.cat([y, val_labels], dim=0)

        predictions = y_pred.argmax(dim=1).cpu()
        actual = y.cpu()

        # Accuracy
        acc = accuracy_score(actual, predictions)
        balanced_acc = balanced_accuracy_score(actual, predictions)

        # ROC-AUC
        auc_result = roc_auc_score(actual, predictions)

        # Precision
        precision = 0
        if 1 in predictions:
            precision = precision_score(actual, predictions)

        # Recall
        recall = recall_score(actual, predictions)

        print(f"accuracy: {acc:.4f}, balanced_acc: {balanced_acc:.4f}, AUC: {auc_result:.4f}, precision: {precision:.4f}, recall: {recall:.4f}")

        # Confusion matrix
        cm = confusion_matrix(actual, predictions)

        fig, ax = plt.subplots()
        sns.heatmap(cm, ax=ax, square=True, annot=True, fmt='d', cbar=False, cmap='Blues')
        ax.set_xlabel('predicted label')
        ax.set_ylabel('true label', rotation=0)
        ax.set_xticklabels(['Control', 'Dementia'], rotation=0)
        ax.set_yticklabels(['Control', 'Dementia'], rotation=0)
        ax.set_title(f'Validation set confusion matrix for epoch {epoch}')
        plt.show()
evaluate(device, eval_model, val_loader)
test_train_loader = DataLoader(train_ds, batch_size=5, num_workers=12, pin_memory=torch.cuda.is_available())
evaluate(device, eval_model, test_train_loader)
(unique, counts) = np.unique(y_train, return_counts=True)
frequencies = np.asarray((unique, counts)).T
print(frequencies)
print(f'{frequencies[0,1]/(np.sum(frequencies[:,1]))*100:.2f}% of the data is control')
(unique, counts) = np.unique(y_test, return_counts=True)
frequencies = np.asarray((unique, counts)).T
print(frequencies)
print(f'{frequencies[0,1]/(np.sum(frequencies[:,1]))*100:.2f}% of the data is control')<jupyter_output>[[  0 176]
 [  1  80]]
68.75% of the data is control
{'fraction_non_alphanumeric': 0.1314231415832176, 'fraction_numerical': 0.02024705796359915, 'mean_word_length': 3.286340852130326, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '314149', 'license': 'MIT', 'n_tokens_mistral': 4724, 'n_tokens_neox': 4526, 'n_words': 942, 'path': 'CNN/monai2.ipynb', 'repo_name': 'Jakl3/Polygence', 'size': '13681'}