Upload 4 files

utils/BiasFieldCorrection.py

@@ -0,0 +1,91 @@
+import argparse
+import os
+import glob
+import ants
+
+def abp_n4(image, intensity_truncation=(0.025, 0.975, 256), mask=None, 
+           use_n3=False):
+    """
+    Apply intensity truncation and bias correction to an image.
+
+    Parameters:
+    - image: Image to be processed (ANTsImage object).
+    - intensity_truncation: Tuple (lower quantile, upper quantile, number of bins).
+    - mask: Optional mask for bias correction (ANTsImage object).
+    - use_n3: Use N3 bias correction instead of N4 if True.
+
+    Returns:
+    - Processed image (ANTsImage object).
+    """
+    if not isinstance(intensity_truncation, (list, tuple)) or \
+       len(intensity_truncation) != 3:
+        raise ValueError("intensity_truncation must be list/tuple with 3 values")
+
+    # Apply intensity truncation
+    truncated_image = ants.iMath(image, "TruncateIntensity", 
+                                 intensity_truncation[0], intensity_truncation[1],
+                                 intensity_truncation[2])
+
+    # Apply bias correction
+    if use_n3:
+        corrected_image = ants.n3_bias_field_correction(truncated_image, mask=mask)
+    else:
+        corrected_image = ants.n4_bias_field_correction(truncated_image, mask=mask)
+
+    return corrected_image
+
+
+def preprocess_image(file_path, output_path, preprocess_type):
+    """
+    Read, preprocess, and write an image based on the specified method.
+
+    Parameters:
+    - file_path: Path to the input image file.
+    - output_path: Path to the output folder.
+    - preprocess_type: Preprocessing method (n3, n4, or abp).
+    """
+    image = ants.image_read(file_path)
+
+    if preprocess_type == "n3":
+        processed_image = ants.n3_bias_field_correction(image)
+    elif preprocess_type == "n4":
+        processed_image = ants.n4_bias_field_correction(image)
+    elif preprocess_type == "abp":
+        processed_image = abp_n4(image)
+    else:
+        raise ValueError("Invalid preprocess type")
+
+    ants.image_write(processed_image, os.path.join(output_path, 
+                                                   os.path.basename(file_path)))
+
+
+def main():
+    """
+    Main function to handle command line arguments and process images.
+    """
+    parser = argparse.ArgumentParser(
+        description="Image Preprocessing Script for Bias Field Correction")
+    parser.add_argument('--input', type=str, required=True, 
+                        help="Input file or folder path")
+    parser.add_argument('--output', type=str, default='output', 
+                        help="Output folder path")
+    parser.add_argument('--type', type=str, default='abp', 
+                        choices=['n3', 'n4', 'abp'], 
+                        help="Type of preprocessing (n3, n4, abp)")
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.output):
+        os.makedirs(args.output)
+
+    files = [args.input] if os.path.isfile(args.input) else \
+            glob.glob(os.path.join(args.input, '*'))
+
+    for file in files:
+        preprocess_image(file, args.output, args.type)
+
+    print("----- Bias Field Correction Process completed -----")
+
+
+if __name__ == '__main__':
+    main()

utils/data_splitter.py

@@ -0,0 +1,79 @@
+import os
+import shutil
+import random
+import argparse
+
+class DataSplitter:
+    """Class to split Nii or nii.gz files into training and validation datasets."""
+
+    def __init__(self, input_folder, train_ratio, train_folder, val_folder):
+        """
+        Initialize the DataSplitter with folder paths and data ratio.
+
+        :param input_folder: Path to the folder containing the Nii files.
+        :param train_ratio: Ratio of the dataset to be used as training data.
+        :param train_folder: Name of the folder to store training data.
+        :param val_folder: Name of the folder to store validation data.
+        """
+        self.input_folder = input_folder
+        self.train_ratio = train_ratio
+        self.train_folder = os.path.join(input_folder, train_folder)
+        self.val_folder = os.path.join(input_folder, val_folder)
+
+    def split_data(self):
+        """
+        Split the data into training and validation datasets and move files
+        to the respective folders.
+        """
+        # List all Nii or nii.gz files in the input folder
+        files = [f for f in os.listdir(self.input_folder)
+                 if f.endswith('.nii') or f.endswith('.nii.gz')]
+        random.shuffle(files)  # Shuffle files for random splitting
+
+        # Determine the split index based on the training ratio
+        split_index = int(len(files) * self.train_ratio)
+        train_files = files[:split_index]
+        val_files = files[split_index:]
+
+        # Create training and validation folders if they don't exist
+        os.makedirs(self.train_folder, exist_ok=True)
+        os.makedirs(self.val_folder, exist_ok=True)
+
+        # Move files to the respective folders
+        for file in train_files:
+            shutil.move(os.path.join(self.input_folder, file),
+                        self.train_folder)
+
+        for file in val_files:
+            shutil.move(os.path.join(self.input_folder, file),
+                        self.val_folder)
+
+        print(f"Files split into {len(train_files)} training "
+              f"and {len(val_files)} validation.")
+
+def main():
+    """
+    Main function to handle command line arguments and initiate data splitting.
+    """
+    parser = argparse.ArgumentParser(
+        description='Split Nii files into training and validation datasets.')
+
+    # Define command line arguments
+    parser.add_argument('--input', type=str, required=True,
+                        help='Input folder with Nii files.')
+    parser.add_argument('--train-ratio', type=float, default=0.8,
+                        help='Training data ratio (default: 0.8).')
+    parser.add_argument('--train-folder', type=str, default='training',
+                        help='Folder for training data (default: "training").')
+    parser.add_argument('--val-folder', type=str, default='validation',
+                        help='Folder for validation data (default: "validation").')
+
+    args = parser.parse_args()
+
+    # Create a DataSplitter instance and split the data
+    splitter = DataSplitter(args.input, args.train_ratio,
+                            args.train_folder, args.val_folder)
+    splitter.split_data()
+
+if __name__ == "__main__":
+    main()

