From 9b22f24e176d16ca591828fe5d6e731579c7a396 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Tue, 22 Oct 2024 13:51:22 +1000
Subject: [PATCH 01/10] Created UNet3D_47473289 folder containing all
 neccessary files for the UNet3D model.

---
 recognition/UNet3D_47473289/README.md  | 1 +
 recognition/UNet3D_47473289/dataset.py | 1 +
 recognition/UNet3D_47473289/modules.py | 1 +
 recognition/UNet3D_47473289/predict.py | 1 +
 recognition/UNet3D_47473289/train.py   | 3 +++
 5 files changed, 7 insertions(+)
 create mode 100644 recognition/UNet3D_47473289/README.md
 create mode 100644 recognition/UNet3D_47473289/dataset.py
 create mode 100644 recognition/UNet3D_47473289/modules.py
 create mode 100644 recognition/UNet3D_47473289/predict.py
 create mode 100644 recognition/UNet3D_47473289/train.py

diff --git a/recognition/UNet3D_47473289/README.md b/recognition/UNet3D_47473289/README.md
new file mode 100644
index 000000000..49019e045
--- /dev/null
+++ b/recognition/UNet3D_47473289/README.md
@@ -0,0 +1 @@
+# Charles Slack COMP3710 Report README
\ No newline at end of file
diff --git a/recognition/UNet3D_47473289/dataset.py b/recognition/UNet3D_47473289/dataset.py
new file mode 100644
index 000000000..5387fb892
--- /dev/null
+++ b/recognition/UNet3D_47473289/dataset.py
@@ -0,0 +1 @@
+# Data loader for leading and preprocessing data.
\ No newline at end of file
diff --git a/recognition/UNet3D_47473289/modules.py b/recognition/UNet3D_47473289/modules.py
new file mode 100644
index 000000000..2ee975302
--- /dev/null
+++ b/recognition/UNet3D_47473289/modules.py
@@ -0,0 +1 @@
+# Source code of the components of the UNet3D Model.
\ No newline at end of file
diff --git a/recognition/UNet3D_47473289/predict.py b/recognition/UNet3D_47473289/predict.py
new file mode 100644
index 000000000..0ea0a3778
--- /dev/null
+++ b/recognition/UNet3D_47473289/predict.py
@@ -0,0 +1 @@
+# Example usage of the trained UNet3D model.
\ No newline at end of file
diff --git a/recognition/UNet3D_47473289/train.py b/recognition/UNet3D_47473289/train.py
new file mode 100644
index 000000000..0aa424cb5
--- /dev/null
+++ b/recognition/UNet3D_47473289/train.py
@@ -0,0 +1,3 @@
+# Source code for training, validating, testing and saving the UNet3D Model.
+from modules import *
+from dataset import *
\ No newline at end of file

From 1638858913faad66f93b5558eecd0abb41fea618 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Fri, 25 Oct 2024 11:09:48 +1000
Subject: [PATCH 02/10] Implemented Improved UNet3D Module

---
 recognition/UNet3D_47473289/modules.py | 114 ++++++++++++++++++++++++-
 1 file changed, 113 insertions(+), 1 deletion(-)

