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:
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 compute_vol_bulk(prediction_dir, dir_struct, label_names, volumes_txt_file):
print("**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)
print("**DONE**")
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:
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