def save_pat_objects(pat_id, phase_id, segmentation, seg_errors, uncertainties, aleatoric,
type_of_map, spacing,
output_dirs, new_spacing=None, do_resample=False, direction=None, origin=None,
pred_probs=None):
if do_resample and new_spacing is None:
raise ValueError("ERROR - if resample tue then new_spacing cannot be None")
if phase_id is not None:
f_prefix = '{}_{}'.format(pat_id, phase_id)
else:
f_prefix = '{}'.format(pat_id)
fname = path.join(output_dirs['pred_labels'], f_prefix + '.nii.gz')
if do_resample:
segmentation = apply_2d_zoom_3d(segmentation.astype(np.int16), spacing, order=0, do_blur=False,
as_type=np.int16, new_spacing=new_spacing)
# print("Resample segmentations ", o_shape, segmentation.shape)
save_nifty(segmentation.astype(np.int16), new_spacing if do_resample else spacing, fname,
direction=direction, origin=origin)
print("INFO - Saved auto mask to {} (sp/new_sp)".format(fname), spacing, new_spacing)
fname = path.join(output_dirs['umaps'], f_prefix + '_{}.nii.gz'.format(type_of_map))
if do_resample:
uncertainties = apply_2d_zoom_3d(uncertainties.astype(np.float32), spacing, do_blur=True,
new_spacing=new_spacing)
# print("Resample uncertainties ", o_shape, uncertainties.shape)
save_nifty(uncertainties.astype(np.float32), new_spacing if do_resample else spacing, fname, direction=direction,
origin=origin)
if seg_errors is not None:
fname = path.join(output_dirs['errors'], f_prefix + '.nii.gz')
if do_resample:
seg_errors = apply_2d_zoom_3d(seg_errors.astype(np.int16), spacing, order=0, do_blur=False,
as_type=np.int16, new_spacing=new_spacing)
save_nifty(seg_errors.astype(np.int16), new_spacing if do_resample else spacing, fname, direction=direction,
origin=origin)
if aleatoric is not None:
fname = path.join(output_dirs['umaps'], f_prefix + '_aleatoric.nii.gz')
if do_resample:
aleatoric = apply_2d_zoom_3d(aleatoric.astype(np.float32), spacing, do_blur=True, new_spacing=new_spacing)
save_nifty(aleatoric.astype(np.float32), new_spacing if do_resample else spacing, fname, direction=direction,
origin=origin)
if pred_probs is not None:
fname = path.join(output_dirs['pred_probs'], f_prefix + '.nii.gz')
pred_new_spacing = np.array([new_spacing[0], 1, new_spacing[1], new_spacing[2]])
pred_spacing = np.array([spacing[0], 1, spacing[1], spacing[2]])
if do_resample:
nclasses = pred_probs.shape[1]
resampled_probs = np.zeros((segmentation.shape[0], pred_probs.shape[1], segmentation.shape[1],
segmentation.shape[2]))
for cls_idx in np.arange(nclasses):
resampled_probs[:, cls_idx] = apply_2d_zoom_3d(pred_probs[:, cls_idx].astype(np.float32), spacing,
do_blur=False, new_spacing=new_spacing)
pred_probs = resampled_probs
print("WARNING do resample ", do_resample)
save_nifty(pred_probs.astype(np.float32), pred_new_spacing if do_resample else pred_spacing, fname, direction=direction,
origin=origin)
def _do_resample(self, img, gt_lbl, sp):
img = apply_2d_zoom_3d(img,
sp,
do_blur=True,
new_spacing=self.new_spacing)
gt_lbl = apply_2d_zoom_3d(gt_lbl,
sp,
order=0,
do_blur=False,
as_type=np.int,
new_spacing=self.new_spacing)
return img, gt_lbl
def save_segmentations(segmentation,
pat_id,
output_dirs,
spacing,
do_resample,
new_spacing=None,
frame_id=None,
direction=None,
origin=None):
if frame_id is not None:
f_prefix = '{}_{:02d}'.format(pat_id, frame_id)
else:
f_prefix = '{}'.format(pat_id)
fname = path.join(output_dirs['pred_labels'], f_prefix + '.nii.gz')
if do_resample:
if segmentation.ndim == 3:
segmentation = apply_2d_zoom_3d(segmentation.astype(np.int16),
spacing,
order=0,
do_blur=False,
as_type=np.int16,
new_spacing=new_spacing)
else:
# assuming 4d data
num_time_points = segmentation.shape[0]
resized_segs = None
for t in np.arange(num_time_points):
seg = segmentation[t]
new_seg = apply_2d_zoom_3d(seg.astype(np.int16),
spacing,
order=0,
do_blur=False,
as_type=np.int16,
new_spacing=new_spacing)
if resized_segs is None:
resized_segs = np.expand_dims(new_seg, axis=0)
else:
resized_segs = np.vstack(
(resized_segs, np.expand_dims(new_seg, axis=0)))
# print("Resample segmentations ", o_shape, segmentation.shape)
save_nifty(segmentation.astype(np.int16),
new_spacing if do_resample else spacing,
fname,
direction=direction,
origin=origin)
def get_pred_probs(src_data_path,
cardiac_phase,
mc_dropout=False,
patient_id=None,
meta_info=True):
"""
returns softmax probabilities as numpy array of shape [z, nclasses, y, x] and spacing
"""
search_path = get_search_mask(src_data_path, "pred_probs", cardiac_phase,
mc_dropout, patient_id)
pred_probs = {}
file_list = glob.glob(search_path)
if len(file_list) == 0:
raise ValueError("ERROR - No search result for {}".format(search_path))
file_list.sort()
for fname in file_list:
pat_id, _ = (os.path.splitext(os.path.basename(fname))[0]).split("_")
img = sitk.ReadImage(fname)
org_sp = img.GetSpacing()[::-1]
softmax_probs = sitk.GetArrayFromImage(img).astype(np.float32)
# We store pred probs in the original spacing (according to the image info).
# Nevertheless, in case the spacing of x, y is below 1mm we resample to 1.4 isotropic. We do the same
# in the ACDC dataset object
if org_sp[-1] < 1.:
nclasses = softmax_probs.shape[1]
t_sp = np.array([org_sp[0], 1.4, 1.4]).astype(np.float64)
t_org_sp = np.array([org_sp[0], org_sp[2],
org_sp[3]]).astype(np.float64)
spacing = np.array([org_sp[0], org_sp[1], 1.4,
1.4]).astype(np.float64)
resampled_probs = []
for cls_idx in np.arange(nclasses):
resampled_probs.append(
apply_2d_zoom_3d(softmax_probs[:,
cls_idx].astype(np.float32),
t_org_sp,
do_blur=False,
new_spacing=t_sp))
softmax_probs = np.concatenate(
([arr[:, None] for arr in resampled_probs]), axis=1)
else:
spacing = org_sp
if meta_info:
pred_probs[pat_id] = {
'pred_probs': softmax_probs,
'spacing': spacing,
'original_spacing': org_sp
}
else:
pred_probs[pat_id] = softmax_probs
del softmax_probs
return pred_probs
def get_images_nifti(data_src_dir, patid=None, resample=False, rescale=False):
if patid is None:
search_path = os.path.join(data_src_dir, "*.nii.gz")
else:
search_path = os.path.join(data_src_dir, patid + "*.nii.gz")
file_list = glob.glob(search_path)
if len(file_list) == 0:
raise ValueError("ERROR - No search result for {}".format(search_path))
else:
print("INFO - Found {} nifti files in {}".format(
len(file_list), data_src_dir))
file_list.sort()
for fname in file_list:
img = sitk.ReadImage(fname)
img_arr = sitk.GetArrayFromImage(img).astype(np.float32)
pat_id, _ = (os.path.splitext(os.path.basename(fname))[0]).split(".")
if img_arr.ndim == 4:
num_of_frames = img_arr.shape[0]
spacing = img.GetSpacing()[::-1][1:]
original_spacing = img.GetSpacing()[::-1][1:]
else:
num_of_frames = 1
img_arr = np.expand_dims(img_arr, axis=0)
spacing = img.GetSpacing()[::-1]
original_spacing = img.GetSpacing()[::-1]
print("INFO - Processing patid {}, #frames {}".format(
pat_id, num_of_frames))
if resample or spacing[-1] < 1.:
do_resample = True
new_spacing = np.array([original_spacing[0], 1.4,
1.4]).astype(np.float32)
spacing = new_spacing
else:
do_resample = False
new_spacing = None
for frame_id in range(num_of_frames):
img_np = img_arr[frame_id]
if do_resample:
img_np = apply_2d_zoom_3d(img_np,
original_spacing,
new_spacing,
do_blur=True)
if rescale:
img_np = rescale_intensities(img_np).astype(np.float32)
yield {
'image': img_np,
'spacing': spacing,
'reference': img_np,
'patient_id': pat_id,
'frame_id': frame_id,
'cardiac_phase': 'none',
'structures': [],
'original_spacing': original_spacing,
'num_of_frames': num_of_frames
}
def acdc_validation_fold_image4d(fold,
root_dir,
patid=None,
resample=False,
rescale=False,
file_suffix='4d.nii.gz'):
if isinstance(patid, str):
patid = int(patid.strip('patient'))
print("acdc_validation_fold_image4d ", fold)
allpatnumbers = np.arange(1, 101)
foldmask = np.tile(np.arange(4)[::-1].repeat(5), 5)
validation_nums = allpatnumbers[foldmask == fold]
if patid is not None:
validation_nums = [pid for pid in validation_nums if patid == pid]
if len(validation_nums) == 0:
raise ValueError(
"ERROR - acdc validation fold - patid {} not found in validation set"
.format(patid))
for patnum in validation_nums:
img = ACDCImage(patnum,
root_dir=root_dir,
resample=True,
scale_intensities=False,
resample_zaxis=False,
file_suffix=file_suffix)
img4d_arr = img.image4d()
num_of_frames = img4d_arr.shape[0]
spacing = img.voxel_spacing()[1:]
original_spacing = img.voxel_spacing()[1:]
for frame_id in range(num_of_frames):
img_np = img4d_arr[frame_id]
if resample:
print("WARNING - Resampling")
spacing = np.array([
original_spacing[0], ACDCImage.new_spacing[1],
ACDCImage.new_spacing[2]
]).astype(np.float32)
img_np = apply_2d_zoom_3d(img_np,
spacing,
do_blur=True,
new_spacing=img.new_spacing)
if rescale:
img_np = rescale_intensities(img_np).astype(np.float32)
yield {
'image': img_np,
'spacing': spacing,
'reference': img_np,
'patient_id': img.patient_id,
'frame_id': frame_id,
'cardiac_phase': 'ED',
'structures': [],
'original_spacing': original_spacing,
'num_of_frames': num_of_frames
}
def get_umaps(src_data_path, cardiac_phase, mc_dropout=False, patient_id=None):
"""
returns uncertainty map (e- or b-map) as numpy array of shape [z, y, x] and spacing
"""
if mc_dropout:
input_dir = "_bmap" + file_ext
else:
input_dir = "_emap" + file_ext
if patient_id is None:
input_dir = os.path.sep + "*_{}".format(cardiac_phase) + input_dir
else:
input_dir = os.path.sep + patient_id + "_{}".format(
cardiac_phase) + input_dir
search_path = src_data_path + os.sep + OUTPUT_DIRS[
'umaps'] + os.sep + input_dir
umaps = {}
file_list = glob.glob(search_path)
if len(file_list) == 0:
raise ValueError("ERROR - No search result for {}".format(search_path))
file_list.sort()
for fname in file_list:
pat_id, _, _ = (os.path.splitext(
os.path.basename(fname))[0]).split("_")
img = sitk.ReadImage(fname)
org_sp = img.GetSpacing()[::-1]
data = sitk.GetArrayFromImage(img).astype(np.float32)
# hacky!!! For entropy maps it happens that some voxels have tiny negative values e-08.
# But we need to make sure these are zero
# data[data < 0] = 0
# We store pred labels (and uncertainty maps) in the original spacing (according to the image info).
# Nevertheless, in case the spacing of x, y is below 1mm we resample to 1.4 isotropic. We do the same
# in the ACDC dataset object
if org_sp[-1] < 1.:
spacing = np.array([org_sp[0], 1.4, 1.4]).astype(np.float64)
data = apply_2d_zoom_3d(data,
org_sp,
do_blur=True,
new_spacing=spacing)
else:
spacing = org_sp
umaps[pat_id] = {
'umap': data,
'spacing': spacing,
'original_spacing': org_sp
}
del data
return umaps
def get_pred_labels(src_data_path,
cardiac_phase,
mc_dropout=False,
patient_id=None,
one_hot=False,
meta_info=True):
"""
returns predicted labels as numpy array of shape [z, y, x] and spacing
"""
search_path = get_search_mask(src_data_path, "pred_labels", cardiac_phase,
mc_dropout, patient_id)
pred_labels = {}
file_list = glob.glob(search_path)
if len(file_list) == 0:
raise ValueError("ERROR - No search result for {}".format(search_path))
file_list.sort()
for fname in file_list:
pat_id, _ = (os.path.splitext(os.path.basename(fname))[0]).split("_")
img = sitk.ReadImage(fname)
org_sp = img.GetSpacing()[::-1]
predictions = sitk.GetArrayFromImage(img).astype(np.long)
# We store pred labels (and uncertainty maps) in the original spacing (according to the image info).
# Nevertheless, in case the spacing of x, y is below 1mm we resample to 1.4 isotropic. We do the same
# in the ACDC dataset object
if org_sp[-1] < 1.:
spacing = np.array([org_sp[0], 1.4, 1.4]).astype(np.float64)
predictions = apply_2d_zoom_3d(predictions,
org_sp,
order=0,
do_blur=False,
as_type=np.int,
new_spacing=spacing)
else:
spacing = org_sp
if one_hot:
predictions = one_hot_encoding(predictions)
if meta_info:
pred_labels[pat_id] = {
'pred_labels': predictions,
'spacing': spacing,
'original_spacing': org_sp
}
else:
pred_labels[pat_id] = predictions
del predictions
return pred_labels
def save_pred_probs(pat_id, phase_id, pred_probs, spacing,
output_dirs, new_spacing=None, do_resample=False, direction=None, origin=None):
if phase_id is not None:
f_prefix = '{}_{}'.format(pat_id, phase_id)
else:
f_prefix = '{}'.format(pat_id)
fname = path.join(output_dirs['pred_probs'], f_prefix + '.nii.gz')
pred_new_spacing = np.array([new_spacing[0], 1, new_spacing[1], new_spacing[2]])
pred_spacing = np.array([spacing[0], 1, spacing[1], spacing[2]])
if do_resample:
nclasses = pred_probs.shape[1]
resampled_probs = None
for cls_idx in np.arange(nclasses):
new_probs = apply_2d_zoom_3d(pred_probs[:, cls_idx].astype(np.float32), spacing,
do_blur=False, new_spacing=new_spacing)
if resampled_probs is None:
resampled_probs = np.zeros((new_probs.shape[0], pred_probs.shape[1], new_probs.shape[1],
new_probs.shape[2]))
resampled_probs[:, cls_idx] = new_probs
pred_probs = resampled_probs
save_nifty(pred_probs.astype(np.float32), pred_new_spacing if do_resample else pred_spacing, fname, direction=direction,
origin=origin)