diff --git a/recognition/2D_Unet_47423510/README.md b/recognition/2D_Unet_47423510/README.md new file mode 100644 index 000000000..495ce0d02 --- /dev/null +++ b/recognition/2D_Unet_47423510/README.md @@ -0,0 +1,31 @@ +# 2D UNet For Predicting Prostate Cancer from HipMRI Scans +**NOTE** -- Unable to conduct tests as training data was made unavaliable (https://edstem.org/au/courses/18266/discussion/2340286) & email to shakes. +## Prostate Cancer HipMRI Application +The purpose of this CNN UNet model is to assist in the diagnosis of prostate cancer through analysis of Hip MRI scans. The model's purpose is to be trained on images, such that it can segment and identify points of interest within the image. By training the Unet model in such a way that it recognises propoerties of cancer and malignent cells, it can assist profesionals in pinpointing the precise location of these issues. +## Algorithm Description + +### UNet Model +The UNet model is a convoluted neural network that works by taking image segmentation, encoding it via several convolution layers. Then decoding it via the transpose convolution, while taking raw imports from the opposite encoding stage. That is to say, that a 2D Unet Model, for each layer of convolution stores the resultant convoluted image, and then concatenates these convoluted images with the same depth and width images obtained during the decoding steps. In doing this trianing can be conducted faster as more of the initial image is retained during training steps. + +### Encoding steps +Each level of encoding consists of 2 convolution steps, used for capturing the context of the image, and then a pooling step which reduces the dimensions of the image in preperation of the next step. This dimension reduction results in a increase in channels porportional to the level of dimension reduction. This spacial reduction is done until a desired bottleneck dimension, with a large number of convolution channels, is reached. Each convolution step is done using the Tensorflow.keras library's inbuilt Convolution function (https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2D). The kernel size used was 3x3, with the convolution being padded such that the output was the same dimensions as the input. The activation function chosen was a rectifier function, such that only the positive components of the arguments are used. + +### Decoding steps +Once the bottleneck is reached, 2 more convolution steps are done, but no further dimension reduction is conducted on the image, in preperation of the restoration of the image. During the decode steps of the model, The reduced image is transformed using Tensorflow.keras library's transpose convolution (https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2DTranspose), such that a dimension reduction step is undone. Resulting in a halving of channels, but a doubling of image size. The image is then concatenated with a snapshot of the image at the encoded layer of the same dimension. this combined image is then convoluted twice using the same parameters as encode, before undergoing the next decoding step. Once the image is of the original image's dimension it is convoluted once with a 1x1 kernel through 1 channel, using softmax activation so a segmented image is obtained. + +The model is then trained, such that each channel's weighting and bias can be accurately estimated by the model. Resulting in a reduction of the loss value. The loss function chosen was the catagorical crossentropy (https://www.tensorflow.org/api_docs/python/tf/keras/losses/categorical_crossentropy). + +## ~~ Example usage~~ + +### Preprocessing +The program handles all several aspects of dataset loading and manipulation prior to the main training loop to ensure efficent training procedure. + +## Dependicies +**Python3** +**Tensorflow** +**Numpy** +**matplotlib.pyplot** +**sklearn.metrics** + +## References +relevant references are linked throughout the report, alongside the lecture example for a UNet and succesful pull requests from previous years. diff --git a/recognition/2D_Unet_47423510/dataset.py b/recognition/2D_Unet_47423510/dataset.py new file mode 100644 index 000000000..927e27a3a --- /dev/null +++ b/recognition/2D_Unet_47423510/dataset.py @@ -0,0 +1,50 @@ +# This is for the dataset + +import numpy as np +import nibabel as nib +from tqdm import tqdm + +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 + +def load_data_2D(imageNames, normImage=False, categorical=False, dtype=np.float32, getAffines=False, early_stop=False): + affines = [] + num = len(imageNames) + first_case = nib.load(imageNames[0]).get_fdata(caching='unchanged') + if (len(first_case.shape) == 3): + first_case = first_case[:,:,0] + 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') + affine = niftiImage.affine + if (len(inImage.shape) == 3): + inImage = inImage[:, :, 0] + inImage = inImage.astype(dtype) + if normImage: + inImage = ((inImage - inImage.mean()) / inImage.std()) + if categorical: + inImage = 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 \ No newline at end of file diff --git a/recognition/2D_Unet_47423510/modules.py b/recognition/2D_Unet_47423510/modules.py new file mode 100644 index 000000000..f4cc939d6 --- /dev/null +++ b/recognition/2D_Unet_47423510/modules.py @@ -0,0 +1,59 @@ +import numpy as np +import tensorflow as tf +from tensorflow.keras import layers, models + + +def init_unet(tensor_in, depth=32): + """ + Function to initial the convolution layers of the UNet. + parameters: + - tensor_in - the initial tensor filter to use for the model + - depth - the number of channels/depth to use at the top layer of extraction (default 32 (based on lecture)) + Return: + NULL + """ + # For each layer of the UNET filter convolute twice and then pool. + # Complete 2 convolutions to condense the image + # Model based on https://colab.research.google.com/drive/1K2kiAJSCa6IiahKxfAIv4SQ4BFq7YDYO?usp=sharing + encode1 = layers.Conv2D(depth, 3, activation='relu', padding='same')(tensor_in) + encode1 = layers.Conv2D(depth, 3, activation='relu', padding='same')(encode1) + # Complete dimension reduction + pool1 = layers.MaxPooling2D(pool_size=(2,2), padding='same')(encode1) + + encode2 = layers.Conv2D(depth*2, 3, activation='relu', padding='same')(pool1) + encode2 = layers.Conv2D(depth*2, 3, activation='relu', padding='same')(encode2) + pool2 = layers.MaxPooling2D(pool_size=(2,2), padding='same')(encode2) + + encode3 = layers.Conv2D(depth*4, 3, activation='relu', padding='same')(pool2) + encode3 = layers.Conv2D(depth*4, 3, activation='relu', padding='same')(encode3) + pool3 = layers.MaxPooling2D(pool_size=(2,2), padding='same')(encode3) + + encode4 = layers.Conv2D(depth*8, 3, activation='relu', padding='same')(pool3) + encode4 = layers.Conv2D(depth*8, 3, activation='relu', padding='same')(encode4) + pool4 = layers.MaxPooling2D(pool_size=(2,2), padding='same')(encode4) + + encode5 = layers.Conv2D(depth*16, 3, activation='relu', padding='same')(pool4) + encode5 = layers.Conv2D(depth*16, 3, activation='relu', padding='same')(encode5) + + up1 = layers.Conv2DTranspose(depth*8, 2, strides=2, padding='same')(encode5) + merge1 = layers.concatenate([encode4, up1]) + decode1 = layers.Conv2D(depth*8, 3, activation='relu', padding='same')(merge1) + decode1 = layers.Conv2D(depth*8, 3, activation='relu', padding='same')(decode1) + + up2 = layers.Conv2DTranspose(depth*4, 2, strides=2, padding='same')(decode1) + merge2 = layers.concatenate([encode3, up2]) + decode2 = layers.Conv2D(depth*4, 3, activation='relu', padding='same')(merge2) + decode2 = layers.Conv2D(depth*4, 3, activation='relu', padding='same')(decode2) + + up3 = layers.Conv2DTranspose(depth*2, 2, strides=2, padding='same')(decode2) + merge3 = layers.concatenate([encode2, up3]) + decode3 = layers.Conv2D(depth*2, 3, activation='relu', padding='same')(merge3) + decode3 = layers.Conv2D(depth*2, 3, activation='relu', padding='same')(decode3) + + up4 = layers.Conv2DTranspose(depth, 2, strides=2, padding='same')(decode3) + merge4 = layers.concatenate([encode1, up4]) + decode4 = layers.Conv2D(depth, 3, activation='relu', padding='same')(merge4) + decode4 = layers.Conv2D(depth, 3, activation='relu', padding='same')(decode4) + tensor_out = layers.Conv2D(1, 1, activation='softmax', padding='same')(decode4) + model = models.Model(inputs=tensor_in, outputs=tensor_out) + return model \ No newline at end of file diff --git a/recognition/2D_Unet_47423510/predict.py b/recognition/2D_Unet_47423510/predict.py new file mode 100644 index 000000000..9fb7909db --- /dev/null +++ b/recognition/2D_Unet_47423510/predict.py @@ -0,0 +1,48 @@ +import numpy as np +import tensorflow as tf +from sklearn.metrics import classification_report +from dataset import load_data_2D +from train import train_unet +import matplotlib.pyplot as plt + + +# Function to take a trained unet model and use it for predictions. +def unet_predict(model, data_test): + """ + Function to predict the model + Paramaters: + - model - the trained unet model to use for the predictions + - data_test - the data set to use for predictions + """ + predictions = model.predict(data_test) + # Get predicted segmented images + print(np.argmax(predictions, axis=1)) + print(classification_report(data_test, predictions)) + + # Create a plot to plot predictions against tests + # reference: https://matplotlib.org/stable/gallery/subplots_axes_and_figures/subplots_demo.html + _, ax = plt.subplots(len(predictions), 2) + for i in range(len(predictions)): + ax[0, i].set_title(f"[{i}] prediction") + ax[0, i].im_show(predictions[i].squeeze()) + ax[1, i].set_title("Test Data") + ax[1, i].im_show(data_test[i].squeeze()) + plt.show() + + +def main(): + # load and evaluate model reference: https://www.tensorflow.org/tutorials/keras/save_and_load + # load the trained model + model = tf.keras.models.load_model('mri_unet.keras') + model.summary() + data_train = load_data_2D() + data_validate = load_data_2D() + data_test = load_data_2D() + _, acc = model.evaluate(data_train, data_validate, verbose=2) + print('Restored model, accuracy: {:5.2f}%'.format(100 * acc)) + unet_predict(model, data_test) + + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/recognition/2D_Unet_47423510/train.py b/recognition/2D_Unet_47423510/train.py new file mode 100644 index 000000000..b3e69d06c --- /dev/null +++ b/recognition/2D_Unet_47423510/train.py @@ -0,0 +1,46 @@ +# This is for the training algorithm +import numpy as np +import tensorflow as tf +from tensorflow.keras import layers, models +from tensorflow.keras import backend as K +from tensorflow.keras.callbacks import EarlyStopping + +from modules import init_unet +from dataset import load_data_2D, to_channels + +# https://medium.com/@vipul.sarode007/u-net-unleashed-a-step-by-step-guide-on-implementing-and-training-your-own-segmentation-model-in-a90ed89399c6 +# Function to calculate the mathmatical dice coefficient +def dice(y_true, y_predicted): + """ + Function to calculate the max label dice coefficient + + """ + return (2 * K.sum(y_true*y_predicted, axis = -1)) / (K.sum(y_true, axis = -1) + K.sum(y_predicted, axis = -1)) + + +#Function to train the Unet model, using tensorflow's inbuilt keras.model.fit() function +def train_unet(data_train, data_train_seg, data_validate, data_validate_seg, epochs=32, batch_size=8): + """ + Training function for the unet based on provided data and segmented data. + """ + init_inputs = layers.Input(shape=(256, 128, 1)) + model = init_unet(init_inputs) + early_stopping = EarlyStopping(monitor = 'val_loss', patience = 3, restore_best_weights = True) + # Compile and train the UNET model on the training set and validation set + # Using the loss function for ones hot encoding + model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy', dice]) + model.fit(data_train, data_train_seg, epochs=epochs, batch_size=batch_size, validation_data=(data_validate, data_validate_seg), callbacks=[early_stopping], verbose=1) + # save model reference: https://www.tensorflow.org/tutorials/keras/save_and_load + model.save('mri_unet.keras') + +def main(): + data_train = load_data_2D() + data_train_seg = load_data_2D() + + data_validate = load_data_2D() + data_validate_seg = load_data_2D() + + train_unet(data_train, data_train_seg, data_validate, data_validate_seg) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/recognition/README.md b/recognition/README.md deleted file mode 100644 index 32c99e899..000000000 --- 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