import os
from quicknat import QuickNat
import nibabel as nib
import numpy as np
import logging
import torch
import csv
import utils.common_utils as common_utils
import utils.data_utils as du
log = logging.getLogger(__name__)
def dice_confusion_matrix(vol_output, ground_truth, num_classes, no_samples=10, mode='train'):
dice_cm = torch.zeros(num_classes, num_classes)
if mode == 'train':
samples = np.random.choice(len(vol_output), no_samples)
vol_output, ground_truth = vol_output[samples], ground_truth[samples]
for i in range(num_classes):
GT = (ground_truth == i).float()
for j in range(num_classes):
Pred = (vol_output == j).float()
inter = torch.sum(torch.mul(GT, Pred))
union = torch.sum(GT) + torch.sum(Pred) + 0.0001
dice_cm[i, j] = 2 * torch.div(inter, union)
avg_dice = torch.mean(torch.diagflat(dice_cm))
return avg_dice, dice_cm
def dice_score_perclass(vol_output, ground_truth, num_classes, no_samples=10, mode='train'):
dice_perclass = torch.zeros(num_classes)
if mode == 'train':
samples = np.random.choice(len(vol_output), no_samples)
vol_output, ground_truth = vol_output[samples], ground_truth[samples]
for i in range(num_classes):
GT = (ground_truth == i).float()
Pred = (vol_output == i).float()
inter = torch.sum(torch.mul(GT, Pred))
union = torch.sum(GT) + torch.sum(Pred) + 0.0001
dice_perclass[i] = (2 * torch.div(inter, union))
return dice_perclass
def compute_volume(prediction_map, label_map, ID):
num_cls = len(label_map) - 1
volume_dict = {}
volume_dict['vol_ID'] = ID
for i in range(num_cls):
binarized_pred = (prediction_map == i).astype(float)
volume_dict[label_map[i + 1]] = np.sum(binarized_pred)
return volume_dict
def _write_csv_table(name, prediction_path, dict_list, label_names):
file_name = name
file_path = os.path.join(prediction_path, file_name)
# Save volume_dict as csv file in the prediction_path
with open(file_path, 'w') as f:
writer = csv.DictWriter(f, fieldnames=label_names)
writer.writeheader()
for data in dict_list:
writer.writerow(data)
def compute_structure_uncertainty(mc_pred_list, label_map, ID):
num_cls = len(label_map) - 1
cvs_dict = {}
cvs_dict['vol_ID'] = ID
iou_dict = {}
iou_dict['vol_ID'] = ID
for c in range(num_cls):
mc_vol = []
inter = (mc_pred_list[0] == c).astype('int')
union = (mc_pred_list[0] == c).astype('int')
mc_vol.append(inter.sum())
for s in range(1, len(mc_pred_list)):
nxt = (mc_pred_list[s] == c).astype('int')
mc_vol.append(nxt.sum())
inter = np.multiply(inter, nxt)
union = (np.add(union, nxt) > 0).astype('int')
s_inter, s_union = np.sum(inter), np.sum(union)
if s_inter == 0 and s_union == 0:
iou_dict[label_map[c + 1]] = 1
elif s_inter > 0 and s_union == 0 or s_inter == 0 and s_union > 0:
iou_dict[label_map[c + 1]] = 0
else:
iou_dict[label_map[c + 1]] = np.divide(s_inter, s_union)
mc_vol = np.array(mc_vol)
cvs_dict[label_map[c + 1]] = np.std(mc_vol) / np.mean(mc_vol)
return iou_dict, cvs_dict
def evaluate_dice_score(model_path, num_classes, data_dir, label_dir, volumes_txt_file, remap_config, orientation,
prediction_path, data_id, device=0, logWriter=None, mode='eval'):
log.info("**Starting evaluation. Please check tensorboard for plots if a logWriter is provided in arguments**")
batch_size = 20
with open(volumes_txt_file) as file_handle:
volumes_to_use = file_handle.read().splitlines()
cuda_available = torch.cuda.is_available()
# First, are we attempting to run on a GPU?
if type(device) == int:
# if CUDA available, follow through, else warn and fallback to CPU
if cuda_available:
model = torch.load(model_path)
torch.cuda.empty_cache()
model.cuda(device)
else:
log.warning(
'CUDA is not available, trying with CPU.' + \
'This can take much longer (> 1 hour). Cancel and ' + \
'investigate if this behavior is not desired.'
)
# switch device to 'cpu'
device = 'cpu'
# If device is 'cpu' or CUDA not available
if (type(device)==str) or not cuda_available:
model = torch.load(
model_path,
map_location=torch.device(device)
)
model.eval()
common_utils.create_if_not(prediction_path)
volume_dice_score_list = []
log.info("Evaluating now...")
file_paths = du.load_file_paths(data_dir, label_dir, data_id, volumes_txt_file)
with torch.no_grad():
for vol_idx, file_path in enumerate(file_paths):
volume, labelmap, class_weights, weights, header = du.load_and_preprocess(file_path,
orientation=orientation,
remap_config=remap_config)
volume = volume if len(volume.shape) == 4 else volume[:, np.newaxis, :, :]
volume, labelmap = torch.tensor(volume).type(torch.FloatTensor), torch.tensor(labelmap).type(
torch.LongTensor)
volume_prediction = []
for i in range(0, len(volume), batch_size):
batch_x, batch_y = volume[i: i + batch_size], labelmap[i:i + batch_size]
if cuda_available and (type(device)==int):
batch_x = batch_x.cuda(device)
out = model(batch_x)
_, batch_output = torch.max(out, dim=1)
volume_prediction.append(batch_output)
volume_prediction = torch.cat(volume_prediction)
volume_dice_score = dice_score_perclass(volume_prediction, labelmap.cuda(device), num_classes, mode=mode)
volume_prediction = (volume_prediction.cpu().numpy()).astype('float32')
#Copy header affine
Mat = np.array([
header['srow_x'],
header['srow_y'],
header['srow_z'],
[0,0,0,1]
])
# Apply original image affine to prediction volume
nifti_img = nib.MGHImage(np.squeeze(volume_prediction), Mat, header=header)
nib.save(nifti_img, os.path.join(prediction_path, volumes_to_use[vol_idx] + str('.mgz')))
if logWriter:
logWriter.plot_dice_score('val', 'eval_dice_score', volume_dice_score, volumes_to_use[vol_idx], vol_idx)
volume_dice_score = volume_dice_score.cpu().numpy()
volume_dice_score_list.append(volume_dice_score)
log.info(volume_dice_score, np.mean(volume_dice_score))
dice_score_arr = np.asarray(volume_dice_score_list)
avg_dice_score = np.mean(dice_score_arr)
log.info("Mean of dice score : " + str(avg_dice_score))
class_dist = [dice_score_arr[:, c] for c in range(num_classes)]
if logWriter:
logWriter.plot_eval_box_plot('eval_dice_score_box_plot', class_dist, 'Box plot Dice Score')
log.info("DONE")
return avg_dice_score, class_dist
def _segment_vol(file_path, model, orientation, batch_size, cuda_available, device):
volume, header = du.load_and_preprocess_eval(file_path,
orientation=orientation)
volume = volume if len(volume.shape) == 4 else volume[:, np.newaxis, :, :]
volume = torch.tensor(volume).type(torch.FloatTensor)
volume_pred = []
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)
out = model(batch_x)
# _, batch_output = torch.max(out, dim=1)
volume_pred.append(out)
volume_pred = torch.cat(volume_pred)
_, volume_prediction = torch.max(volume_pred, dim=1)
volume_prediction = (volume_prediction.cpu().numpy()).astype('float32')
volume_prediction = np.squeeze(volume_prediction)
if orientation == "COR":
volume_prediction = volume_prediction.transpose((1, 2, 0))
volume_pred = volume_pred.permute((2, 1, 3, 0))
elif orientation == "AXI":
volume_prediction = volume_prediction.transpose((2, 0, 1))
volume_pred = volume_pred.permute((3, 1, 0, 2))
return volume_pred, volume_prediction, header
def _segment_vol_unc(file_path, model, orientation, batch_size, mc_samples, cuda_available, device):
volume, header = du.load_and_preprocess_eval(file_path,
orientation=orientation)
volume = volume if len(volume.shape) == 4 else volume[:, np.newaxis, :, :]
volume = torch.tensor(volume).type(torch.FloatTensor)
mc_pred_list = []
for j in range(mc_samples):
volume_pred = []
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)
out = model.predict(batch_x, enable_dropout=True, out_prob=True)
# _, batch_output = torch.max(out, dim=1)
volume_pred.append(out)
volume_pred = torch.cat(volume_pred)
_, volume_prediction = torch.max(volume_pred, dim=1)
volume_prediction = (volume_prediction.cpu().numpy()).astype('float32')
volume_prediction = np.squeeze(volume_prediction)
if orientation == "COR":
volume_prediction = volume_prediction.transpose((1, 2, 0))
volume_pred = volume_pred.permute((2, 1, 3, 0))
elif orientation == "AXI":
volume_prediction = volume_prediction.transpose((2, 0, 1))
volume_pred = volume_pred.permute((3, 1, 0, 2))
mc_pred_list.append(volume_prediction)
if j == 0:
expected_pred = (1 / mc_samples) * volume_pred
else:
expected_pred += (1 / mc_samples) * volume_pred
_, final_seg = torch.max(expected_pred, dim=1)
final_seg = (final_seg.cpu().numpy()).astype('float32')
final_seg = np.squeeze(final_seg)
return expected_pred, final_seg, mc_pred_list, header
def compute_vol_bulk(prediction_dir, dir_struct, label_names, volumes_txt_file):
log.info("**Computing volume estimates**")
with open(volumes_txt_file) as file_handle:
volumes_to_use = file_handle.read().splitlines()
file_paths = du.load_file_paths_eval(prediction_dir, volumes_txt_file, dir_struct)
volume_dict_list = []
for vol_idx, file_path in enumerate(file_paths):
volume_prediction, header = du.load_and_preprocess_eval(file_path, "SAG", notlabel=False)
per_volume_dict = compute_volume(volume_prediction, label_names, volumes_to_use[vol_idx])
volume_dict_list.append(per_volume_dict)
_write_csv_table('volume_estimates.csv', prediction_dir, volume_dict_list, label_names)
log.info("**DONE**")
def evaluate(coronal_model_path, volumes_txt_file, data_dir, device, prediction_path, batch_size, orientation,
label_names, dir_struct, need_unc=False, mc_samples=0, exit_on_error=False):
log.info("**Starting evaluation**")
with open(volumes_txt_file) as file_handle:
volumes_to_use = file_handle.read().splitlines()
cuda_available = torch.cuda.is_available()
# First, are we attempting to run on a GPU?
if type(device) == int:
# if CUDA available, follow through, else warn and fallback to CPU
if cuda_available:
model = torch.load(coronal_model_path)
torch.cuda.empty_cache()
model.cuda(device)
else:
log.warning(
'CUDA is not available, trying with CPU. ' + \
'This can take much longer (> 1 hour). Cancel and ' + \
'investigate if this behavior is not desired.'
)
# switch device to 'cpu'
device = 'cpu'
# If device is 'cpu' or CUDA not available
if (type(device)==str) or not cuda_available:
model = torch.load(
coronal_model_path,
map_location=torch.device(device)
)
model.eval()
common_utils.create_if_not(prediction_path)
log.info("Evaluating now...")
file_paths = du.load_file_paths_eval(data_dir, volumes_txt_file, dir_struct)
with torch.no_grad():
volume_dict_list = []
cvs_dict_list = []
iou_dict_list = []
for vol_idx, file_path in enumerate(file_paths):
try:
if need_unc == "True":
_, volume_prediction, mc_pred_list, header = _segment_vol_unc(file_path, model, orientation,
batch_size, mc_samples,
cuda_available, device)
iou_dict, cvs_dict = compute_structure_uncertainty(mc_pred_list, label_names,
volumes_to_use[vol_idx])
cvs_dict_list.append(cvs_dict)
iou_dict_list.append(iou_dict)
else:
_, volume_prediction, header = _segment_vol(file_path, model, orientation, batch_size,
cuda_available,
device)
#Copy header affine
Mat = np.array([
header['srow_x'],
header['srow_y'],
header['srow_z'],
[0,0,0,1]
])
# Apply original image affine to prediction volume
nifti_img = nib.Nifti1Image(volume_prediction, Mat, header=header)
log.info("Processed: " + volumes_to_use[vol_idx] + " " + str(vol_idx + 1) + " out of " + str(
len(file_paths)))
save_file = os.path.join(prediction_path, volumes_to_use[vol_idx])
if '.nii' not in save_file:
save_file += '.nii.gz'
nib.save(nifti_img, save_file)
per_volume_dict = compute_volume(volume_prediction, label_names, volumes_to_use[vol_idx])
volume_dict_list.append(per_volume_dict)
except FileNotFoundError as exp:
log.error("Error in reading the file ...")
log.exception(exp)
if exit_on_error:
raise(exp)
except Exception as exp:
log.exception(exp)
if exit_on_error:
raise(exp)
_write_csv_table('volume_estimates.csv', prediction_path, volume_dict_list, label_names)
if need_unc == "True":
_write_csv_table('cvs_uncertainty.csv', prediction_path, cvs_dict_list, label_names)
_write_csv_table('iou_uncertainty.csv', prediction_path, iou_dict_list, label_names)
log.info("DONE")
def evaluate2view(coronal_model_path, axial_model_path, volumes_txt_file, data_dir, device, prediction_path, batch_size,
label_names, dir_struct, need_unc=False, mc_samples=0, exit_on_error=False):
log.info("**Starting evaluation**")
with open(volumes_txt_file) as file_handle:
volumes_to_use = file_handle.read().splitlines()
cuda_available = torch.cuda.is_available()
if type(device) == int:
# if CUDA available, follow through, else warn and fallback to CPU
if cuda_available:
model1 = torch.load(coronal_model_path)
model2 = torch.load(axial_model_path)
torch.cuda.empty_cache()
model1.cuda(device)
model2.cuda(device)
else:
log.warning(
'CUDA is not available, trying with CPU.' + \
'This can take much longer (> 1 hour). Cancel and ' + \
'investigate if this behavior is not desired.'
)
if (type(device)==str) or not cuda_available:
model1 = torch.load(
coronal_model_path,
map_location=torch.device(device)
)
model2 = torch.load(
axial_model_path,
map_location=torch.device(device)
)
model1.eval()
model2.eval()
common_utils.create_if_not(prediction_path)
log.info("Evaluating now...")
file_paths = du.load_file_paths_eval(data_dir, volumes_txt_file, dir_struct)
with torch.no_grad():
volume_dict_list = []
cvs_dict_list = []
iou_dict_list = []
for vol_idx, file_path in enumerate(file_paths):
try:
if need_unc == "True":
volume_prediction_cor, _, mc_pred_list_cor, header = _segment_vol_unc(file_path, model1, "COR",
batch_size, mc_samples,
cuda_available, device)
volume_prediction_axi, _, mc_pred_list_axi, header = _segment_vol_unc(file_path, model2, "AXI",
batch_size, mc_samples,
cuda_available, device)
mc_pred_list = mc_pred_list_cor + mc_pred_list_axi
iou_dict, cvs_dict = compute_structure_uncertainty(mc_pred_list, label_names,
volumes_to_use[vol_idx])
cvs_dict_list.append(cvs_dict)
iou_dict_list.append(iou_dict)
else:
volume_prediction_cor, _, header = _segment_vol(file_path, model1, "COR", batch_size,
cuda_available,
device)
volume_prediction_axi, _, header = _segment_vol(file_path, model2, "AXI", batch_size,
cuda_available,
device)
_, volume_prediction = torch.max(volume_prediction_axi + volume_prediction_cor, dim=1)
volume_prediction = (volume_prediction.cpu().numpy()).astype('float32')
volume_prediction = np.squeeze(volume_prediction)
#Copy header affine
Mat = np.array([
header['srow_x'],
header['srow_y'],
header['srow_z'],
[0,0,0,1]
])
# Apply original image affine to prediction volume
nifti_img = nib.Nifti1Image(volume_prediction, Mat, header=header)
log.info("Processed: " + volumes_to_use[vol_idx] + " " + str(vol_idx + 1) + " out of " + str(
len(file_paths)))
nib.save(nifti_img, os.path.join(prediction_path, volumes_to_use[vol_idx] + str('.nii.gz')))
per_volume_dict = compute_volume(volume_prediction, label_names, volumes_to_use[vol_idx])
volume_dict_list.append(per_volume_dict)
except FileNotFoundError as exp:
log.error("Error in reading the file ...")
log.exception(exp)
if exit_on_error:
raise(exp)
except Exception as exp:
log.exception(exp)
if exit_on_error:
raise(exp)
# log.info("Other kind o error!")
_write_csv_table('volume_estimates.csv', prediction_path, volume_dict_list, label_names)
if need_unc == "True":
_write_csv_table('cvs_uncertainty.csv', prediction_path, cvs_dict_list, label_names)
_write_csv_table('iou_uncertainty.csv', prediction_path, iou_dict_list, label_names)
log.info("DONE")