diff --git a/recognition/UNet3D_47473289/modules.py b/recognition/UNet3D_47473289/modules.py
index 2ee975302..dc0e05569 100644
--- a/recognition/UNet3D_47473289/modules.py
+++ b/recognition/UNet3D_47473289/modules.py
@@ -1 +1,113 @@
-# Source code of the components of the UNet3D Model.
\ No newline at end of file
+# Source code of the components of the UNet3D Model.
+
+import torch
+import torch.nn as nn
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(ConvBlock, self).__init__()
+        self.block = nn.Sequential(
+            nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm3d(out_channels),
+            nn.ReLU(),
+            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm3d(out_channels),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+class UpConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(UpConv, self).__init__()
+        self.up = nn.ConvTranspose3d(in_channels, out_channels, kernel_size=2, stride=2)
+
+    def forward(self, x):
+        return self.up(x)
+
+class ImprovedUNet3D(nn.Module):
+    def __init__(self, in_channels=1, out_channels=1):
+        super(ImprovedUNet3D, self).__init__()
+        print("UNET INIT")
+        # Contracting path
+        self.conv1 = ConvBlock(in_channels, 64)
+        self.pool1 = nn.MaxPool3d(2)
+
+        self.conv2 = ConvBlock(64, 128)
+        self.pool2 = nn.MaxPool3d(2)
+
+        self.conv3 = ConvBlock(128, 256)
+        self.pool3 = nn.MaxPool3d(2)
+
+        self.conv4 = ConvBlock(256, 512)
+        self.pool4 = nn.MaxPool3d(2)
+
+        self.bottleneck = ConvBlock(512, 1024)
+
+        # Expansive path
+        self.upconv4 = UpConv(1024, 512)
+        self.conv_up4 = ConvBlock(1024, 512)
+
+        self.upconv3 = UpConv(512, 256)
+        self.conv_up3 = ConvBlock(512, 256)
+
+        self.upconv2 = UpConv(256, 128)
+        self.conv_up2 = ConvBlock(256, 128)
+
+        self.upconv1 = UpConv(128, 64)
+        self.conv_up1 = ConvBlock(128, 64)
+
+        self.final_conv = nn.Conv3d(64, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        # Contracting path
+        x1 = self.conv1(x)
+        #print("CONV1DONE")
+        p1 = self.pool1(x1)
+
+        x2 = self.conv2(p1)
+        #print("2")
+        p2 = self.pool2(x2)
+
+        x3 = self.conv3(p2)
+        #print("3")
+        p3 = self.pool3(x3)
+
+        x4 = self.conv4(p3)
+        #print("4")
+        p4 = self.pool4(x4)
+
+        # Bottleneck
+        bn = self.bottleneck(p4)
+        #print("BN")
+
+        # Expansive path
+        u4 = self.upconv4(bn)
+        #print("upconv1")
+        u4 = torch.cat([u4, x4], dim=1)
+        u4 = self.conv_up4(u4)
+        #print("conv up 1")
+
+        u3 = self.upconv3(u4)
+        #print("upconv2")
+        u3 = torch.cat([u3, x3], dim=1)
+        
+        u3 = self.conv_up3(u3)
+        #print("conv up 2")
+
+        u2 = self.upconv2(u3)
+        #print("upconv3")
+        u2 = torch.cat([u2, x2], dim=1)
+        
+        u2 = self.conv_up2(u2)
+        #print("conv up 3")
+
+        u1 = self.upconv1(u2)
+        #print("upconv4")
+        u1 = torch.cat([u1, x1], dim=1)
+        u1 = self.conv_up1(u1)
+        #print("conv up 4")
+
+        return self.final_conv(u1)
+    

From 9559713b7470b66d64e750f31b744d9da56aec60 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Fri, 25 Oct 2024 11:10:10 +1000
Subject: [PATCH 03/10] Implemented DSC Module

---
 recognition/UNet3D_47473289/modules.py | 38 ++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/recognition/UNet3D_47473289/modules.py b/recognition/UNet3D_47473289/modules.py
index dc0e05569..6bdb60a2d 100644
--- a/recognition/UNet3D_47473289/modules.py
+++ b/recognition/UNet3D_47473289/modules.py
@@ -111,3 +111,41 @@ def forward(self, x):
 
         return self.final_conv(u1)
     
+class DSC(nn.Module):
+    """
+        Implementation of Dice-Sorensen Coefficient loss
+        2 * |predict * target|/|predict|+|target|
+    """
+    def __init__(self):
+        super(DSC, self).__init__()
+        # Smooth to avoid 0 division
+        self.smooth = 1.0 
+
+    def forward(self, y_pred, y_true):
+        assert y_pred.size() == y_true.size()
+        # Flattening tensors
+        #y_pred = y_pred[:, 0].contiguous().view(-1)
+        #y_true = y_true[:, 0].contiguous().view(-1)
+        intersection = (y_pred * y_true).sum(dim=[2,3,4])
+        # Calculate DSC
+        dsc = (2. * intersection + self.smooth) / (
+            y_pred.sum(dim=[2,3,4]) + y_true.sum(dim=[2,3,4]) + self.smooth
+        )
+        """
+        # Returning DSC Loss
+
+        assert y_pred.size() == y_true.size(), "Input sizes must be the same"
+
+        # Compute intersection and sums directly on the original tensors
+        intersection = (y_pred * y_true).sum(dim=[2, 3, 4])  # Sum along spatial dimensions
+
+        # Compute the Dice Score
+        dsc = (2. * intersection + self.smooth) / (
+            y_pred.sum(dim=[2, 3, 4]) + y_true.sum(dim=[2, 3, 4]) + self.smooth
+        ) READ THISSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSS SUM ALONG SPATIAL DIMENSIONS
+
+        # Average over the batch
+        dsc = dsc.mean()"""
+        dsc = dsc.mean()
+        print(1. - dsc)
+        return 1. - dsc

From 9e0061c27019944e7a820104da8f65265c6b72f2 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Fri, 25 Oct 2024 11:10:55 +1000
Subject: [PATCH 04/10] Implemented loading 3d data from nifti files

---
 recognition/UNet3D_47473289/dataset.py | 118 ++++++++++++++++++++++++-
 1 file changed, 117 insertions(+), 1 deletion(-)

diff --git a/recognition/UNet3D_47473289/dataset.py b/recognition/UNet3D_47473289/dataset.py
index 5387fb892..a215488d3 100644
--- a/recognition/UNet3D_47473289/dataset.py
+++ b/recognition/UNet3D_47473289/dataset.py
@@ -1 +1,117 @@
-# Data loader for leading and preprocessing data.
\ No newline at end of file
+# Data loader for leading and preprocessing data.
+from torch.utils.data import DataLoader, TensorDataset, Dataset
+import torch
+import nibabel as nib
+import numpy as np
+import tqdm as tqdm
+from sklearn.model_selection import train_test_split
+from scipy.ndimage import zoom
+import torch.nn as nn
+import os
+import torchvision.transforms as transforms
+import glob
+
+
+
+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
+
+transform = transforms.Compose([
+    transforms.Resize((128, 128)),
+   # transforms.Normalize(0.5, 0.5),
+])
+l_transform = transforms.Compose([
+    transforms.Resize((128, 128)),
+])
+
+
+def load_data_3D(imageNames, normImage=False , categorical=False, dtype=np.float32, 
+                    getAffines=False, orient=False, early_stop=False):
+    '''
+    Load medical image data from names , cases list provided into a list for each .
+
+    This function pre - allocates 5 D arrays for conv3d to avoid excessive memory ↘
+    usage .
+
+    normImage : bool ( normalise the image 0.0 -1.0)
+    orient : Apply orientation and resample image ? Good for images with large slice ↘
+    thickness or anisotropic resolution
+    dtype : Type of the data . If dtype = np . uint8 , it is assumed that the data is ↘
+    labels
+    early_stop : Stop loading pre - maturely ? Leaves arrays mostly empty , for quick ↘
+    loading and testing scripts .
+    '''
+    affines = []
+
+
+    # ~ interp = ' continuous '
+    interp = 'linear'
+    if dtype == np.uint8: # assume labels
+        interp = 'nearest'
+
+    # get fixed size
+    num = len(imageNames)
+
+        # ~ testResultName = "oriented.nii.gz"
+        # ~ niftiImage.to_filename(testResultName)
+    first_case = nib.load(imageNames[0]).get_fdata(caching='unchanged')
+
+    if len(first_case.shape) == 4:
+        first_case = first_case [:,:,:,0] # sometimes extra dims , remove
+    if categorical:
+        first_case = to_channels(first_case, dtype=dtype)
+        rows, cols, depth, channels = first_case.shape
+        images = np.zeros((num, rows, cols, depth, channels) , dtype=dtype)
+
+    else:
+        rows, cols, depth = first_case.shape
+        images = np.zeros((num, rows, cols, depth), dtype=dtype)
+
+
+    for i ,inName in enumerate(tqdm.tqdm(imageNames)):
+        niftiImage = nib.load(inName)
+        if len(first_case.shape) == 3:
+            first_case = np.expand_dims(first_case, axis=0)
+
+        inImage = niftiImage.get_fdata(caching='unchanged') # read disk only
+        affine = niftiImage.affine
+        
+        if len (inImage.shape) == 4:
+            inImage = inImage [:,:,:,0] # sometimes extra dims in HipMRI_study data
+
+        inImage = inImage [:,:,:depth] # clip slices
+        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
+            images [i,:inImage.shape[0],:inImage.shape[1],:inImage.shape [2],:inImage.shape[3]] = inImage # with pad
+
+
+      
+        else :
+            # ~ images [i,:,:,:] = inImage
+
+            images [i,:inImage.shape[0],:inImage.shape[1],:inImage.shape[2]] = inImage # with pad
+            
+        affines.append(affine)
+        if i > 20 and early_stop:
+            print("STOPPED EARly")
+            break
+
+    if getAffines:
+        return images, affines
+    else:
+        print("Returned images")
+        return images
+

From aaa8d6f213a4bcb66bdd4971e11fa4b7c1f0ec09 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Fri, 25 Oct 2024 11:11:29 +1000
Subject: [PATCH 05/10] Implemented custom dataset for loading labels and 3d
 images

---
 recognition/UNet3D_47473289/dataset.py | 38 ++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/recognition/UNet3D_47473289/dataset.py b/recognition/UNet3D_47473289/dataset.py
index a215488d3..f8a705c88 100644
--- a/recognition/UNet3D_47473289/dataset.py
+++ b/recognition/UNet3D_47473289/dataset.py
@@ -115,3 +115,41 @@ def load_data_3D(imageNames, normImage=False , categorical=False, dtype=np.float
         print("Returned images")
         return images
 
+class MyCustomDataset(Dataset):
+    def __init__(self):
+        # load all nii handle in a list
+        #self.image_paths = glob.glob(f'{"/Users/charl/Documents/3710Report/PatternAnalysis-2024/recognition/semantic_MRs_anon"}/**/*.nii.gz', recursive=True)
+        #self.label_paths = glob.glob(f'{"/Users/charl/Documents/3710Report/PatternAnalysis-2024/recognition/semantic_labels_only"}/**/*.nii.gz', recursive=True)
+        self.label_paths = glob.glob(f'{"/home/groups/comp3710/HipMRI_Study_open/semantic_labels_only"}/**/*.nii.gz', recursive=True)
+        self.image_paths = glob.glob(f'{"/home/groups/comp3710/HipMRI_Study_open/semantic_MRs"}/**/*.nii.gz', recursive=True)
+        print(len(self.image_paths))
+        print(len(self.label_paths))
+  
+ 
+        self.resize = transforms.Compose([transforms.Resize((256, 128))])
+        self.up = torch.nn.Upsample(size=(128,128,128))
+
+        self.classes = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5}
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, i):
+        image = load_data_3D([self.image_paths[i]],early_stop=True)
+        label = load_data_3D([self.label_paths[i]],early_stop=True)
+        print(image.shape)
+        image = torch.tensor(image).float().unsqueeze(1)
+        label = torch.tensor([self.classes.get(cla.item(), 0) for cla in label.flatten()]).reshape(label.shape)
+        label = nn.functional.one_hot(label.squeeze(0), num_classes=6).permute(3,1,0,2).float()
+        label = label.unsqueeze(0)
+        image = self.up(image).squeeze(0)
+        label = self.up(label).squeeze(0)
+        print(image.shape)
+        print(label.shape)
+
+        return image, label
+
+
+
+dataset = MyCustomDataset()
+print(dataset[10])
+print(dataset[10][1].shape)
\ No newline at end of file

From 32a20dc240c66bf12a78b146e227e96410501657 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Fri, 25 Oct 2024 11:13:26 +1000
Subject: [PATCH 06/10] Implemented model training and validation using adam
 optimizer and a sequential scheduler, loading data from custom dataset into
 train validate and test sets

---
 recognition/UNet3D_47473289/train.py | 94 +++++++++++++++++++++++++++-
 1 file changed, 93 insertions(+), 1 deletion(-)

diff --git a/recognition/UNet3D_47473289/train.py b/recognition/UNet3D_47473289/train.py
index 0aa424cb5..7ae682786 100644
--- a/recognition/UNet3D_47473289/train.py
+++ b/recognition/UNet3D_47473289/train.py
@@ -1,3 +1,95 @@
 # Source code for training, validating, testing and saving the UNet3D Model.
 from modules import *
-from dataset import *
\ No newline at end of file
+from dataset import *
+import torch
+import torch.optim as optim
+from sklearn.model_selection import train_test_split
+import time
+import glob
+torch.cuda.empty_cache()
+
+
+cust_dataset = MyCustomDataset()
+train, val, test = torch.utils.data.random_split(cust_dataset, [0.2, 0.1, 0.7])
+train = DataLoader(train)
+val = DataLoader(val)
+test = DataLoader(test)
+model = ImprovedUNet3D(in_channels=1, out_channels=6).cuda()
+optimizer = optim.Adam(model.parameters(), lr=1e-4)
+
+
+criterion = DSC()
+
+learn_rate = 0.1
+# Alternates minimum and maximum learning rate
+l_sched_1 = optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.005, max_lr = learn_rate, step_size_up=15, step_size_down=15, mode='triangular', verbose= False)
+# Linearly reduces learning rate
+l_sched_2 = optim.lr_scheduler.LinearLR(optimizer , start_factor=0.005/learn_rate, end_factor=0.005/5, verbose=False)
+# Sequentially applies sched 1 and 2 
+scheduler = optim.lr_scheduler.SequentialLR(optimizer, schedulers=[l_sched_1, l_sched_2], milestones=[30])
+
+
+# No. of Epochs (Number of times all training data is run)
+n_epochs = 1
+train_losses = []
+print(">> Training <<")
+start = time.time()
+for epoch in range(n_epochs):
+    epstart = time.time()
+    # Set training mode
+    model.train()
+    running_loss = 0.0
+    print(len(train))
+    for i, (inputs, labels) in enumerate(train):
+        itstart = time.time()
+
+        print(inputs.shape)
+        print(labels.shape)
+        print(labels.size)
+        print(inputs.size)
+        #if (i > 2):
+         #   break
+        inputs = inputs.cuda()
+        labels = labels.cuda()
+        
+        optimizer.zero_grad()
+        #print("ZERO")
+        output = model(inputs)
+        #print("OUTPUT")
+        # Using dsc class to calculate loss
+        print(output.shape)
+        loss = criterion(output, labels)
+        # Backward pass
+        #print("CRIT")
+        
+        loss.backward()
+        #print("BACK")
+        optimizer.step()
+        print("STEP")
+        itend = time.time()
+        print(f"Iteration #{i} took {itend - itstart} seconds")
+        running_loss += loss.item()
+    scheduler.step()
+    print("SCEDSTEP")
+    # Set evaluation mode for validation
+    train_loss = running_loss / len(train)
+    train_losses.append(train_loss)
+    model.eval()
+    scores = []
+    with torch.no_grad():
+        for i, (val_inputs, val_labels, ) in enumerate(val):
+            val_inputs = val_inputs.cuda()
+            val_labels = val_labels.cuda()
+            output = model(val_inputs)
+            # Calculating Dice score for each pass
+            score = 1 - criterion(output, val_labels)
+            scores.append(score)
+    # Average validation score for all batches
+    mean_score = sum(scores) / len(scores)
+    epend = time.time()
+    print(f"Epoch {epoch + 1}/{n_epochs}, Training Loss: {loss.item()}, Validation DSC: {mean_score}")
+    print(f"Epoch {epoch + 1}/{n_epochs} took {epend - epstart} seconds")
+
+end = time.time()
+total_time = end - start
+print(f"Training and Validation took {total_time} seconds or {total_time/60} minutes.")

