Commit 010773a3 authored by Andrei-Claudiu Roibu's avatar Andrei-Claudiu Roibu 🖥
Browse files

added function for evaluating on all 3 paths

parent 8171f623
......@@ -10,6 +10,8 @@ To use content from this folder, import the functions and instantiate them as yo
from utils.data_evaluation_utils import function_name
TODO: Might be worth adding some information on uncertaintiy estimation, later down the line
"""
import os
......@@ -387,7 +389,7 @@ def evaluate_two_paths(trained_model1_path,
predicted_volume = (predicted_volume.cpu().numpy()).astype('float32')
predicted_volume = np.squeeze(predicted_volume)
# Generate New Header Affine
header_affines = np.array([header['srow_x'], header['srow_y'], header['srow_z'], [0,0,0,1]])
......@@ -418,7 +420,139 @@ def evaluate_two_paths(trained_model1_path,
log.info("rsfMRI Generation Complete")
def evaluate_all_paths(trained_model1_path,
trained_model2_path,
trained_model3_path,
number_of_classes,
data_directory,
targets_directory,
data_list,
orientation1,
orientation2,
orientation3,
prediction_output_path,
batch_size,
device= 0,
LogWriter= None,
mode='evaluate',
exit_on_error = False):
"""Two Path Evaluator
This function generates the rsfMRI map for an input running on on a single axis or path
Args:
trained_model1_path (str): Path to the location of the trained model1
trained_model2_path (str): Path to the location of the trained model2
trained_model3_path (str): Path to the location of the trained model3
number_of_classes (int): Number of classes
data_directory (str): Path to input data directory
targets_directory (str): Path to labelled data (Y-equivalent)
data_list (str): Path to a .txt file containing the input files for consideration
orientation1 (str): String detailing the current view 1 (COR, SAG, AXL)
orientation2 (str): String detailing the current view 2 (COR, SAG, AXL)
orientation3 (str): String detailing the current view 2 (COR, SAG, AXL)
prediction_output_path (str): Output prediction path
batch_size (int): Size of batch to be evaluated
device (str/int): Device type used for training (int - GPU id, str- CPU)
LogWriter (class): Log Writer class for the BrainMapper U-net
mode (str): Current run mode or phase
exit_on_error (bool): Flag that triggers the raising of an exception
Returns:
None
Raises:
FileNotFoundError: Error in reading the provided file!
Exception: Error code execution!
"""
log.info("Started Evaluation. Check tensorboard for plots (if a LogWriter is provided)")
with open(data_list) as data_list_file:
volumes_to_be_used = data_list_file.read().splitlines()
# Test if cuda is available and attempt to run on GPU
cuda_available = torch.cuda.is_available()
if type(device) == int:
if cuda_available:
model1 = torch.load(trained_model1_path)
model2 = torch.load(trained_model2_path)
model3 = torch.load(trained_model3_path)
torch.cuda.empty_cache()
model1.cuda(device)
model2.cuda(device)
model3.cuda(device)
else:
log.warning("CUDA not available. Switching to CPU. Investigate behaviour!")
device = 'cpu'
if (type(device) == str) or not cuda_available:
model1 = torch.load(trained_model1_path, map_location=torch.device(device))
model2 = torch.load(trained_model2_path, map_location=torch.device(device))
model3 = torch.load(trained_model3_path, map_location=torch.device(device))
model1.eval()
model2.eval()
model3.eval()
# Create the prediction path folder if this is not available
data_utils.create_folder(prediction_output_path)
# Initiate the evaluation
log.info("rsfMRI Generation Started")
file_paths = data_utils.load_file_paths(data_directory, data_list)
with torch.no_grad():
for volume_index, file_path in enumerate(file_paths):
try:
# Generate volume & header
output_volume1, _, header = _generate_volume(file_path, model1, orientation1, batch_size, device, cuda_available)
output_volume2, _, header = _generate_volume(file_path, model2, orientation2, batch_size, device, cuda_available)
output_volume3, _, header = _generate_volume(file_path, model3, orientation3, batch_size, device, cuda_available)
# Agregate the two views
# TODO: Consider other ways of combining
# In the original implementation, authors went for summation, however, there are other options
# Other options: take the mean? Some other summtion? Take the max?
_, predicted_volume = torch.max(output_volume1 + output_volume2 + output_volume3, dim=1)
predicted_volume = (predicted_volume.cpu().numpy()).astype('float32')
predicted_volume = np.squeeze(predicted_volume)
# Generate New Header Affine
header_affines = np.array([header['srow_x'], header['srow_y'], header['srow_z'], [0,0,0,1]])
# We apply the affine and save the image
output_nifti_image = nib.Nifti1Image(predicted_volume, header_affines, header= header)
output_nifti_path = os.path.join(prediction_output_path, volumes_to_be_used[volume_index] + str('.mgz'))
if '.nii' not in output_nifti_path:
output_nifti_path += '.nii.gz'
nib.save(output_nifti_image, output_nifti_path)
log.info("Processed: "+ volumes_to_be_used[volume_index] + " " + str(volume_index + 1) + " out of " + str(len(volumes_to_be_used)))
except FileNotFoundError as exception_expression:
log.error("Error in reading the provided file!")
log.exception(exception_expression)
if exit_on_error:
raise(exception_expression)
except Exception as exception_expression:
log.error("Error code execution!")
log.exception(exception_expression)
if exit_on_error:
raise(exception_expression)
log.info("rsfMRI Generation Complete")
def _generate_volume(file_path, model, orientation, batch_size, device, cuda_available):
"""rsfMRI Volume Generator
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment