added regression for evaluation

run.py

@@ -300,12 +300,14 @@ def evaluate_mapping(mapping_evaluation_parameters):
     mapping_data_file = mapping_evaluation_parameters['mapping_data_file']
     data_list = mapping_evaluation_parameters['data_list']
     prediction_output_path = mapping_evaluation_parameters['prediction_output_path']
+    dmri_mean_mask_path = mapping_evaluation_parameters['dmri_mean_mask_path']
+    rsfmri_mean_mask_path = mapping_evaluation_parameters['rsfmri_mean_mask_path']
     device = mapping_evaluation_parameters['device']
     exit_on_error = mapping_evaluation_parameters['exit_on_error']
     brain_mask_path = mapping_evaluation_parameters['brain_mask_path']
-    mean_mask_path = mapping_evaluation_parameters['mean_mask_path']
     mean_reduction = mapping_evaluation_parameters['mean_reduction']
     scaling_factors = mapping_evaluation_parameters['scaling_factors']
+    regression_factors = mapping_evaluation_parameters['regression_factors']
 
     evaluations.evaluate_mapping(trained_model_path,
                      data_directory,
@@ -313,9 +315,11 @@ def evaluate_mapping(mapping_evaluation_parameters):
                      data_list,
                      prediction_output_path,
                      brain_mask_path,
-                                 mean_mask_path,
+                     dmri_mean_mask_path,
+                     rsfmri_mean_mask_path,
                      mean_reduction,
                      scaling_factors,
+                     regression_factors,
                      device=device,
                      exit_on_error=exit_on_error)
 

settings.ini

@@ -28,13 +28,13 @@ final_model_output_file = "CU3D17-3.pth.tar"
 training_batch_size = 5
 validation_batch_size = 5
 use_pre_trained = False
-learning_rate = 1e-1
+learning_rate = 1e-5
 optimizer_beta = (0.9, 0.999)
 optimizer_epsilon = 1e-8
 optimizer_weigth_decay = 1e-5
 number_of_epochs = 200
 loss_log_period = 50
-learning_rate_scheduler_step_size = 5
+learning_rate_scheduler_step_size = 4
 learning_rate_scheduler_gamma = 1e-1
 use_last_checkpoint = False
 

settings_evaluation.ini

@@ -5,8 +5,11 @@ data_directory = "/well/win-biobank/projects/imaging/data/data3/subjectsAll/"
 mapping_data_file = "dMRI/autoptx_preproc/tractsNormSummed.nii.gz"
 data_list = "datasets/test.txt"
 brain_mask_path = "utils/MNI152_T1_2mm_brain_mask.nii.gz"
+rsfmri_mean_mask_path = "utils/mean_dr_stage2.nii.gz"
+dmri_mean_mask_path = "utils/mean_tractsNormSummed_downsampled.nii.gz"
 mean_mask_path = "utils/mean_dr_stage2.nii.gz"
 scaling_factors = "datasets/scaling_factors.pkl"
+regression_factors = "datasets/regression_weights.pkl"
 mean_reduction = True
 device = 0
 exit_on_error = True

utils/data_evaluation_utils.py

@@ -22,6 +22,7 @@ import logging
 import utils.data_utils as data_utils
 import nibabel as nib
 import matplotlib.pyplot as plt
+import pandas as pd
 from fsl.data.image import Image
 
 log = logging.getLogger(__name__)
@@ -151,9 +152,11 @@ def evaluate_mapping(trained_model_path,
                      data_list,
                      prediction_output_path,
                      brain_mask_path,
-                     mean_mask_path,
+                     dmri_mean_mask_path,
+                     rsfmri_mean_mask_path,
                      mean_reduction,
                      scaling_factors,
+                     regression_factors,
                      device=0,
                      mode='evaluate',
                      exit_on_error=False):