From b6b87511e28806f7f794d8069754d25c6a3f1a82 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Fri, 25 Oct 2024 11:43:43 +1000
Subject: [PATCH 07/10] Filled README.md file with Overview, Model description
 and dependencies

---
 recognition/UNet3D_47473289/README.md | 29 ++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/recognition/UNet3D_47473289/README.md b/recognition/UNet3D_47473289/README.md
index 49019e045..b2d2457c5 100644
--- a/recognition/UNet3D_47473289/README.md
+++ b/recognition/UNet3D_47473289/README.md
@@ -1 +1,28 @@
-# Charles Slack COMP3710 Report README
\ No newline at end of file
+# Improved 3D Unet for Prostate MRI Segmentation
+
+## Overview
+
+The implementation of an improved 3D UNet segmenting the (downsampled) Prostate 3D data set with all labels having a minimum Dice similarity coefficient of 0.7 on the test set.The data is augmented using the appropriate transforms in PyTorch and is loaded from a Nifti file format ".nii.gz"
+
+
+## Improved UNet3D Model
+
+The Improved 3D UNet model is a CNN for 3D image segmentation.
+The model is based on the keras brats UNet model using an encoder and decoder with a bottleneck inbetween. The encoder applies convolutions to the input images and reduces its dimensions to identify patterns in the images. The decoder applies up-convolutions to the downsampled images while concatenating at each step  with the relative encoder convolution in its respective layer. The bottleneck is the connection between the encoder and the decoder.
+
+The dataset is split into training data, validation data and testing data. 
+These datasets are a subset of the total data loaded, split into 70% training, 20% validation and 10% testing.
+
+Training the model has the goal of minimising the DSC loss on the training data, used to maximise
+the DSC in evaluation mode on the validation set. Appropriate transformations are applied to the 
+data to maximise the speed of training.
+
+## Dependencies
+
+PyTorch version: 2.4.0
+Numpy version: 1.26.4
+Tqdm version: 4.66.4
+Nibabel version: 5.3.1
+Python version: 3.12.3
+
+## Results
\ No newline at end of file

From 928d0765285c170b1d591cbba31ef3b8e7974b07 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Fri, 25 Oct 2024 11:45:49 +1000
Subject: [PATCH 08/10] Cleaned dataset style and removed unused code

---
 recognition/UNet3D_47473289/dataset.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/recognition/UNet3D_47473289/dataset.py b/recognition/UNet3D_47473289/dataset.py
index f8a705c88..48f6739a8 100644
--- a/recognition/UNet3D_47473289/dataset.py
+++ b/recognition/UNet3D_47473289/dataset.py
@@ -7,10 +7,17 @@
 from sklearn.model_selection import train_test_split
 from scipy.ndimage import zoom
 import torch.nn as nn
-import os
 import torchvision.transforms as transforms
 import glob
-
+import sklearn
+import scipy
+import sys
+
+print("PyTorch version:", torch.__version__)
+print("Numpy version:", np.__version__)
+print("Tqdm version:", tqdm.__version__)
+print("Nibabel version:", nib.__version__)
+print("Python version:", sys.version)
 
 
 def to_channels ( arr : np . ndarray , dtype = np . uint8 ) -> np . ndarray :

From ab41535de5c0a9fc5859fc3f896200743a272148 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Fri, 25 Oct 2024 11:48:27 +1000
Subject: [PATCH 09/10] Cleaned style and removed unused or unnecessary code