utils/validation.py

@@ -0,0 +1,85 @@
+import csv
+import torch
+import numpy as np
+from skimage.metrics import structural_similarity as compare_ssim
+import math
+
+class ImageQualityMetrics:
+    """Class to compute image quality metrics like SSIM, PSNR, and MSE."""
+
+    @staticmethod
+    @staticmethod
+    def ssim_3d(img1, img2):
+        """Calculate SSIM for each 2D slice in the 3D image and return the average."""
+        ssim_vals = []
+        for i in range(img1.shape[1]):  # Depth
+            slice1 = img1[0, i, :, :]
+            slice2 = img2[0, i, :, :]
+            ssim_val = compare_ssim(slice1, slice2, data_range=slice1.max() - slice1.min())
+            ssim_vals.append(ssim_val)
+        return np.mean(ssim_vals)
+    
+    @staticmethod
+    def psnr(img1, img2):
+        """Calculate PSNR (Peak Signal-to-Noise Ratio) between two images."""
+        mse = torch.mean((img1 - img2) ** 2)
+        if mse == 0:
+            return math.inf
+        return 20 * math.log10(img1.max() - img1.min()) - 10 * math.log10(mse)
+
+    @staticmethod
+    def mse(img1, img2):
+        """Calculate MSE (Mean Squared Error) between two images."""
+        return torch.mean((img1 - img2) ** 2)
+
+
+class ValidationRecorder:
+    """Class to handle validation process and record the metrics."""
+
+    def __init__(self, csv_file_path):
+        """Initialize the recorder with the path to the CSV file."""
+        self.csv_file_path = csv_file_path
+
+    def initialize_csv(self):
+        """Initialize the CSV file with headers."""
+        with open(self.csv_file_path, mode='w', newline='') as file:
+            writer = csv.writer(file)
+            writer.writerow(['Epoch', 'Loss', 'SSIM', 'PSNR', 'MSE'])
+
+    def validate_and_record(self, epoch, dataloader, device, generator, 
+                            criterion_g):
+        """Validate the model and record the metrics in the CSV file."""
+        generator.eval()
+        total_loss, total_ssim, total_psnr, total_mse = 0.0, 0.0, 0.0, 0.0
+
+        with torch.no_grad():
+            for _, (low_res, high_res) in enumerate(dataloader):
+                low_res, high_res = low_res.to(device), high_res.to(device)
+                fake_images = generator(low_res)
+
+                loss = criterion_g(fake_images, high_res)
+                total_loss += loss.item()
+
+                for j in range(high_res.size(0)):
+                    ssim_val = ImageQualityMetrics.ssim_3d(
+                        high_res[j].cpu().numpy(), fake_images[j].cpu().numpy())
+                    psnr_val = ImageQualityMetrics.psnr(
+                        high_res[j], fake_images[j])
+                    mse_val = ImageQualityMetrics.mse(
+                        high_res[j], fake_images[j])
+                    total_ssim += ssim_val
+                    total_psnr += psnr_val
+                    total_mse += mse_val.item()
+
+        avg_loss = total_loss / len(dataloader)
+        avg_ssim = total_ssim / (len(dataloader) * dataloader.batch_size)
+        avg_psnr = total_psnr / (len(dataloader) * dataloader.batch_size)
+        avg_mse = total_mse / (len(dataloader) * dataloader.batch_size)
+
+        self._write_to_csv(epoch, avg_loss, avg_ssim, avg_psnr, avg_mse)
+
+    def _write_to_csv(self, epoch, avg_loss, avg_ssim, avg_psnr, avg_mse):
+        """Write the validation metrics to the CSV file."""
+        with open(self.csv_file_path, mode='a', newline='') as file:
+            writer = csv.writer(file)
+            writer.writerow([epoch, avg_loss, avg_ssim, avg_psnr, avg_mse])