@@ -168,9 +171,11 @@ def evaluate_mapping(trained_model_path,
         data_list (str): Path to a .txt file containing the input files for consideration
         prediction_output_path (str): Output prediction path
         brain_mask_path (str): Path to the MNI brain mask file
-        mean_mask_path (str): Path to the dualreg subject mean mask
+        dmri_mean_mask_path (str): Path to the dualreg subject mean mask
+        rsfmri_mean_mask_path (str): Path to the summed tract mean mask
         mean_reduction (bool): Flag indicating if the targets should be de-meaned using the mean_mask_path
         scaling_factors (str): Path to the scaling factor file
+        regression_factors (str): Path to the linear regression weights file
         device (str/int): Device type used for training (int - GPU id, str- CPU)
         mode (str): Current run mode or phase
         exit_on_error (bool): Flag that triggers the raising of an exception
@@ -212,7 +217,7 @@ def evaluate_mapping(trained_model_path,
     # Initiate the evaluation
 
     log.info("rsfMRI Generation Started")
-    file_paths = data_utils.load_file_paths(
+    file_paths, volumes_to_be_used = data_utils.load_file_paths(
         data_directory, data_list, mapping_data_file)
 
     with torch.no_grad():
@@ -221,8 +226,14 @@ def evaluate_mapping(trained_model_path,
             try:
                 print("Mapping Volume {}/{}".format(volume_index+1, len(file_paths)))
                 # Generate volume & header
+
+                subject = volumes_to_be_used[volume_index]
+
                 predicted_complete_volume, predicted_volume, header, xform = _generate_volume_map(
-                    file_path, model, device, cuda_available, brain_mask_path, mean_mask_path, scaling_factors, mean_reduction)
+                    file_path, subject, model, 
+                    device, cuda_available, brain_mask_path, 
+                    dmri_mean_mask_path, rsfmri_mean_mask_path, scaling_factors, 
+                    regression_factors, mean_reduction)
 
                 # Generate New Header Affine
 
@@ -273,19 +284,32 @@ def evaluate_mapping(trained_model_path,
     log.info("rsfMRI Generation Complete")
 
 
-def _generate_volume_map(file_path, model, device, cuda_available, brain_mask_path, mean_mask_path, scaling_factors, mean_reduction=False):
+def _generate_volume_map(file_path, 
+                        subject, 
+                        model, 
+                        device, 
+                        cuda_available, 
+                        brain_mask_path, 
+                        dmri_mean_mask_path, 
+                        rsfmri_mean_mask_path, 
+                        scaling_factors, 
+                        regression_factors,
+                        mean_reduction=False):
     """rsfMRI Volume Generator
 
     This function uses the trained model to generate a new volume
 
     Args:
         file_path (str): Path to the desired file
+        subject (str): Subject ID of the subject volume to be regressed
         model (class): BrainMapper model class
         device (str/int): Device type used for training (int - GPU id, str- CPU)
         cuda_available (bool): Flag indicating if a cuda-enabled GPU is present
         brain_mask_path (str): Path to the MNI brain mask file
-        mean_mask_path (str): Path to the dualreg subject mean mask
+        dmri_mean_mask_path (str): Path to the group mean volume
+        rsfmri_mean_mask_path (str): Path to the dualreg subject mean mask
         scaling_factors (str): Path to the scaling factor file
+        regression_factors (str): Path to the linear regression weights file
         mean_reduction (bool): Flag indicating if the targets should be de-meaned using the mean_mask_path
 
     Returns
@@ -296,6 +320,8 @@ def _generate_volume_map(file_path, model, device, cuda_available, brain_mask_pa
     volume, header, xform = data_utils.load_and_preprocess_evaluation(
         file_path)
 
+    volume = _regress_input(volume, subject, dmri_mean_mask_path, regression_factors)
+
     print('volume range:', np.min(volume), np.max(volume))
 
     volume = _scale_input(volume, scaling_factors)
@@ -323,8 +349,9 @@ def _generate_volume_map(file_path, model, device, cuda_available, brain_mask_pa
     MNI152_T1_2mm_brain_mask = Image(brain_mask_path).data
 
     if mean_reduction == True:
-        mean_mask = Image(mean_mask_path).data[:, :, :, 0]
-        predicted_complete_volume = np.add(output, mean_mask)
+        mean_mask = Image(rsfmri_mean_mask_path).data[:, :, :, 0]
+        weight = pd.read_pickle(regression_factors).loc[subject]['w_rsfMRI']
+        predicted_complete_volume = np.add(output, np.multiply(weight, mean_mask))
 
         print('predicted_complete_volume', np.min(predicted_complete_volume), np.max(predicted_complete_volume))
 
@@ -373,6 +400,30 @@ def _scale_input(volume, scaling_factors):
     return scaled_volume
 
 
+def _regress_input(volume, subject, dmri_mean_mask_path, regression_factors):
+    """ Inputn Regression
+
+    This function regresse the group mean from the input volume using the saved regression weights.
+
+    TODO: This function repressents only a temporary solution. For deployment, a NN needs to be trained which predicts the relevant scaling factors.
+
+    Args:
+        volume (np.array): Unregressed volume
+        subject (str): Subject ID of the subject volume to be regressed
+        dmri_mean_mask_path (str): Path to the group mean volume
+        regression_factors (str): Path to the linear regression weights file
+
+    Returns:
+        regressed_volume (np.array): Linear regressed volume
+
+    """
+
+    weight = pd.read_pickle(regression_factors).loc[subject]['w_dMRI']
+    group_mean = Image(dmri_mean_mask_path).data
+    regressed_volume = np.subtract(volume, np.multiply(weight, group_mean))
+
+    return regressed_volume
+
 def _rescale_output(volume, scaling_factors):
     """Output Rescaling
 

utils/data_utils.py

@@ -608,7 +608,7 @@ def load_file_paths(data_directory, data_list, mapping_data_file, targets_direct
         file_paths = [[os.path.join(data_directory, volume, mapping_data_file), os.path.join(
             targets_directory, volume)] for volume in volumes_to_be_used]
 
-    return file_paths
+    return file_paths, volumes_to_be_used 
 
 
 def load_and_preprocess_evaluation(file_path):