Commit 9345f994 authored by Andrei-Claudiu Roibu's avatar Andrei-Claudiu Roibu 🖥
Browse files

added function that generates a new volume

parent fab045c7
......@@ -120,11 +120,11 @@ def evaluate_dice_score(trained_model_path,
volume_dice_score_list = []
log.info("Dice Score Evaluation Started")
file_paths = data_utils.load_file_paths(data_directory, targets_directory, data_list)
file_paths = data_utils.load_file_paths(data_directory, data_list, targets_directory)
with torch.no_grad():
for volume_index, file_path in enumerate(file_paths):
volume, label_map, header = data_utils.load_and_preprocess(file_paths, orientation)
volume, label_map, header = data_utils.load_and_preprocess(file_path, orientation)
if len(volume.shape) == 4:
volume = volume
......@@ -143,7 +143,7 @@ def evaluate_dice_score(trained_model_path,
output = model(input_slice)
_, slice_output = torch.max(output, dim=1)\
_, slice_output = torch.max(output, dim=1)
# This needs to be checked - torch.max returns max values and locations.
# For segmentations, we are interested in the locations
......@@ -188,8 +188,151 @@ def evaluate_dice_score(trained_model_path,
return average_dice_score
def evaluate_single_path():
pass
def evaluate_single_path(trained_model_path,
number_of_classes,
data_directory,
targets_directory,
data_list,
orientation,
prediction_output_path,
batch_size,
device= 0,
LogWriter= None,
mode='evaluate',
exit_on_error = False):
"""Single Path Evaluator
This function generates the rsfMRI map for an input running on on a single axis or path
Args:
trained_model_path (str): Path to the location of the trained model
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
orientation (str): String detailing the current view (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
"""
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:
model = torch.load(trained_model_path)
torch.cuda.empty_cache()
model.cuda(device)
else:
log.warning("CUDA not available. Switching to CPU. Investigate behaviour!")
device = 'cpu'
if (type(device) == str) or not cuda_available:
model = torch.load(trained_model_path, map_location=torch.device(device))
model.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
predicted_volume, header = _generate_volume(file_path, model, orientation, batch_size, device, cuda_available)
# get the header
# copy affine
# apply original image affine to new image
# save the image
except:
pass
log.info("rsfMRI Generation Complete")
def evaluate_multiple_paths():
pass
\ No newline at end of file
pass
def _generate_volume(file_path, model, orientation, batch_size, device, cuda_available):
"""rsfMRI Volume Generator
This function uses the trained model to generate a new volume
Args:
file_path (str): Path to the desired file
model (class): BrainMapper model class
orientation (str): String detailing the current view (COR, SAG, AXL)
batch_size (int): Size of batch to be evaluated
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
Returns
predicted_volume (np.array): Array containing the information regarding the generated volume
header (class): 'nibabel.nifti1.Nifti1Header' class object, containing volume metadata
Raises:
None
"""
input_volume, header = data_utils.load_and_preprocess_evaluation(file_path, orientation)
if len(volume.shape) == 4:
volume = volume
else:
volume = volume[:, np.newaxis, :, :]
volume = torch.tensor(volume).type(torch.FloatTensor)
output_volume = []
for i in range(0, len(volume), batch_size):
batch_x = volume[i: i+batch_size]
if cuda_available and (type(device)==int):
batch_x = batch_x.cuda(device)
output = model(batch_x)
output_volume.append(output)
_, predicted_volume = torch.max(output_volume, dim=1)
# This needs to be checked - torch.max returns max values and locations.
# For segmentations, we are interested in the locations
# For the functional data, we might be interested in the actual values.
# The strength of the value represents the strength of the activation
# A threshold might also be required!
predicted_volume = (predicted_volume.cpu().numpy()).astype('float32')
predicted_volume = np.squeeze(predicted_volume)
if orientation == 'coronal':
predicted_volume = predicted_volume.transpose((1,2,0))
elif orientation == 'axial':
predicted_volume = predicted_volume.transpose((2,0,1))
elif:
predicted_volume = predicted_volume
return 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