diff --git a/README.md b/README.md deleted file mode 100644 index 3a10f6515f..0000000000 --- a/README.md +++ /dev/null @@ -1,19 +0,0 @@ -# Pattern Analysis -Pattern Analysis of various datasets by COMP3710 students in 2024 at the University of Queensland. - -We create pattern recognition and image processing library for Tensorflow (TF), PyTorch or JAX. - -This library is created and maintained by The University of Queensland [COMP3710](https://my.uq.edu.au/programs-courses/course.html?course_code=comp3710) students. - -The library includes the following implemented in Tensorflow: -* fractals -* recognition problems - -In the recognition folder, you will find many recognition problems solved including: -* segmentation -* classification -* graph neural networks -* StyleGAN -* Stable diffusion -* transformers -etc. diff --git a/recognition/README.md b/recognition/README.md deleted file mode 100644 index 32c99e8993..0000000000 --- a/recognition/README.md +++ /dev/null @@ -1,10 +0,0 @@ -# Recognition Tasks -Various recognition tasks solved in deep learning frameworks. - -Tasks may include: -* Image Segmentation -* Object detection -* Graph node classification -* Image super resolution -* Disease classification -* Generative modelling with StyleGAN and Stable Diffusion \ No newline at end of file diff --git a/recognition/unet_hipmri_s4646244/README.md b/recognition/unet_hipmri_s4646244/README.md new file mode 100644 index 0000000000..e08090ead3 --- /dev/null +++ b/recognition/unet_hipmri_s4646244/README.md @@ -0,0 +1,83 @@ +# 2D Unet to segment MRI images of prostate cancer +This repository contains a TensorFlow Keras implimentation of a Binary classification Unet to segment prostate cancer MRI scans. + +#### Files: +dataset.py +- A file used to read in the .nii files into training, validation and test sets +modules.py + +- A file containing the Unet architecture, including the Encoder, decoder and bottleneck + +train.py +- A file containing the functions to train the unet and save the model + +predict.py +- A file to show how well the Unet predicts test images + +#### Folders +Unet_images +- A folder containing the testing results of running different epochs + +### Model +###### Filters +The filters applied during the encoder step are [64,128,256,512] and [512,256,128,64] for the decoder. +###### Encoder +The encoder takes the input tensor and for each filter in the list completes two convolutions with a 3x3 kernel size and a relu activation function. +It then saves the resulting tensor into the skip connection so it can be used in the decoder. +A 2x2 max pooling is then used and the resulting tensor output. +This happens four times, once for each filter. + +###### Bottleneck +The bottleneck step applies two convolutions using a 3x3 kernel and the relu activation function. +it then outputs the tensor to be used in the decoder + +###### Decoder +The decoder then takes this tensor, completes an up convolution using a 2x2 kernel and a relu function +The relevent skip connection is then concatenated with the tensor from the upconv. +This concated tensor is then fed through two convolutions using a 3x3 kernel and relu activiation +One last convolution using the sigmoid function is used to complete the binary classification. + +### Training +The Unet is trained with an adam optimiser with an initial learning rate of 0.0001. It uses a combined loss function that consists of binary cross entropy and dice loss as this resulted in the best performance. +The training also uses early stoppage, this stops the training when the validation loss performance does not increase after three consecutive epochs. +A learning rate scheduler is also used to monitor the validation loss, if the validation does not improve after two epochs the learning rate is halved. +In the results twelve epochs were used to reduce the training time however only 6 were run as the model stopped early. +It also uses a validation set to evaluate performance after each epoch, helping to monitor overfitting by checking how well the model performs on new data. + +### Testing +After the model is trained it then is given new data to segment. This data is + +### Performance +The model was run using a batch size of four and twelve epochs, only six were run as it stopped early. +The results for the model being trained on three and six epochs can be seen in the Unet_images folder. + +The model is very good at binary segmentaion of the prostate cancer images, If i had more time I would convert it to do multi class segmentation. +It segments most of the regions well except very small areas. +It can be seen that the mean dice test score is just above 0.65 however there are many datapoints that fall below this region. +When looking at the dice coefficents over each epoch it sharply increases and then slowly tapers off, this is the same for the loss function except it sharply decreases. + +#### Required dependencies +- TensorFlow (for Keras layers, models, and callbacks) +- NumPy (for numerical operations) +- Matplotlib (for plotting) +- NiBabel (for neuroimaging data handling) +- tqdm (for progress bars) +- scikit-image (for image transformations) +- pathlib (for filesystem path manipulations) + +#### Future improvements +The dataset given is for multiclass segmentation. The implimentation of my unet and training is only for binary classification, I tried implimenting the multiclass model by using a softmax activation function rather than sigmoid and modifying my train.py to handle the multiple classes however I could not get it working. In the future I will look into modifying the implimentation so that it can do multiclass segmentation. + +#### How to run +Before running the model the relevent files paths need to be added into dataset.py +Once this is done all that is needed to be run is the predict.py file with no arguments. +This will train, validate, test and print the results of the model. + +If a powerful graphics card is in your system, you may be able to increase the batch size in train.py this will result in faster training. + +### References +Reference for Dice coefficient metric implementation in the train.py function +Stack Overflow. "Dice coefficient not increasing for U-Net image segmentation." +https://stackoverflow.com/questions/67018431/dice-coefficent-not-increasing-for-u-net-image-segmentation + + diff --git a/recognition/unet_hipmri_s4646244/Unet_images/12epochs/DiceDistribution.png b/recognition/unet_hipmri_s4646244/Unet_images/12epochs/DiceDistribution.png new file mode 100644 index 0000000000..c1eaeede80 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/12epochs/DiceDistribution.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/12epochs/DiceImprovement.png b/recognition/unet_hipmri_s4646244/Unet_images/12epochs/DiceImprovement.png new file mode 100644 index 0000000000..7091a9d607 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/12epochs/DiceImprovement.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/12epochs/everyTestImageDice.png b/recognition/unet_hipmri_s4646244/Unet_images/12epochs/everyTestImageDice.png new file mode 100644 index 0000000000..1be3846b06 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/12epochs/everyTestImageDice.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/12epochs/predictedseg.png b/recognition/unet_hipmri_s4646244/Unet_images/12epochs/predictedseg.png new file mode 100644 index 0000000000..56b9e0469c Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/12epochs/predictedseg.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Distribution_Dice.png b/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Distribution_Dice.png new file mode 100644 index 0000000000..333d466d04 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Distribution_Dice.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Epoch_improvement.png b/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Epoch_improvement.png new file mode 100644 index 0000000000..84f9f41005 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Epoch_improvement.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Five_original_true_pred.png b/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Five_original_true_pred.png new file mode 100644 index 0000000000..d84bba3709 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Five_original_true_pred.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Overall_performance_test.png b/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Overall_performance_test.png new file mode 100644 index 0000000000..8899b797b7 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/3epochs/Overall_performance_test.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/6epochs/DiceImprovement.png b/recognition/unet_hipmri_s4646244/Unet_images/6epochs/DiceImprovement.png new file mode 100644 index 0000000000..0e637f9da4 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/6epochs/DiceImprovement.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/6epochs/Diceoverall.png b/recognition/unet_hipmri_s4646244/Unet_images/6epochs/Diceoverall.png new file mode 100644 index 0000000000..66f4365a46 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/6epochs/Diceoverall.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/6epochs/modelVSeachImage.png b/recognition/unet_hipmri_s4646244/Unet_images/6epochs/modelVSeachImage.png new file mode 100644 index 0000000000..18d1ade4c5 Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/6epochs/modelVSeachImage.png differ diff --git a/recognition/unet_hipmri_s4646244/Unet_images/6epochs/pred_real.png b/recognition/unet_hipmri_s4646244/Unet_images/6epochs/pred_real.png new file mode 100644 index 0000000000..b7ce8f263a Binary files /dev/null and b/recognition/unet_hipmri_s4646244/Unet_images/6epochs/pred_real.png differ diff --git a/recognition/unet_hipmri_s4646244/dataset.py b/recognition/unet_hipmri_s4646244/dataset.py new file mode 100644 index 0000000000..18817e7afd --- /dev/null +++ b/recognition/unet_hipmri_s4646244/dataset.py @@ -0,0 +1,107 @@ +import numpy as np +import nibabel as nib +from tqdm import tqdm +import skimage.transform as skTrans +from pathlib import Path +import tensorflow as tf + +def to_channels(arr: np.ndarray, dtype=np.uint8) -> np.ndarray: + channels = np.unique(arr) + res = np.zeros(arr.shape + (len(channels),), dtype=dtype) + for c in channels: + c = int(c) + res[..., c:c + 1][arr == c] = 1 + return res + +# load medical image functions +def load_data_2D(imageNames, normImage=False, categorical=False, dtype=np.float32, getAffines=False, early_stop=False): + ''' + Load medical image data from names, cases list provided into a list for each. + + This function pre-allocates 4D arrays for conv2d to avoid excessive memory usage. + + normImage: bool (normalise the image 0.0 -1.0) + early_stop: Stop loading prematurely, leaves arrays mostly empty, for quick loading and testing scripts. + ''' + affines = [] + + # get fixed size + num = len(imageNames) + first_case = nib.load(imageNames[0]).get_fdata(caching='unchanged') + if len(first_case.shape) == 3: + first_case = first_case[:, :, 0] # sometimes extra dims, remove + if categorical: + first_case = to_channels(first_case, dtype=dtype) + rows, cols, channels = first_case.shape + images = np.zeros((num, rows, cols, channels), dtype=dtype) + else: + rows, cols = first_case.shape + images = np.zeros((num, rows, cols), dtype=dtype) + + for i, inName in enumerate(tqdm(imageNames)): + niftiImage = nib.load(inName) + inImage = niftiImage.get_fdata(caching='unchanged') # read disk only + + affine = niftiImage.affine + if len(inImage.shape) == 3: + inImage = inImage[:, :, 0] # sometimes extra dims in HipMRI_study data + + # Converts the image to a 256,128 image + inImage = skTrans.resize(inImage, (256, 128), order=1, preserve_range=True) + + inImage = inImage.astype(dtype) + if normImage: + # ~ inImage = inImage / np.linalg.norm(inImage) + # ~ inImage = 255. * inImage / inImage.max() + inImage = (inImage - inImage.mean()) / inImage.std() + + if categorical: + inImage = utils.to_channels(inImage, dtype=dtype) + images[i, :, :, :] = inImage + + else: + images[i, :, :] = inImage + + affines.append(affine) + if i > 20 and early_stop: + break + + if getAffines: + return images, affines + else: + return images + +testDir = '/home/kankuna/Documents/COMP3710DATA/HipMRI_study_keras_slices_data/keras_slices_test/' +trainDir = '/home/kankuna/Documents/COMP3710DATA/HipMRI_study_keras_slices_data/keras_slices_train.large/' +validateDir = '/home/kankuna/Documents/COMP3710DATA/HipMRI_study_keras_slices_data/keras_slices_validate/' + +testSegDir = '/home/kankuna/Documents/COMP3710DATA/HipMRI_study_keras_slices_data/keras_slices_seg_test/' +trainSegDir = '/home/kankuna/Documents/COMP3710DATA/HipMRI_study_keras_slices_data/keras_slices_seg_train.large/' +validateSegDir = '/home/kankuna/Documents/COMP3710DATA/HipMRI_study_keras_slices_data/keras_slices_seg_validate/' + +# Load the scans +from pathlib import Path + +# Load the test images +testListNii = sorted(Path(testDir).glob('*.nii')) +testImages = load_data_2D(testListNii, normImage=True, categorical=False) + +# Load the training images +trainListNii = sorted(Path(trainDir).glob('*.nii')) +trainImages = load_data_2D(trainListNii, normImage=True, categorical=False) + +# Load the validation images +validateListNii = sorted(Path(validateDir).glob('*.nii')) +validateImages = load_data_2D(validateListNii, normImage=True, categorical=False) + +# Load the segmented test scans +testSegListNii = sorted(Path(testSegDir).glob('*.nii')) +testSegImages = load_data_2D(testSegListNii, normImage=True, categorical=False) + +# Load the segmented training scans +trainSegListNii = sorted(Path(trainSegDir).glob('*.nii')) +trainSegImages = load_data_2D(trainSegListNii, normImage=True, categorical=False) + +# Load the segmented validation scans +validateSegListNii = sorted(Path(validateSegDir).glob('*.nii')) +validateSegImages = load_data_2D(validateSegListNii, normImage=True, categorical=False) diff --git a/recognition/unet_hipmri_s4646244/modules.py b/recognition/unet_hipmri_s4646244/modules.py new file mode 100644 index 0000000000..3e15d7973e --- /dev/null +++ b/recognition/unet_hipmri_s4646244/modules.py @@ -0,0 +1,51 @@ +from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Conv2DTranspose, concatenate # type: ignore +from tensorflow.keras.models import Model # type: ignore + +# List of filters to be applied in the encoding and reverse in the decoding step +filterList = [64,128,256,512] + +# Function that does the encoding of the Unet by applying convolutions and max pooling to the given tensor at each filter level +# Parameters: inputTensor, a tf tensor that will be encoded +# Returns: (tensor, skipConnectionList), a tuple containing the resulting encoded tensor and a list of the skip connections at each step +def encoder(inputTensor): + skipConnectionList = [] + tensor = inputTensor + for filter in filterList: + firstConv = Conv2D(filter, kernel_size = (3,3), padding = 'same', strides = 1, activation = 'relu')(tensor) + secondConv = Conv2D(filter, kernel_size = (3,3), padding = 'same', strides = 1, activation = 'relu')(firstConv) + skipConnectionList.append(secondConv) + tensor = MaxPooling2D(pool_size = (2,2), padding = 'same')(secondConv) + return tensor, skipConnectionList + +# Function that does the decoding of the Unet by applying an up convolution followed by concatting the tensor with the skip connection and then applying convolutions. +# Parameters: inputTensor, a tf tensor that will be decoded +# Returns: tensor, a decoded tensor +def decoder(skipConnectionList, inputTensor): + tensor = inputTensor + for filter in reversed(filterList): + upConv = Conv2DTranspose(filter, kernel_size = (2,2), padding = 'same', activation = 'relu', strides = 2)(tensor) + skipConnection = skipConnectionList.pop() + concatTensor = concatenate([upConv, skipConnection]) + firstConv = Conv2D(filter, kernel_size = (3,3), padding = 'same', strides = 1, activation = 'relu')(concatTensor) + secondConv = Conv2D(filter, kernel_size = (3,3), padding = 'same', strides = 1, activation = 'relu')(firstConv) + tensor = secondConv + finalConv = Conv2D(1, kernel_size=(1, 1), padding='same', strides=1, activation='sigmoid')(tensor) + return finalConv + +# Function that applies the bottleneck of the unet. +# Parameters: inputTensor, a tf tensor that will be decoded +# Returns: tensor, a tensor that has gone through two convolutions +def bottleneck(inputTensor): + firstConv = Conv2D(1024, kernel_size = (3,3), padding = 'same', strides = 1, activation = 'relu')(inputTensor) + secondConv = Conv2D(1024, kernel_size = (3,3), padding = 'same', strides = 1, activation = 'relu')(firstConv) + tensor = secondConv + return tensor + +# Function that Applies the encoder, bottleneck and decoder into one unet model +# Returns a keras model of the unet +def unet(): + inputs = Input(shape = (256, 128, 1)) + encodedResult, skipConnectionList = encoder(inputs) + bottleneckResult = bottleneck(encodedResult) + decodedResult = decoder(skipConnectionList, bottleneckResult) + return Model(inputs=[inputs], outputs=[decodedResult]) diff --git a/recognition/unet_hipmri_s4646244/predict.py b/recognition/unet_hipmri_s4646244/predict.py new file mode 100644 index 0000000000..7fb0fa370d --- /dev/null +++ b/recognition/unet_hipmri_s4646244/predict.py @@ -0,0 +1,86 @@ +import train +import matplotlib.pyplot as plt +from train import unetModel, dice_metric, trainResults +from dataset import testImages, testSegImages +import numpy as np + +testPredictedSeg = unetModel.predict(testImages) +print(np.unique(testPredictedSeg)) + +#Function to find the dice score of each set of actual segments and predicted +def calculate_dice_scores(y_true, y_pred): + dice_scores = [] + for i in range(len(y_true)): + y_pred_squeezed = np.squeeze(y_pred[i]) + score = dice_metric(y_pred_squeezed, y_true[i]).numpy() + dice_scores.append(score) + return dice_scores + +dice_scores = calculate_dice_scores(testSegImages, testPredictedSeg) +dice_scores = np.array(dice_scores) + +#Print the actual image, actual segment and predicted segment +fig, pos = plt.subplots(5, 3, figsize=(15, 25)) +for i in range(5): + # Display original image + pos[i, 0].imshow(testImages[i].squeeze()) + pos[i, 0].set_title(f'Original image {i+1}') + pos[i, 0].axis('off') + + # Display actual segmentation + pos[i, 1].imshow(testSegImages[i].squeeze()) + pos[i, 1].set_title(f'Actual segmentation {i+1}') + pos[i, 1].axis('off') + + # Display predicted segmentation + pos[i, 2].imshow(testPredictedSeg[i].squeeze()) + pos[i, 2].set_title(f'Predicted segmentation {i+1}') + pos[i, 2].axis('off') +plt.tight_layout() +plt.show() + +#print the dice scores for each image and the distribution +plt.figure(figsize=(12, 6)) +plt.plot(dice_scores, marker='o', linestyle='None', color='b') +plt.title('Dice Scores for Each Test Image') +plt.xlabel('Test Image Index') +plt.ylabel('Dice Score') +plt.ylim(0, 1) +plt.yticks(np.linspace(0, 1, num=11)) +plt.grid() +plt.axhline(y=np.mean(dice_scores), color='r', linestyle='--', label='Mean Dice Score') +plt.legend() +plt.show() + +plt.figure(figsize=(12, 6)) +plt.hist(dice_scores, bins=10, color='c', edgecolor='black', alpha=0.7) +plt.title('Distribution of Dice Scores') +plt.xlabel('Dice Score') +plt.ylabel('Frequency') +plt.xlim(0, 1) +plt.grid() +plt.show() + +# Plotting Loss and Dice Coefficient +plt.figure(figsize=(12, 5)) +plt.subplot(1, 2, 1) +plt.plot(trainResults.history['loss'], label='Training Loss') +plt.plot(trainResults.history['val_loss'], label='Validation Loss') +plt.title('Loss Over Epochs') +plt.xlabel('Epochs') +plt.ylabel('Loss') +plt.legend() +plt.grid() + +# Plot Dice Coefficient +plt.subplot(1, 2, 2) +plt.plot(trainResults.history['dice_metric'], label='Training Dice Coefficient') +plt.plot(trainResults.history['val_dice_metric'], label='Validation Dice Coefficient') +plt.title('Dice Coefficient Over Epochs') +plt.xlabel('Epochs') +plt.ylabel('Dice Coefficient') +plt.legend() +plt.grid() + +plt.tight_layout() +plt.show() diff --git a/recognition/unet_hipmri_s4646244/train.py b/recognition/unet_hipmri_s4646244/train.py new file mode 100644 index 0000000000..9aa18908a0 --- /dev/null +++ b/recognition/unet_hipmri_s4646244/train.py @@ -0,0 +1,47 @@ +import tensorflow as tf +import numpy as np +from tensorflow.keras import backend as K # type: ignore +from tensorflow.keras.optimizers import Adam # type: ignore +from tensorflow.keras.losses import BinaryCrossentropy # type: ignore +from tensorflow.keras.metrics import BinaryAccuracy # type: ignore +from tensorflow.keras.callbacks import ReduceLROnPlateau # type: ignore +from tensorflow.keras.callbacks import EarlyStopping # type: ignore +from modules import unet +from dataset import testImages, trainImages, validateImages, testSegImages, trainSegImages, validateSegImages + +# Function to calculate the dice ceofficient +# Reference for Dice coefficient metric implementation: +# https://stackoverflow.com/questions/67018431/dice-coefficent-not-increasing-for-u-net-image-segmentation +def dice_metric(y_pred, y_true): + intersection = K.sum(K.sum(K.abs(y_true * y_pred), axis=-1)) + union = K.sum(K.sum(K.abs(y_true) + K.abs(y_pred), axis=-1)) + return 2*intersection / union + +# function to find the dice loss +def dice_loss(y_true, y_pred): + return 1 - dice_metric(y_true, y_pred) + +#Function to take the combined loss from binary cross entropy and dice +def combined_loss(y_true, y_pred): + return dice_loss(y_true, y_pred) + BinaryCrossentropy()(y_true, y_pred) + +#learning rate scheduler and early stoppage +LearningRateScheduler = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=2, min_lr=1e-6, verbose=1) +earlyStopping = EarlyStopping(monitor = 'val_loss', patience = 3, restore_best_weights = True) + +#Create a unet instance +unetModel = unet() + +#compile the unet with a the adam optimiser and a learning rate of 0.0001. using the combined loss and the dice metric +unetModel.compile(optimizer=Adam(learning_rate=0.0001), loss=combined_loss, metrics=[dice_metric]) + +#print the unet model information +#unetModel.summary() + +# Run the training on the model +trainResults = unetModel.fit(trainImages, trainSegImages, validation_data = (validateImages, validateSegImages), + batch_size = 4, epochs=12, callbacks=[earlyStopping, LearningRateScheduler], verbose=1) + +# Run the trained model on the test datasets +testResults = unetModel.evaluate(testImages, testSegImages, batch_size = 1) +unetModel.save('unet_model.keras')