From 9279a8ed770b61e117a6b8a2c31fd8f3ffeec26c Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Wed, 9 Oct 2024 19:08:58 +1000 Subject: [PATCH 01/10] chore: created base files for model --- recognition/modules.py | 46 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 recognition/modules.py diff --git a/recognition/modules.py b/recognition/modules.py new file mode 100644 index 000000000..42d44e0a3 --- /dev/null +++ b/recognition/modules.py @@ -0,0 +1,46 @@ +import torch +import torch.nn as nn + +class encoderBlock(nn.Module): + #convolution+ReLU+maxPool + def __init__(self, in_channels, out_channels): + super(encoderBlock, self).__init__() + self.conv = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size = 3, padding = 1), + nn.ReLU(inplace=True), + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + nn.ReLU(inplace=True), + ) + self.pool = nn.MaxPool2d(kernel_size=2, stride=2) + + def forward(self, x): + x = self.conv(x) + pooled = self.pool(x) + return x, pooled + +def decoderBlock(): + #convolution+concatenation+convolution+ReLU + def __init__(self, inchannels, out_channels): + super(decoderBlock, self).__init() + self.upconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride = 2) + self.conv = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size = 3, padding = 1), + nn.ReLU(inplace=True) + nn.Conv2d(out_channels, out_channels, kernel_size = 3, padding = 1), + nn.ReLU(inplace=True), + ) + + def forward(self, x, skip_features): + x = self.upconv(x) + x = torch.cat((x, skip_features), dim = 1) + x = self.conv(x) + return x + +def bottleneck(): + #convolutionalLayer1 + #convolutionalLayer2 + +def skipConnections(): + +def outputLayer(): + From 5735885159ad8aed5f811e3ce404afe925e8049a Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Wed, 9 Oct 2024 19:28:27 +1000 Subject: [PATCH 02/10] chore: additional files for 2d unet model --- recognition/HipMRI_2dUnet_s4858392/README.md | 0 recognition/dataset.py | 0 recognition/predict.py | 0 recognition/train.py | 0 4 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 recognition/HipMRI_2dUnet_s4858392/README.md create mode 100644 recognition/dataset.py create mode 100644 recognition/predict.py create mode 100644 recognition/train.py diff --git a/recognition/HipMRI_2dUnet_s4858392/README.md b/recognition/HipMRI_2dUnet_s4858392/README.md new file mode 100644 index 000000000..e69de29bb diff --git a/recognition/dataset.py b/recognition/dataset.py new file mode 100644 index 000000000..e69de29bb diff --git a/recognition/predict.py b/recognition/predict.py new file mode 100644 index 000000000..e69de29bb diff --git a/recognition/train.py b/recognition/train.py new file mode 100644 index 000000000..e69de29bb From 2c4561e1104a93120d0be59ace30b1e4a7ff8f88 Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Thu, 31 Oct 2024 19:21:46 +1000 Subject: [PATCH 03/10] feat: unet 2d --- recognition/HipMRI_2dUnet_s4858392/modules.py | 64 +++++++++++ .../HipMRI_2dUnet_s4858392/newdataset.py | 49 ++++++++ recognition/HipMRI_2dUnet_s4858392/train.py | 107 ++++++++++++++++++ 3 files changed, 220 insertions(+) create mode 100644 recognition/HipMRI_2dUnet_s4858392/modules.py create mode 100644 recognition/HipMRI_2dUnet_s4858392/newdataset.py create mode 100644 recognition/HipMRI_2dUnet_s4858392/train.py diff --git a/recognition/HipMRI_2dUnet_s4858392/modules.py b/recognition/HipMRI_2dUnet_s4858392/modules.py new file mode 100644 index 000000000..103e6a055 --- /dev/null +++ b/recognition/HipMRI_2dUnet_s4858392/modules.py @@ -0,0 +1,64 @@ +import torch +import torch.nn as nn +import torchvision.transforms.functional as TF + +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, 3, 1, 1, bias = False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace = True), + nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias = False), + nn.BatchNorm2d(out_channels), + nn.ReLU(inplace = True), + ) + def forward(self, x): + return self.conv(x) + +class uNet(nn.Module): + def __init__( + self, in_channels = 3, out_channels = 1, features=[64,128,256,512], + ): + super(uNet, self).__init__() + self.ups = nn.ModuleList() + self.downs = nn.ModuleList() + self.pool = nn.MaxPool2d(kernel_size=2, stride=2) + + #uNet down part + for feature in features: + self.downs.append(doubleConv(in_channels, feature)) + in_channels = feature + + #uNet up part + for feature in reversed(features): + self.ups.append( + nn.ConvTranspose2d( + feature*2, feature, kernel_size = 2, stride=2 + ) + ) + self.ups.append(doubleConv(feature*2, feature)) + + self.bottleneck = doubleConv(features[-1],features[-1]*2) + self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1) + def forward(self, x): + skipConnections = [] + + for down in self.downs: + x = down(x) + skipConnections.append(x) + x = self.pool(x) + + x = self.bottleneck(x) + skipConnections = skipConnections[::-1] + + for idx in range(0, len(self.ups), 2): + x = self.ups[idx](x) + skipConnection = skipConnections[idx//2] + if x.shape != skipConnection.shape: + x = TF.resize(x, size = skipConnection.shape[2:]) + concatSkip = torch.cat((skipConnection, x), dim=1) + x = self.ups[idx+1](concatSkip) + + return self.final_conv(x) + diff --git a/recognition/HipMRI_2dUnet_s4858392/newdataset.py b/recognition/HipMRI_2dUnet_s4858392/newdataset.py new file mode 100644 index 000000000..014f49c7e --- /dev/null +++ b/recognition/HipMRI_2dUnet_s4858392/newdataset.py @@ -0,0 +1,49 @@ +import torch +from torch.utils.data import Dataset +import numpy as np +import nibabel as nib +from tqdm import tqdm +from niftiload import load_data_2D +import os + +class NIFTIDataset(Dataset): + def __init__(self, imageDir, normImage=False, categorical=False, dtype=np.float32, early_stop=False): + """ + Custom Dataset for loading medical images and corresponding affines. + + Args: + - imageNames (list): List of file paths to medical images. + - normImage (bool): Whether to normalize the image. + - categorical (bool): Whether to convert images to one-hot encoding. + - dtype: Data type of the images. + - early_stop (bool): Whether to stop loading prematurely (for testing). + """ + self.imageNames = [os.path.join(imageDir, fname) for fname in os.listdir(imageDir) if fname.endswith('.nii.gz')] + + # Load images and affines using the load_data_2D function + self.images, self.affines = load_data_2D(self.imageNames, normImage=normImage, + categorical=categorical, dtype=dtype, + getAffines=True, early_stop=early_stop) + + def __len__(self): + """Returns the number of images in the dataset.""" + return len(self.images) + + def __getitem__(self, idx): + """ + Get image and affine at index `idx`. + + Args: + - idx (int): Index of the image to fetch. + + Returns: + - image (Tensor): Image at the specified index. + - affine (ndarray): Affine matrix of the image. + """ + image = self.images[idx] + affine = self.affines[idx] + + # Convert the image to a PyTorch tensor + image_tensor = torch.tensor(image, dtype=torch.float32) + + return image_tensor, affine diff --git a/recognition/HipMRI_2dUnet_s4858392/train.py b/recognition/HipMRI_2dUnet_s4858392/train.py new file mode 100644 index 000000000..b315f67e0 --- /dev/null +++ b/recognition/HipMRI_2dUnet_s4858392/train.py @@ -0,0 +1,107 @@ + +import torch +import albumentations as A +import numpy as np +from albumentations.pytorch import ToTensorV2 +from tqdm import tqdm +import torch.nn as nn +import torch.optim as optim +from dataset import MRIDataset +from torch.utils.data import DataLoader +from modules import uNet +from newdataset import NIFTIDataset + +#Hyperparameters +LEARNING_RATE = 1e-4 +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +BATCH_SIZE = 16 +NUM_EPOCHS = 3 +NUM_WORKERS = 2 +IMAGE_HEIGHT = 128 +IMAGE_WIDTH = 256 +PIN_MEMORY = True +LOAD_MODEL = False + +if __name__ == '__main__': + #images dataset + train_image_dir = "/Users/kylow/Dev/PatternAnalysis-2024/recognition/HipMRI_2dUnet_s4858392/HipMRI_study_keras_slices_data/keras_slices_train" + val_image_dir = "/Users/kylow/Dev/PatternAnalysis-2024/recognition/HipMRI_2dUnet_s4858392/HipMRI_study_keras_slices_data/keras_slices_validate" + + #mask dataset + train_seg_dir = "/Users/kylow/Dev/PatternAnalysis-2024/recognition/HipMRI_2dUnet_s4858392/HipMRI_study_keras_slices_data/keras_slices_seg_train" + val_seg_dir = "//Users/kylow/Dev/PatternAnalysis-2024/recognition/HipMRI_2dUnet_s4858392/HipMRI_study_keras_slices_data/keras_slices_seg_validate" + + train_dataset = NIFTIDataset(imageDir=train_image_dir) + val_dataset = NIFTIDataset(imageDir=val_image_dir) + + train_loader = DataLoader(train_dataset, batch_size=16, shuffle=False, num_workers=4) + val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False, num_workers=4) + + model = uNet(in_channels=1, out_channels=1).to(DEVICE) + optimizer = optim.Adam(model.parameters(), lr= LEARNING_RATE) + criterion = nn.BCEWithLogitsLoss() + + def train(model, loader, optimizer, criterion, devices): + model.train() + running_loss = 0.0 + + for images, masks in loader: + images, masks = images.to(DEVICE), masks.to(DEVICE) + + # Zero the gradients + optimizer.zero_grad() + + # Forward pass + outputs = model(images) + loss = criterion(outputs, masks) + + # Backward pass and optimization + loss.backward() + optimizer.step() + + running_loss += loss.item() + + avg_loss = running_loss / len(loader) + return avg_loss + + def validate(model, loader, criterion, device): + model.eval() + val_loss = 0.0 + with torch.no_grad(): + for images, masks in loader: + images, masks = images.to(device) + outputs = model(images) + loss = criterion(outputs, masks) + val_loss += loss.item() + + avg_loss = val_loss / len(loader) + return avg_loss + + def train(model, loader, optimizer, criterion, device): + model.train() + running_loss = 0.0 + for images, masks in loader: + images, masks = images.to(device), masks.to(device) + + # Zero the gradients + optimizer.zero_grad() + + # Forward pass + outputs = model(images) + loss = criterion(outputs, masks) + + # Backward pass and optimization + loss.backward() + optimizer.step() + + running_loss += loss.item() + + avg_loss = running_loss / len(loader) + return avg_loss + +# Training loop + for epoch in range(NUM_EPOCHS): + train_loss = train(model, train_loader, optimizer, criterion, DEVICE) + val_loss = validate(model, val_loader, criterion, DEVICE) + + print(f"Epoch {epoch+1}/{NUM_EPOCHS}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}") \ No newline at end of file From 6e4d2db435450986614bb27dee3f36a14144d504 Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Fri, 1 Nov 2024 23:08:08 +1000 Subject: [PATCH 04/10] chore: fixed known bugs and housekeeping --- recognition/HipMRI_2dUnet_s4858392/train.py | 76 +++++++++------------ recognition/dataset.py | 0 recognition/modules.py | 46 ------------- recognition/predict.py | 0 4 files changed, 32 insertions(+), 90 deletions(-) delete mode 100644 recognition/dataset.py delete mode 100644 recognition/modules.py delete mode 100644 recognition/predict.py diff --git a/recognition/HipMRI_2dUnet_s4858392/train.py b/recognition/HipMRI_2dUnet_s4858392/train.py index b315f67e0..0fcacc60b 100644 --- a/recognition/HipMRI_2dUnet_s4858392/train.py +++ b/recognition/HipMRI_2dUnet_s4858392/train.py @@ -1,4 +1,3 @@ - import torch import albumentations as A import numpy as np @@ -6,10 +5,10 @@ from tqdm import tqdm import torch.nn as nn import torch.optim as optim -from dataset import MRIDataset from torch.utils.data import DataLoader from modules import uNet from newdataset import NIFTIDataset +import torch.nn.functional as F #Hyperparameters LEARNING_RATE = 1e-4 @@ -24,54 +23,40 @@ if __name__ == '__main__': #images dataset - train_image_dir = "/Users/kylow/Dev/PatternAnalysis-2024/recognition/HipMRI_2dUnet_s4858392/HipMRI_study_keras_slices_data/keras_slices_train" - val_image_dir = "/Users/kylow/Dev/PatternAnalysis-2024/recognition/HipMRI_2dUnet_s4858392/HipMRI_study_keras_slices_data/keras_slices_validate" + train_image_dir = "/home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_train" + val_image_dir = "/home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_validate" #mask dataset - train_seg_dir = "/Users/kylow/Dev/PatternAnalysis-2024/recognition/HipMRI_2dUnet_s4858392/HipMRI_study_keras_slices_data/keras_slices_seg_train" - val_seg_dir = "//Users/kylow/Dev/PatternAnalysis-2024/recognition/HipMRI_2dUnet_s4858392/HipMRI_study_keras_slices_data/keras_slices_seg_validate" + train_seg_dir = "/home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_seg_train" + val_seg_dir = "/home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_seg_validate" train_dataset = NIFTIDataset(imageDir=train_image_dir) val_dataset = NIFTIDataset(imageDir=val_image_dir) - train_loader = DataLoader(train_dataset, batch_size=16, shuffle=False, num_workers=4) - val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False, num_workers=4) + train_loader = DataLoader(train_dataset, batch_size=1, shuffle=False, num_workers=1) + val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1) model = uNet(in_channels=1, out_channels=1).to(DEVICE) optimizer = optim.Adam(model.parameters(), lr= LEARNING_RATE) criterion = nn.BCEWithLogitsLoss() - def train(model, loader, optimizer, criterion, devices): - model.train() - running_loss = 0.0 - - for images, masks in loader: - images, masks = images.to(DEVICE), masks.to(DEVICE) - - # Zero the gradients - optimizer.zero_grad() - - # Forward pass - outputs = model(images) - loss = criterion(outputs, masks) - - # Backward pass and optimization - loss.backward() - optimizer.step() - - running_loss += loss.item() - - avg_loss = running_loss / len(loader) - return avg_loss - def validate(model, loader, criterion, device): model.eval() val_loss = 0.0 + total_dice_score = 0.0 with torch.no_grad(): - for images, masks in loader: - images, masks = images.to(device) + for batch in loader: + + images, masks = batch + images = images.to(device) + masks =masks.to(device) + if images.dim() ==3: + images = images.unsqueeze(1) + if masks.dim() ==3: + masks = masks.unsqueeze(1) outputs = model(images) - loss = criterion(outputs, masks) + masks_resized = F.interpolate(masks, size=(256, 128), mode="nearest") + loss = criterion(outputs, masks_resized) val_loss += loss.item() avg_loss = val_loss / len(loader) @@ -81,27 +66,30 @@ def train(model, loader, optimizer, criterion, device): model.train() running_loss = 0.0 for images, masks in loader: + if len(images.shape) ==3: + images = images.unsqueeze(1) images, masks = images.to(device), masks.to(device) - + if masks.dim()==3: + masks = masks.unsqueeze(1) # Zero the gradients optimizer.zero_grad() - - # Forward pass outputs = model(images) - loss = criterion(outputs, masks) - - # Backward pass and optimization + # Forward pass + if outputs.dim() ==4: + outputs = outputs.view(outputs.size(0),-1,outputs.size(2),outputs.size(3)) + masks_resized = F.interpolate(masks, size=(256, 128), mode="nearest") + loss = criterion(outputs, masks_resized) + loss.backward() optimizer.step() - + running_loss += loss.item() avg_loss = running_loss / len(loader) return avg_loss - -# Training loop +#Training Loop for epoch in range(NUM_EPOCHS): train_loss = train(model, train_loader, optimizer, criterion, DEVICE) val_loss = validate(model, val_loader, criterion, DEVICE) - + print(f"Epoch {epoch+1}/{NUM_EPOCHS}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}") \ No newline at end of file diff --git a/recognition/dataset.py b/recognition/dataset.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/recognition/modules.py b/recognition/modules.py deleted file mode 100644 index 42d44e0a3..000000000 --- a/recognition/modules.py +++ /dev/null @@ -1,46 +0,0 @@ -import torch -import torch.nn as nn - -class encoderBlock(nn.Module): - #convolution+ReLU+maxPool - def __init__(self, in_channels, out_channels): - super(encoderBlock, self).__init__() - self.conv = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size = 3, padding = 1), - nn.ReLU(inplace=True), - nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), - nn.ReLU(inplace=True), - ) - self.pool = nn.MaxPool2d(kernel_size=2, stride=2) - - def forward(self, x): - x = self.conv(x) - pooled = self.pool(x) - return x, pooled - -def decoderBlock(): - #convolution+concatenation+convolution+ReLU - def __init__(self, inchannels, out_channels): - super(decoderBlock, self).__init() - self.upconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride = 2) - self.conv = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size = 3, padding = 1), - nn.ReLU(inplace=True) - nn.Conv2d(out_channels, out_channels, kernel_size = 3, padding = 1), - nn.ReLU(inplace=True), - ) - - def forward(self, x, skip_features): - x = self.upconv(x) - x = torch.cat((x, skip_features), dim = 1) - x = self.conv(x) - return x - -def bottleneck(): - #convolutionalLayer1 - #convolutionalLayer2 - -def skipConnections(): - -def outputLayer(): - diff --git a/recognition/predict.py b/recognition/predict.py deleted file mode 100644 index e69de29bb..000000000 From bc496037767e0674721f55307d14d8e142f31408 Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Fri, 1 Nov 2024 23:40:52 +1000 Subject: [PATCH 05/10] feat: added predict.py --- .../HipMRI_2dUnet_s4858392/niftiload.py | 68 +++++++++++++++ recognition/HipMRI_2dUnet_s4858392/predict.py | 87 +++++++++++++++++++ recognition/HipMRI_2dUnet_s4858392/train.py | 12 ++- 3 files changed, 165 insertions(+), 2 deletions(-) create mode 100644 recognition/HipMRI_2dUnet_s4858392/niftiload.py create mode 100644 recognition/HipMRI_2dUnet_s4858392/predict.py diff --git a/recognition/HipMRI_2dUnet_s4858392/niftiload.py b/recognition/HipMRI_2dUnet_s4858392/niftiload.py new file mode 100644 index 000000000..7df81f8b1 --- /dev/null +++ b/recognition/HipMRI_2dUnet_s4858392/niftiload.py @@ -0,0 +1,68 @@ +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 + +# 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 (normalize 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(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 + inImage = inImage.astype(dtype) + + if normImage: + # Normalize the image + inImage = (inImage - inImage.mean()) / inImage.std() + + if categorical: + inImage = to_channels(inImage, dtype=dtype) + images[i, :, :, :] = inImage + else: + start_x = (inImage.shape[0] - 256) // 2 # Vertical center + start_y = (inImage.shape[1] - 128) // 2 # Horizontal center + cropped_image = inImage[start_x:start_x + 256, start_y:start_y + 128] + images[i, :, :] = cropped_image + 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/HipMRI_2dUnet_s4858392/predict.py b/recognition/HipMRI_2dUnet_s4858392/predict.py new file mode 100644 index 000000000..124cd198d --- /dev/null +++ b/recognition/HipMRI_2dUnet_s4858392/predict.py @@ -0,0 +1,87 @@ +import os +import torch +import nibabel as nib +import numpy as np +from torch.utils.data import DataLoader +from modules import uNet +from newdataset import NIFTIDataset +import torch.nn.functional as F + + +# Hyperparameters +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +MODEL_PATH = "models/unet_model.pth" +OUTPUT_DIR = "/home/Student/s4858392/PAR/results" +GROUND_TRUTH_DIR = "/home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_seg_train" + +def load_model(model_path): + model = uNet(in_channels=1, out_channels=1).to(DEVICE) + model.load_state_dict(torch.load(model_path)) + model.eval() # Set the model to evaluation mode + return model + + +def predict(model, dataset_loader): + predictions = [] + with torch.no_grad(): + for images in dataset_loader: + images = images.to(DEVICE) + if images.dim() == 3: # Check if it's a 3D tensor + images = images.unsqueeze(1) # Add a channel dimension if necessary + + outputs = model(images) + outputs = torch.sigmoid(outputs) # Apply sigmoid to get probabilities + + # Assuming binary segmentation; thresholding to get binary masks + binary_mask = (outputs > 0.5).float() # Threshold at 0.5 + predictions.append(binary_mask.cpu().numpy()) + + return predictions + +def save_predictions(predictions, output_dir): + os.makedirs(output_dir, exist_ok=True) + for i, pred in enumerate(predictions): + img = nib.Nifti1Image(pred[0], np.eye(4)) + nib.save(img, os.path.join(output_dir, f"predicted_mask_{i}.nii.gz")) + +def load_ground_truth(ground_truth_dir): + ground_truth_masks = [] + for file_name in os.listdir(ground_truth_dir): + if file_name.endswith(".nii.gz"): + img = nib.load(os.path.join(ground_truth_dir, file_name)).get_fdata() + ground_truth_masks.append(img) + return ground_truth_masks + +def dice_score(pred, target): + smooth = 1e-6 + pred = pred.flatten() + target = target.flatten() + intersection = np.sum(pred * target) + return (2. * intersection + smooth) / (np.sum(pred) + np.sum(target) + smooth) + + +if __name__ == '__main__': + # Load the trained model + model = load_model(MODEL_PATH) + + # Prepare your new dataset (adjust the path to your new images) + test_image_dir = "//home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_train" + test_dataset = NIFTIDataset(imageDir=test_image_dir) + test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False) + ground_truth_masks = load_ground_truth(GROUND_TRUTH_DIR) + + # Make predictions + predictions = predict(model, test_loader) + + # Save predictions to a directory + save_predictions(predictions, output_dir="/home/Student/s4858392/PAR/results") + print("Predictions saved to 'results' directory.") + + dice_scores = [] + for pred, gt in zip(predictions, ground_truth_masks): + score = dice_score(pred[0], gt) + dice_scores.append(score) + print(f"Dice Score: {score:.4f}") + + avg_dice_score = np.mean(dice_scores) + print(f"Average Dice Score: {avg_dice_score:.4f}") \ No newline at end of file diff --git a/recognition/HipMRI_2dUnet_s4858392/train.py b/recognition/HipMRI_2dUnet_s4858392/train.py index 0fcacc60b..ab3cbef30 100644 --- a/recognition/HipMRI_2dUnet_s4858392/train.py +++ b/recognition/HipMRI_2dUnet_s4858392/train.py @@ -9,17 +9,22 @@ from modules import uNet from newdataset import NIFTIDataset import torch.nn.functional as F +import os #Hyperparameters LEARNING_RATE = 1e-4 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" BATCH_SIZE = 16 -NUM_EPOCHS = 3 +NUM_EPOCHS = 5 NUM_WORKERS = 2 IMAGE_HEIGHT = 128 IMAGE_WIDTH = 256 PIN_MEMORY = True LOAD_MODEL = False +MODEL_DIR = "models" +os.makedirs(MODEL_DIR, exist_ok=True) +MODEL_PATH = os.path.join(MODEL_DIR, "unet_model.pth") + if __name__ == '__main__': #images dataset @@ -92,4 +97,7 @@ def train(model, loader, optimizer, criterion, device): train_loss = train(model, train_loader, optimizer, criterion, DEVICE) val_loss = validate(model, val_loader, criterion, DEVICE) - print(f"Epoch {epoch+1}/{NUM_EPOCHS}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}") \ No newline at end of file + print(f"Epoch {epoch+1}/{NUM_EPOCHS}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}") + + torch.save(model.state_dict(), MODEL_PATH) + print(f"Model saved to {MODEL_PATH}") \ No newline at end of file From 01391801ac58e4cdff506be605c18d0ea6dde904 Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Sat, 2 Nov 2024 00:00:30 +1000 Subject: [PATCH 06/10] feat: include dice score into predict.pu --- recognition/HipMRI_2dUnet_s4858392/predict.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/recognition/HipMRI_2dUnet_s4858392/predict.py b/recognition/HipMRI_2dUnet_s4858392/predict.py index 124cd198d..f49368633 100644 --- a/recognition/HipMRI_2dUnet_s4858392/predict.py +++ b/recognition/HipMRI_2dUnet_s4858392/predict.py @@ -12,7 +12,7 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu" MODEL_PATH = "models/unet_model.pth" OUTPUT_DIR = "/home/Student/s4858392/PAR/results" -GROUND_TRUTH_DIR = "/home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_seg_train" +GROUND_TRUTH_DIR = "/home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_seg_test" def load_model(model_path): model = uNet(in_channels=1, out_channels=1).to(DEVICE) @@ -22,9 +22,14 @@ def load_model(model_path): def predict(model, dataset_loader): + model.eval() predictions = [] with torch.no_grad(): - for images in dataset_loader: + for batch in dataset_loader: + images, _ = batch + if isinstance(images, list): + images = torch.stack(images) + images = images.to(DEVICE) if images.dim() == 3: # Check if it's a 3D tensor images = images.unsqueeze(1) # Add a channel dimension if necessary @@ -65,7 +70,7 @@ def dice_score(pred, target): model = load_model(MODEL_PATH) # Prepare your new dataset (adjust the path to your new images) - test_image_dir = "//home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_train" + test_image_dir = "//home/groups/comp3710/HipMRI_Study_open/keras_slices_data/keras_slices_test" test_dataset = NIFTIDataset(imageDir=test_image_dir) test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False) ground_truth_masks = load_ground_truth(GROUND_TRUTH_DIR) From 010905924eb2aaae83d867793706eee25e9d595a Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Sat, 2 Nov 2024 01:15:23 +1000 Subject: [PATCH 07/10] Update README.md --- README.md | 47 ++++++++++++++++++++++++++++++++--------------- 1 file changed, 32 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 3a10f6515..a0d032353 100644 --- a/README.md +++ b/README.md @@ -1,19 +1,36 @@ -# Pattern Analysis -Pattern Analysis of various datasets by COMP3710 students in 2024 at the University of Queensland. +# U-Net Image Segmentation -We create pattern recognition and image processing library for Tensorflow (TF), PyTorch or JAX. +## Overview +This project implements a U-Net architecture for MRI segmentation tasks, mainly on Hip MRIs. The U-Net model is made to learn the features to provide accurate segmentations. This project has data preproccessing, training, predictions and evaluation metrics included. -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. +## Table of Contents +- [Features](#features) +- [Usage](#usage) +- [Hyperparameters](#hyperparameters) +- [Training and Evaluation](#training-and-evaluation) +- [Results](#results) -The library includes the following implemented in Tensorflow: -* fractals -* recognition problems +## Features +- U-Net architecture +- Design to effectively segment MRI images of Nifti format +- Overfitting measures with early stopping +- Hyperparameters are adjustable +- Evaluation metrics including Dice Score and Loss +- CUDA ready with Cpu backup +- Predictions on new unseen images + +## Installation +Clone the repository and install the required dependencies + +## Hyperparameters +- Learning Rate: 1e-5 +- Epochs = 5 +- Batch size = 4 + +## Training and Evaluation +- Average Training Loss: ~0.001 +- Average Evaluation Loss: ~0.04 + +## Results +- Average Dice Score: TBC -In the recognition folder, you will find many recognition problems solved including: -* segmentation -* classification -* graph neural networks -* StyleGAN -* Stable diffusion -* transformers -etc. From f287bfaa79e297126a620fbf8798c0aadc10eb10 Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Sat, 2 Nov 2024 01:23:42 +1000 Subject: [PATCH 08/10] Update README.md --- recognition/HipMRI_2dUnet_s4858392/README.md | 35 ++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/recognition/HipMRI_2dUnet_s4858392/README.md b/recognition/HipMRI_2dUnet_s4858392/README.md index e69de29bb..838e1cfb1 100644 --- a/recognition/HipMRI_2dUnet_s4858392/README.md +++ b/recognition/HipMRI_2dUnet_s4858392/README.md @@ -0,0 +1,35 @@ +# U-Net Image Segmentation + +## Overview +This project implements a U-Net architecture for MRI segmentation tasks, mainly on Hip MRIs. The U-Net model is made to learn the features to provide accurate segmentations. This project has data preproccessing, training, predictions and evaluation metrics included. + +## Table of Contents +- [Features](#features) +- [Usage](#usage) +- [Hyperparameters](#hyperparameters) +- [Training and Evaluation](#training-and-evaluation) +- [Results](#results) + +## Features +- U-Net architecture +- Design to effectively segment MRI images of Nifti format +- Overfitting measures with early stopping +- Hyperparameters are adjustable +- Evaluation metrics including Dice Score and Loss +- CUDA ready with Cpu backup +- Predictions on new unseen images + +## Installation +Clone the repository and install the required dependencies + +## Hyperparameters +- Learning Rate: 1e-5 +- Epochs = 5 +- Batch size = 4 + +## Training and Evaluation +- Average Training Loss: ~0.001 +- Average Evaluation Loss: ~0.04 + +## Results +- Average Dice Score: TBC From a7af3a9fdb2a8eb4b89c3a1990e4e262ffca8656 Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Mon, 4 Nov 2024 00:53:13 +1000 Subject: [PATCH 09/10] Update README.md --- README.md | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index a0d032353..77b87f36d 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,12 @@ # U-Net Image Segmentation ## Overview -This project implements a U-Net architecture for MRI segmentation tasks, mainly on Hip MRIs. The U-Net model is made to learn the features to provide accurate segmentations. This project has data preproccessing, training, predictions and evaluation metrics included. +This project implements a U-Net architecture for MRI segmentation tasks, mainly on Hip MRIs. The U-Net model consist of two main parts, an encoder and decoder, the encoder progressibely reduces the spatial dimensions of the input image while extracting features, while the decoder reconstructs the image by upsampling and refining features. The model uses skip connections between the encoder and decoder to retain spatial information. A bottleneck layer sits between the encoder and decoder, connecting the two. +This project has data preproccessing, training, predictions and evaluation metrics included. ## Table of Contents - [Features](#features) -- [Usage](#usage) +- [Dependencies](#dependencies) - [Hyperparameters](#hyperparameters) - [Training and Evaluation](#training-and-evaluation) - [Results](#results) @@ -19,8 +20,12 @@ This project implements a U-Net architecture for MRI segmentation tasks, mainly - CUDA ready with Cpu backup - Predictions on new unseen images -## Installation -Clone the repository and install the required dependencies +## Dependencies +- Pytorch 2.3.0 +- Numpy 1.26.4 +- Nibabel +- Albumentations 1.2.0 +- tqdm 4.66.5 ## Hyperparameters - Learning Rate: 1e-5 From 6bbec349e141b749b25be4d1ee73ec42fa02be77 Mon Sep 17 00:00:00 2001 From: Kylow797 <87619255+Kylow797@users.noreply.github.com> Date: Mon, 4 Nov 2024 00:54:19 +1000 Subject: [PATCH 10/10] Update README.md --- recognition/HipMRI_2dUnet_s4858392/README.md | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/recognition/HipMRI_2dUnet_s4858392/README.md b/recognition/HipMRI_2dUnet_s4858392/README.md index 838e1cfb1..917d218b3 100644 --- a/recognition/HipMRI_2dUnet_s4858392/README.md +++ b/recognition/HipMRI_2dUnet_s4858392/README.md @@ -1,11 +1,12 @@ # U-Net Image Segmentation ## Overview -This project implements a U-Net architecture for MRI segmentation tasks, mainly on Hip MRIs. The U-Net model is made to learn the features to provide accurate segmentations. This project has data preproccessing, training, predictions and evaluation metrics included. +This project implements a U-Net architecture for MRI segmentation tasks, mainly on Hip MRIs. The U-Net model consist of two main parts, an encoder and decoder, the encoder progressibely reduces the spatial dimensions of the input image while extracting features, while the decoder reconstructs the image by upsampling and refining features. The model uses skip connections between the encoder and decoder to retain spatial information. A bottleneck layer sits between the encoder and decoder, connecting the two. +This project has data preproccessing, training, predictions and evaluation metrics included. ## Table of Contents - [Features](#features) -- [Usage](#usage) +- [Dependencies](#dependencies) - [Hyperparameters](#hyperparameters) - [Training and Evaluation](#training-and-evaluation) - [Results](#results) @@ -19,8 +20,12 @@ This project implements a U-Net architecture for MRI segmentation tasks, mainly - CUDA ready with Cpu backup - Predictions on new unseen images -## Installation -Clone the repository and install the required dependencies +## Dependencies +- Pytorch 2.3.0 +- Numpy 1.26.4 +- Nibabel +- Albumentations 1.2.0 +- tqdm 4.66.5 ## Hyperparameters - Learning Rate: 1e-5 @@ -32,4 +37,4 @@ Clone the repository and install the required dependencies - Average Evaluation Loss: ~0.04 ## Results -- Average Dice Score: TBC +- Average Dice Score: ~0.41