2D U-Net HipMRI s46978116

recognition/prostate2dnet_46978116/README.MD

@@ -0,0 +1,104 @@
+# Segmenting HipMRI Study for prostate cancer radiotherapy Data 
+
+In this report, we will attempt to segment HipMRI data based on the part of the body the image depicts with the help of 2D U-Net Convolutional Neural Network. 
+
+# Understanding the Dataset
+
+The dataset consists of 2D MRI slices of the male pelvis collected from a radiation therapy study at Calvary Mater Newcastle Hospital. Each MRI image has a corresponding segmentation image, which we will refer to as the mask, that segments the MRI image based on the part of the body depicted in each segment. There are six classes to these segments: the background, bladder, body, bone, rectum and prostate. Each image is roughly 256 x 128 pixels although some are slightly larger than this so they have all been resized to 256 x 256. Our task is to develop a method that accurately segments MRI data of the male pelvis based on training data collected from these images and their masks. 
+
+# The 2D U-Net Model
+
+The chosen model to learn this segmentation is a 2D U-Net. The 2D U-Net a deep neural network developed by Ronneburg, etal in 2015. This model shows promising performance in segmentation of medical data, making it an ideal fit for the task. The 2D U-Net utilisizes an encoder-decoder structure with skip connections. This makes it effective in computervision segmentation problems as it can capture the spatial information of the data well and retains information through the skip connections between the encoder and decoder layers. The model also trains quickly with high accuracy.
+
+![2D U-Net Architecture](images\2dunet_architecture.png)
+
+The architecture consists of two 2D convolutional layers which are each followed by max pooling layers. The activation function is ReLu and a 2D Batch normalization is used as well. Although Batch normalization was not used in the original paper (because it wasn't developed yet) we have used it here to improve performance. The data is downsampled using these blocks until it reaches a bottleneck at 512 feature channels. Here, it is upscaled back to the original image size using upsampling layers with skip connections to retain information while reaching a detailed resolution. The final layer has 6 outchannels which corresponds with the 6 classes of segments.
+
+# Output and Results
+
+After training on the training set with early stoppage the the dice co-efficients on the test set using cross-entropy loss and minimal transformations:
+
+Class 0: 0.9945
+Class 1: 0.9768
+Class 2: 0.8689
+Class 3: 0.6563
+Class 4: 0.3489
+Class 5: 0.2344
+Average Dice Score: 0.6800
+
+Below are the plots of the training performance:
+
+![Training Performance CrossEntropyLoss](images\graphs\crossplot.png)
+
+And the predictions made by the model:
+
+![Best Predictions [Epoch 17] CrossEntropyLoss (Classes: 0, 1, 2)](images\predictions\bestcross10\epoch_test_classes_0-2.png)
+
+![Best Predictions [Epoch 17] CrossEntropyLoss (Classes: 3, 4, 5)](images\predictions\bestcross10\epoch_test_classes_3-5.png)
+
+
+close up predictions:
+
+
+![Best Predictions Close Up [Epoch 17] CrossEntropyLoss (Classes: 0, 1, 2)](images\predictions\bestcross2\epoch_test_classes_0-2.png)
+
+![Best Predictions Close Up [Epoch 17] CrossEntropyLoss (Classes: 3, 4, 5)](images\predictions\bestcross2\epoch_test_classes_3-5.png)
+
+
+On the left is the MRI image, the middle is real segment for the corresponding MRI image and the right is the prediction from the model.
+
+Notably, the model underperformed in segmenting for class 3, 4 and 5. These classes were identified early on as having poor performance compared to other classes which can be attributed to their very low prevelance rate. The vast majority of MRI slides do not contain these segments and if they do its very small compared to the other classes. 
+
+To try to combat this transformations were performed on the images to try to improve segmenting of these classes such as rotations and zooming. Furthermore, the loss function was changed to directly include the dice loss (1 - dice coefficient) for each class. This dice loss had an added penalty when the co-efficient was below 0.75 to further incentivize the model to learn how to segment class 4 and 5. Below is the test dice coefficients after these changes, showing a very large improvement especially in the problem classes 3, 4, and 5:
+
+Class 0: 0.9863
+Class 1: 0.9777
+Class 2: 0.8820
+Class 3: 0.7254
+Class 4: 0.6833
+Class 5: 0.8711
+Average Dice Score: 0.8543
+
+And the plots of training performance and predictions also show improvement:
+
+![Training Performance DiceLoss](images\graphs\diceplot.png)
+
+The predictions made by the model:
+
+![Best Predictions [Epoch 40] DiceLoss (Classes: 0, 1, 2)](images\predictions\bestdice10\epoch_test_classes_0-2.png)
+
+![Best Predictions [Epoch 40] DiceLoss (Classes: 3, 4, 5)](images\predictions\bestdice10\epoch_test_classes_3-5.png)
+
+
+close up predictions:
+
+![Best Predictions Close Up [Epoch 40] DiceLoss (Classes: 0, 1, 2)](images\predictions\bestdice2\epoch_test_classes_0-2.png)
+
+![Best Predictions Close Up [Epoch 40] DiceLoss (Classes: 3, 4, 5)](images\predictions\bestdice2\epoch_test_classes_3-5.png)
+
+
+# Requirements
+
+torch
+torchvision
+scikit-image
+tqdm
+nibabel
+numpy
+albumentatinos
+matplotlib
+
+# References
+
+[1] O. Ronneberger, P. Fischer, and T. Brox, “U-Net: Convolutional Networks for Biomedical Image Segmentation,” in Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015, ser. Lecture Notes in Computer Science, N. Navab, J. Hornegger, W. M. Wells, and A. F. Frangi, Eds. Cham: Springer International Publishing, 2015, pp. 234–241. Available: https://doi.org/10.48550/arXiv.1505.04597
+
+[2] J. Schmidt, “Creating and training a U-Net model with PyTorch for 2D & 3D semantic segmentation: Dataset building,” Towards Data Science, Dec. 2, 2020. [Online]. Available: https://towardsdatascience.com/creating-and-training-a-u-net-model-with-pytorch-for-2d-3d-semantic-segmentation-dataset-fb1f7f80fe55.
+
+[3] A. Persson, “PyTorch Image Segmentation Tutorial with U-NET: everything from scratch baby,” YouTube, Feb. 23, 2021. [Online]. Available: https://www.youtube.com/watch?v=IHq1t7NxS8k.
+
+[4] Esri, “How U-net works,” ArcGIS API for Python Documentation. [Online]. Available: https://developers.arcgis.com/python/latest/guide/how-unet-works/.
+
+[5] J. Dowling and P. Greer, “Labelled weekly MR images of the male pelvis,” CSIRO Data Access Portal, Sep. 20, 2021. [Online]. Available: https://doi.org/10.25919/45t8-p065.
+
+
+

recognition/prostate2dnet_46978116/dataset.py

@@ -0,0 +1,106 @@
+
+import numpy as np
+import nibabel as nib
+import os
+import torch
+from tqdm import tqdm
+from torch.utils.data import Dataset
+from skimage.transform import resize
+import torchvision.transforms as transforms
+
+
+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 pre - maturely, 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(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
+            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 = 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 torch.tensor(images, dtype = torch.float32), affines
+    else:
+        return torch.tensor(images, dtype = torch.float32)
+
+# Dataset structure for loading images and masks into dataloader
+class ProstateDataset(Dataset):
+    def __init__(self, image_path, mask_path, norm_image=False, transform=None, target_size=(128, 64)):
+        self.transform = transform
+        self.image_path = image_path
+        self.mask_path = mask_path
+
+        # list of paths
+        self.images = os.listdir(self.image_path)
+        self.masks = os.listdir(self.mask_path)
+
+
+        self.target_size = target_size
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        # load with helper 
+        img_pth = os.path.join(self.image_path,self.images[idx])
+        mask_pth= os.path.join(self.mask_path,self.images[idx].replace('case', 'seg'))
+        image = load_data_2D([img_pth], normImage=True)
+        mask = load_data_2D([mask_pth])
+
+                # Apply transformations
+        image = transforms.Resize((256, 256))(image)
+        mask = transforms.Resize((256, 256))(mask)
+
+        mask = mask.long()
+
+
+
+        return image, mask

recognition/prostate2dnet_46978116/images/predictions/test_cross_scores.txt

@@ -0,0 +1,7 @@
+Class 0: 0.9945
+Class 1: 0.9768
+Class 2: 0.8689
+Class 3: 0.6563
+Class 4: 0.3489
+Class 5: 0.2344
+Average Dice Score: 0.6800

recognition/prostate2dnet_46978116/images/predictions/test_dice_scores.txt

@@ -0,0 +1,7 @@
+Class 0: 0.9863
+Class 1: 0.9777
+Class 2: 0.8820
+Class 3: 0.7254
+Class 4: 0.6833
+Class 5: 0.8711
+Average Dice Score: 0.8543

recognition/prostate2dnet_46978116/modules.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as TF
+
+
+class UNet(nn.Module):
+    def __init__(self):
+        super(UNet, self).__init__()
+
+        # Down double conv layers
+        self.down1 = DoubleConv(in_channels=1, out_channels=64)
+        self.down2 = DoubleConv(in_channels=64, out_channels=128)
+        self.down3 = DoubleConv(in_channels=128, out_channels=256)
+
+        # Bottle neck
+        self.down4 = DoubleConv(in_channels=256, out_channels=512)
+
+        #max pool layer
+        self.max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        # Up transpose layers + Double Conv
+        self.up_trans1 = nn.ConvTranspose2d(in_channels=512,out_channels=256,kernel_size=2,stride=2)
+        self.up1 = DoubleConv(512, 256)
+
+
+        self.up_trans2 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=2,stride=2)
+        self.up2 = DoubleConv(256, 128) 
+
+        self.up_trans3 = nn.ConvTranspose2d(in_channels=128,out_channels=64,kernel_size=2,stride=2)
+        self.up3 = DoubleConv(128, 64)
+
+
+        self.final = nn.Conv2d(in_channels=64,out_channels=6,kernel_size=1)
+
+
+    def forward(self, initial):
+
+
+        # down 
+        c1 = self.down1(initial)
+        p1 = self.max_pool(c1)
+
+        c2 = self.down2(p1)
+        p2 = self.max_pool(c2)
+
+        c3 = self.down3(p2)
+        p3 = self.max_pool(c3)
+
+
+
+        c4 = self.down4(p3)
+
+        # upsample
+        t1 = self.up_trans1(c4)
+        d1 = self.up1(torch.cat([t1, c3], 1))
+
+        t2 = self.up_trans2(d1)
+        d2 = self.up2(torch.cat([t2, c2], 1))
+
+        t3 = self.up_trans3(d2)
+        d3 = self.up3(torch.cat([t3, c1], 1))
+
+
+        # output
+        out = self.final(d3)
+        return out
+
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DoubleConv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride = 1, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride= 1, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+
+    def forward(self,x):
+        return self.conv(x)