diff --git a/README.md b/README.md index 3a10f6515..77b87f36d 100644 --- a/README.md +++ b/README.md @@ -1,19 +1,41 @@ -# 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 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. -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) +- [Dependencies](#dependencies) +- [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 + +## Dependencies +- Pytorch 2.3.0 +- Numpy 1.26.4 +- Nibabel +- Albumentations 1.2.0 +- tqdm 4.66.5 + +## 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. diff --git a/recognition/HipMRI_2dUnet_s4858392/README.md b/recognition/HipMRI_2dUnet_s4858392/README.md new file mode 100644 index 000000000..917d218b3 --- /dev/null +++ b/recognition/HipMRI_2dUnet_s4858392/README.md @@ -0,0 +1,40 @@ +# U-Net Image Segmentation + +## Overview +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) +- [Dependencies](#dependencies) +- [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 + +## Dependencies +- Pytorch 2.3.0 +- Numpy 1.26.4 +- Nibabel +- Albumentations 1.2.0 +- tqdm 4.66.5 + +## 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: ~0.41 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/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..f49368633 --- /dev/null +++ b/recognition/HipMRI_2dUnet_s4858392/predict.py @@ -0,0 +1,92 @@ +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_test" + +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): + model.eval() + predictions = [] + with torch.no_grad(): + 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 + + 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_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) + + # 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 new file mode 100644 index 000000000..ab3cbef30 --- /dev/null +++ b/recognition/HipMRI_2dUnet_s4858392/train.py @@ -0,0 +1,103 @@ +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 torch.utils.data import DataLoader +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 = 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 + 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 = "/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=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 validate(model, loader, criterion, device): + model.eval() + val_loss = 0.0 + total_dice_score = 0.0 + with torch.no_grad(): + 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) + 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) + return avg_loss + + 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() + outputs = model(images) + # 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 + 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}") + + torch.save(model.state_dict(), MODEL_PATH) + print(f"Model saved to {MODEL_PATH}") \ No newline at end of file diff --git a/recognition/train.py b/recognition/train.py new file mode 100644 index 000000000..e69de29bb