Commit 8171f623 authored by Andrei-Claudiu Roibu's avatar Andrei-Claudiu Roibu 🖥
Browse files

added evaluation for multiple paths

parent 7520de29
......@@ -262,7 +262,7 @@ def evaluate_single_path(trained_model_path,
for volume_index, file_path in enumerate(file_paths):
try:
# Generate volume & header
predicted_volume, header = _generate_volume(file_path, model, orientation, batch_size, device, cuda_available)
_, predicted_volume, header = _generate_volume(file_path, model, orientation, batch_size, device, cuda_available)
# Generate New Header Affine
......@@ -288,14 +288,136 @@ def evaluate_single_path(trained_model_path,
except Exception as exception_expression:
log.error("Error code execution!")
log.exception(exception_expression):
log.exception(exception_expression)
if exit_on_error:
raise(exception_expression)
log.info("rsfMRI Generation Complete")
def evaluate_multiple_paths():
pass
def evaluate_two_paths(trained_model1_path,
trained_model2_path,
number_of_classes,
data_directory,
targets_directory,
data_list,
orientation1,
orientation2,
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
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)
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)
torch.cuda.empty_cache()
model1.cuda(device)
model2.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))
model1.eval()
model2.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)
# 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, 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
......@@ -338,6 +460,8 @@ def _generate_volume(file_path, model, orientation, batch_size, device, cuda_ava
output = model(batch_x)
output_volume.append(output)
output_volume = torch.cat(output_volume)
_, predicted_volume = torch.max(output_volume, dim=1)
# This needs to be checked - torch.max returns max values and locations.
......@@ -351,10 +475,13 @@ def _generate_volume(file_path, model, orientation, batch_size, device, cuda_ava
if orientation == 'coronal':
predicted_volume = predicted_volume.transpose((1,2,0))
output_volume = output_volume.permute((2,1,3,0))
elif orientation == 'axial':
predicted_volume = predicted_volume.transpose((2,0,1))
elif:
output_volume = output_volume.permute((3,1,0,2))
elif orientation == 'sagittal':
predicted_volume = predicted_volume
output_volume = output_volume
return predicted_volume, header
return output_volume, predicted_volume, header
\ No newline at end of file
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