---
 recognition/UNet3D_47473289/dataset.py | 26 +-----------------
 recognition/UNet3D_47473289/modules.py | 37 ++------------------------
 recognition/UNet3D_47473289/train.py   | 25 ++++-------------
 3 files changed, 8 insertions(+), 80 deletions(-)

diff --git a/recognition/UNet3D_47473289/dataset.py b/recognition/UNet3D_47473289/dataset.py
index 48f6739a8..98b0afc3e 100644
--- a/recognition/UNet3D_47473289/dataset.py
+++ b/recognition/UNet3D_47473289/dataset.py
@@ -9,15 +9,6 @@
 import torch.nn as nn
 import torchvision.transforms as transforms
 import glob
-import sklearn
-import scipy
-import sys
-
-print("PyTorch version:", torch.__version__)
-print("Numpy version:", np.__version__)
-print("Tqdm version:", tqdm.__version__)
-print("Nibabel version:", nib.__version__)
-print("Python version:", sys.version)
 
 
 def to_channels ( arr : np . ndarray , dtype = np . uint8 ) -> np . ndarray :
@@ -29,14 +20,6 @@ def to_channels ( arr : np . ndarray , dtype = np . uint8 ) -> np . ndarray :
 
     return res
 
-transform = transforms.Compose([
-    transforms.Resize((128, 128)),
-   # transforms.Normalize(0.5, 0.5),
-])
-l_transform = transforms.Compose([
-    transforms.Resize((128, 128)),
-])
-
 
 def load_data_3D(imageNames, normImage=False , categorical=False, dtype=np.float32, 
                     getAffines=False, orient=False, early_stop=False):
@@ -56,7 +39,6 @@ def load_data_3D(imageNames, normImage=False , categorical=False, dtype=np.float
     '''
     affines = []
 
-
     # ~ interp = ' continuous '
     interp = 'linear'
     if dtype == np.uint8: # assume labels
@@ -103,12 +85,9 @@ def load_data_3D(imageNames, normImage=False , categorical=False, dtype=np.float
             inImage = to_channels(inImage, dtype=dtype)
             # ~ images [i,:,:,:,:] = inImage
             images [i,:inImage.shape[0],:inImage.shape[1],:inImage.shape [2],:inImage.shape[3]] = inImage # with pad
-
-
       
         else :
             # ~ images [i,:,:,:] = inImage
-
             images [i,:inImage.shape[0],:inImage.shape[1],:inImage.shape[2]] = inImage # with pad
             
         affines.append(affine)
@@ -129,11 +108,8 @@ def __init__(self):
         #self.label_paths = glob.glob(f'{"/Users/charl/Documents/3710Report/PatternAnalysis-2024/recognition/semantic_labels_only"}/**/*.nii.gz', recursive=True)
         self.label_paths = glob.glob(f'{"/home/groups/comp3710/HipMRI_Study_open/semantic_labels_only"}/**/*.nii.gz', recursive=True)
         self.image_paths = glob.glob(f'{"/home/groups/comp3710/HipMRI_Study_open/semantic_MRs"}/**/*.nii.gz', recursive=True)
-        print(len(self.image_paths))
-        print(len(self.label_paths))
   
- 
-        self.resize = transforms.Compose([transforms.Resize((256, 128))])
+        self.transform = transforms.Compose([transforms.Resize((256, 128))])
         self.up = torch.nn.Upsample(size=(128,128,128))
 
         self.classes = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5}
diff --git a/recognition/UNet3D_47473289/modules.py b/recognition/UNet3D_47473289/modules.py
index 6bdb60a2d..a3e95be2f 100644
--- a/recognition/UNet3D_47473289/modules.py
+++ b/recognition/UNet3D_47473289/modules.py
@@ -61,53 +61,36 @@ def __init__(self, in_channels=1, out_channels=1):
         self.final_conv = nn.Conv3d(64, out_channels, kernel_size=1)
 
     def forward(self, x):
-        # Contracting path
         x1 = self.conv1(x)
-        #print("CONV1DONE")
         p1 = self.pool1(x1)
 
         x2 = self.conv2(p1)
-        #print("2")
         p2 = self.pool2(x2)
 
         x3 = self.conv3(p2)
-        #print("3")
         p3 = self.pool3(x3)
 
         x4 = self.conv4(p3)
-        #print("4")
         p4 = self.pool4(x4)
 
-        # Bottleneck
         bn = self.bottleneck(p4)
-        #print("BN")
 
-        # Expansive path
         u4 = self.upconv4(bn)
-        #print("upconv1")
         u4 = torch.cat([u4, x4], dim=1)
         u4 = self.conv_up4(u4)
-        #print("conv up 1")
+
 
         u3 = self.upconv3(u4)
-        #print("upconv2")
         u3 = torch.cat([u3, x3], dim=1)
-        
         u3 = self.conv_up3(u3)
-        #print("conv up 2")
 
         u2 = self.upconv2(u3)
-        #print("upconv3")
         u2 = torch.cat([u2, x2], dim=1)
-        
         u2 = self.conv_up2(u2)
-        #print("conv up 3")
 
         u1 = self.upconv1(u2)
-        #print("upconv4")
         u1 = torch.cat([u1, x1], dim=1)
         u1 = self.conv_up1(u1)
-        #print("conv up 4")
 
         return self.final_conv(u1)
     
@@ -123,29 +106,13 @@ def __init__(self):
 
     def forward(self, y_pred, y_true):
         assert y_pred.size() == y_true.size()
-        # Flattening tensors
-        #y_pred = y_pred[:, 0].contiguous().view(-1)
-        #y_true = y_true[:, 0].contiguous().view(-1)
+
         intersection = (y_pred * y_true).sum(dim=[2,3,4])
         # Calculate DSC
         dsc = (2. * intersection + self.smooth) / (
             y_pred.sum(dim=[2,3,4]) + y_true.sum(dim=[2,3,4]) + self.smooth
         )
-        """
-        # Returning DSC Loss
-
-        assert y_pred.size() == y_true.size(), "Input sizes must be the same"
-
-        # Compute intersection and sums directly on the original tensors
-        intersection = (y_pred * y_true).sum(dim=[2, 3, 4])  # Sum along spatial dimensions
-
-        # Compute the Dice Score
-        dsc = (2. * intersection + self.smooth) / (
-            y_pred.sum(dim=[2, 3, 4]) + y_true.sum(dim=[2, 3, 4]) + self.smooth
-        ) READ THISSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSSS SUM ALONG SPATIAL DIMENSIONS
 
-        # Average over the batch
-        dsc = dsc.mean()"""
         dsc = dsc.mean()
         print(1. - dsc)
         return 1. - dsc
diff --git a/recognition/UNet3D_47473289/train.py b/recognition/UNet3D_47473289/train.py
index 7ae682786..2bac6fbda 100644
--- a/recognition/UNet3D_47473289/train.py
+++ b/recognition/UNet3D_47473289/train.py
@@ -41,36 +41,21 @@
     running_loss = 0.0
     print(len(train))
     for i, (inputs, labels) in enumerate(train):
-        itstart = time.time()
 
-        print(inputs.shape)
-        print(labels.shape)
-        print(labels.size)
-        print(inputs.size)
-        #if (i > 2):
-         #   break
         inputs = inputs.cuda()
         labels = labels.cuda()
-        
+
         optimizer.zero_grad()
-        #print("ZERO")
         output = model(inputs)
-        #print("OUTPUT")
-        # Using dsc class to calculate loss
-        print(output.shape)
+
         loss = criterion(output, labels)
-        # Backward pass
-        #print("CRIT")
-        
+
         loss.backward()
-        #print("BACK")
+
         optimizer.step()
-        print("STEP")
-        itend = time.time()
-        print(f"Iteration #{i} took {itend - itstart} seconds")
+
         running_loss += loss.item()
     scheduler.step()
-    print("SCEDSTEP")
     # Set evaluation mode for validation
     train_loss = running_loss / len(train)
     train_losses.append(train_loss)

From f3c41e678e139120ced6a5f5982193adb631d070 Mon Sep 17 00:00:00 2001
From: charlieeslack <charlieslack22@gmail.com>
Date: Sun, 27 Oct 2024 11:45:18 +1000
Subject: [PATCH 10/10] Final tests run ready for pull request

---
 recognition/UNet3D_47473289/README.md  |  5 +-
 recognition/UNet3D_47473289/dataset.py |  6 +--
 recognition/UNet3D_47473289/predict.py | 65 +++++++++++++++++++++++++-
 recognition/UNet3D_47473289/train.py   | 10 ++--
 4 files changed, 77 insertions(+), 9 deletions(-)

diff --git a/recognition/UNet3D_47473289/README.md b/recognition/UNet3D_47473289/README.md
index b2d2457c5..b51351991 100644
--- a/recognition/UNet3D_47473289/README.md
+++ b/recognition/UNet3D_47473289/README.md
@@ -24,5 +24,8 @@ Numpy version: 1.26.4
 Tqdm version: 4.66.4
 Nibabel version: 5.3.1
 Python version: 3.12.3
+Matplotlib version: 3.8.4
 
-## Results
\ No newline at end of file
+## Results
+
+Average DSC across classes = 0.6122
\ No newline at end of file
diff --git a/recognition/UNet3D_47473289/dataset.py b/recognition/UNet3D_47473289/dataset.py
index 98b0afc3e..d8f9fb3f9 100644
--- a/recognition/UNet3D_47473289/dataset.py
+++ b/recognition/UNet3D_47473289/dataset.py
@@ -4,10 +4,9 @@
 import nibabel as nib
 import numpy as np
 import tqdm as tqdm
-from sklearn.model_selection import train_test_split
-from scipy.ndimage import zoom
+#from sklearn.model_selection import train_test_split
+#from scipy.ndimage import zoom
 import torch.nn as nn
-import torchvision.transforms as transforms
 import glob
 
 
@@ -109,7 +108,6 @@ def __init__(self):
         self.label_paths = glob.glob(f'{"/home/groups/comp3710/HipMRI_Study_open/semantic_labels_only"}/**/*.nii.gz', recursive=True)
         self.image_paths = glob.glob(f'{"/home/groups/comp3710/HipMRI_Study_open/semantic_MRs"}/**/*.nii.gz', recursive=True)
   
-        self.transform = transforms.Compose([transforms.Resize((256, 128))])
         self.up = torch.nn.Upsample(size=(128,128,128))
 
         self.classes = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5}
diff --git a/recognition/UNet3D_47473289/predict.py b/recognition/UNet3D_47473289/predict.py
index 0ea0a3778..a489fa2c6 100644
--- a/recognition/UNet3D_47473289/predict.py
+++ b/recognition/UNet3D_47473289/predict.py
@@ -1 +1,64 @@
-# Example usage of the trained UNet3D model.
\ No newline at end of file
+# Example usage of the trained UNet3D model.
+import torch
+import matplotlib.pyplot as plt
+
+from skimage.util import montage
+import numpy as np
+
+
+#dataset = MyCustomDataset()
+#_, test = torch.utils.data.random_split(cust_dataset, [0.9, 0.1])
+#test_loader = DataLoader(test)
+#model = ImprovedUNet3D(in_channels=1, out_channels=6).cuda()
+#model.load_state_dict(torch.load('./new_folder/model.pth'))
+def predict(model, test_loader):
+    model.eval()
+    with torch.no_grad():
+        for test_images, test_masks in test_loader:
+            # moves test_images to be on gpu
+            vis = Visulizer()
+            vis.visualize(test_images, test_masks)
+            
+            test_images = test_images.cuda()
+            test_outputs = model(test_images)
+            predicted_masks = torch.argmax(test_outputs, dim=1)
+
+            plt.figure(figsize=(10, 5))
+            # The base input image
+            plt.subplot(1, 3, 1)
+            plt.imshow(test_images[0].cpu().squeeze(), cmap='gray')
+            plt.title("Input Image")
+
+            # The segmentation mask
+            plt.subplot(1, 3, 2)
+            plt.imshow(test_masks[0].cpu().argmax(dim=0).squeeze(), cmap='gray')
+            plt.title("Segmentation Mask")
+
+            # The predicted Mask
+            plt.subplot(1, 3, 3)
+            plt.imshow(predicted_masks[0].cpu().squeeze(), cmap='gray')
+            plt.title("Predicted Mask")
+
+            # Save image
+            output_path = "output.png"
+            plt.savefig(output_path)
+            plt.show()  
+            break
+            
+
+class Visulizer:
+    def montage_nd(self, image):
+        if len(image.shape)>3:
+            return montage(np.stack([self.montage_nd(img) for img in image],0))
+        elif len(image.shape)==3:
+            return montage(image)
+        else:
+            print('Input less than 3d image, returning original')
+            return image
+
+    def visualize(self, image, mask):
+        fig, axs = plt.subplots(1, 2, figsize = (20, 15 * 2))
+        axs[0].imshow(self.montage_nd(image[...,0]), cmap = 'bone')
+        axs[1].imshow(self.montage_nd(mask[...,0]), cmap = 'bone')
+        plt.show()
+        
diff --git a/recognition/UNet3D_47473289/train.py b/recognition/UNet3D_47473289/train.py
index 2bac6fbda..bac5c9c68 100644
--- a/recognition/UNet3D_47473289/train.py
+++ b/recognition/UNet3D_47473289/train.py
@@ -1,16 +1,18 @@
 # Source code for training, validating, testing and saving the UNet3D Model.
 from modules import *
 from dataset import *
+from predict import *
 import torch
 import torch.optim as optim
 from sklearn.model_selection import train_test_split
 import time
 import glob
+import os
 torch.cuda.empty_cache()
 
 
 cust_dataset = MyCustomDataset()
-train, val, test = torch.utils.data.random_split(cust_dataset, [0.2, 0.1, 0.7])
+train, val, test = torch.utils.data.random_split(cust_dataset, [0.7, 0.2, 0.1])
 train = DataLoader(train)
 val = DataLoader(val)
 test = DataLoader(test)
@@ -30,7 +32,7 @@
 
 
 # No. of Epochs (Number of times all training data is run)
-n_epochs = 1
+n_epochs = 5
 train_losses = []
 print(">> Training <<")
 start = time.time()
@@ -41,7 +43,6 @@
     running_loss = 0.0
     print(len(train))
     for i, (inputs, labels) in enumerate(train):
-
         inputs = inputs.cuda()
         labels = labels.cuda()
 
@@ -75,6 +76,9 @@
     print(f"Epoch {epoch + 1}/{n_epochs}, Training Loss: {loss.item()}, Validation DSC: {mean_score}")
     print(f"Epoch {epoch + 1}/{n_epochs} took {epend - epstart} seconds")
 
+
+predict(model, test)
+
 end = time.time()
 total_time = end - start
 print(f"Training and Validation took {total_time} seconds or {total_time/60} minutes.")