def image_transform(self, fns, fn_paras=None):
"""
Transform the points of a geometry using the computed transformation
Args:
fns: dictionary of the file names: in_im, in_label, out_im, out_label
fn_paras: file name to the parameter map of previously done registration
Returns:
new_label: transformed label
"""
self.readParameterMap(fn_paras)
move_im = preProcess.resample_spacing(sitk.ReadImage(fns['in_im']))[0]
# wrap point set
transformixImageFilter = sitk.TransformixImageFilter()
transformixImageFilter.SetMovingImage(move_im)
transformixImageFilter.SetTransformParameterMap(self.parameter_map)
transformixImageFilter.Execute()
sitk.WriteImage(transformixImageFilter.GetResultImage(), fns['out_im'])
label_im = preProcess.resample_spacing(sitk.ReadImage(
fns['in_label']))[0]
transformixImageFilter = sitk.TransformixImageFilter()
transformixImageFilter.SetMovingImage(label_im)
transformixImageFilter.SetTransformParameterMap(self.parameter_map)
transformixImageFilter.Execute()
sitk.WriteImage(transformixImageFilter.GetResultImage(),
fns['out_label'])
def preProcessImage(img_fn, mask_fn, m, view, size=None):
from skimage.transform import resize
"""
pre-process image and labels to feed into models
"""
img, _ = resample_spacing(img_fn, order=1)
mask, _ = resample_spacing(mask_fn, order=0)
img_vol = sitk.GetArrayFromImage(img)
label_vol = sitk.GetArrayFromImage(mask)
img_vol = RescaleIntensity(img_vol, m, [750, -750])
#remove blank slices
label_vol = np.moveaxis(label_vol, view, 0)
IDs = np.max(np.max(label_vol, axis=-1), axis=-1) == 0
sliced = np.moveaxis(img_vol, view, 0)[~IDs, :, :]
if size is not None: #rescale the view axis
shape = (size, 256, 256)
else:
shape = (sliced.shape[0], 256, 256)
sliced = np.moveaxis(resize(sliced, shape, order=1), 0, view)
mask_sliced = label_vol[~IDs, :, :]
mask_sliced = np.moveaxis(resize(mask_sliced, shape, order=1), 0, view)
return sliced, mask_sliced
def blankSpaces(y_train_filenames_ct, y_train_filenames_mr):
num = len(y_train_filenames_ct)
ratio_mr = []
ratio_ct = []
for i in range(num):
im_mr = resample_spacing(y_train_filenames_mr[i])[0]
mask_mr, ratio = cropMask(sitk.GetArrayFromImage(im_mr), 1.)
print(ratio)
ratio_mr = np.concatenate((ratio_mr, ratio))
im_ct = resample_spacing(y_train_filenames_ct[i])[0]
mask_ct, ratio = cropMask(sitk.GetArrayFromImage(im_ct), 1.)
print(ratio)
ratio_ct = np.concatenate((ratio_ct, ratio))
return ratio_ct, ratio_mr
def resample_prediction(im, orig_im):
#resample prediction so it matches the original image
im_info = resample_spacing(orig_im, order=1)[0]
im.SetSpacing(im_info.GetSpacing())
im.SetOrigin(im_info.GetOrigin())
im.SetDirection(im_info.GetDirection())
return centering(im, orig_im, order=1)
def evaluate_assd(self):
import vtk
from vtk.util.numpy_support import vtk_to_numpy, numpy_to_vtk
def _get_assd(p_surf, g_surf):
dist_fltr = vtk.vtkDistancePolyDataFilter()
dist_fltr.SetInputData(1, p_surf)
dist_fltr.SetInputData(0, g_surf)
dist_fltr.SignedDistanceOff()
dist_fltr.Update()
distance_poly = vtk_to_numpy(
dist_fltr.GetOutput().GetPointData().GetArray(0))
return np.mean(distance_poly), dist_fltr.GetOutput()
ref_im = sitk.ReadImage(self.label_fn)
ref_im = resample_spacing(ref_im,
template_size=(256, 256, 256),
order=0)[0]
ref_im, M = exportSitk2VTK(ref_im)
pred_im = resample_spacing(self.pred,
template_size=(256, 256, 256),
order=0)[0]
pred_im, M = exportSitk2VTK(pred_im)
ref_im_py = swapLabels(vtk_to_numpy(
ref_im.GetPointData().GetScalars()))
ref_im.GetPointData().SetScalars(numpy_to_vtk(ref_im_py))
ids = np.unique(ref_im_py)
pred_poly_list = []
dist_poly_list = []
ref_poly_list = []
dist = [0.] * (len(ids) + 1)
for index, i in enumerate(ids):
if i == 0:
continue
p_s = vtk_marching_cube(pred_im, 0, i)
r_s = vtk_marching_cube(ref_im, 0, i)
dist_ref2pred, d_ref2pred = _get_assd(p_s, r_s)
dist_pred2ref, d_pred2ref = _get_assd(r_s, p_s)
dist[index + 1] = (dist_ref2pred + dist_pred2ref) * 0.5
dist_poly_list.append(d_pred2ref)
pred_poly_list.append(p_s)
ref_poly_list.append(r_s)
dist_poly = appendPolyData(dist_poly_list)
pred_poly = appendPolyData(pred_poly_list)
ref_poly = appendPolyData(ref_poly_list)
dist[0], _ = _get_assd(pred_poly, ref_poly)
return dist
def _augment(train_filenames,
range_adjust,
ratio,
scale_factor=None,
order=1):
if scale_factor is None:
scale_factor = np.mean([
sample_in_range(range_adjust[0]),
sample_in_range(range_adjust[1]),
sample_in_range(range_adjust[2])
] / ratio)
sitkIm, ref_img = resample_spacing(train_filenames, order=order)
image = resample_scale(sitkIm, ref_img, scale_factor, order)
return image, scale_factor
def tf_test(img_fn, label_fn):
import tensorflow as tf
img = RescaleIntensity(sitk.GetArrayFromImage(preProcess.resample_spacing(img_fn)[0]), "mr", [750, -750])
label = sitk.GetArrayFromImage(preProcess.resample_spacing(label_fn, order=0)[0])
label[label==421] = 420
tf_img = tf.placeholder(tf.float32, shape=img.shape)
tf_label = tf.placeholder(tf.int32, shape=label.shape)
tr_img_aug, tr_label_aug= tf_intensity_augmentation(tf_img, tf_label, 8, changeIntensity=True)
with tf.Session() as sess:
out_im, out_label = sess.run([tr_img_aug, tr_label_aug], feed_dict={tf_img: img, tf_label: label})
fig, axes = plt.subplots(2,3)
for i in range(axes.shape[0]):
loc = 50
incr = 40
axes[i][0].imshow(img[loc+incr*(i+1),:,:], cmap='gray')
axes[i][0].axis('off')
axes[i][1].imshow(out_im[loc+incr*(i+1),:,:],cmap='gray')
axes[i][1].axis('off')
axes[i][2].imshow(out_label[loc+incr*(i+1),:,:],cmap='gray')
axes[i][2].axis('off')
plt.show()
return out_im, out_label
def volume_prediction_average(self, size):
self.input_size = size
img_vol = resample_spacing(self.image_fn,
order=1,
template_size=(size, size, size))[0]
self.image_info = {}
self.image_info['spacing'] = img_vol.GetSpacing()
self.image_info['origin'] = img_vol.GetOrigin()
self.image_info['direction'] = img_vol.GetDirection()
img_vol = sitk.GetArrayFromImage(img_vol)
img_vol = RescaleIntensity(img_vol, self.modality, [750, -750])
self.original_shape = img_vol.shape
prob = np.zeros((*self.original_shape, 8))
unique_views = np.unique(self.views)
self.pred_time = 0.
for view in unique_views:
indices = np.where(self.views == view)[0]
predict_shape = [size, size, size, 8]
predict_shape[view] = img_vol.shape[view]
prob_view = np.zeros(predict_shape)
for i in indices:
model_path = self.models[i]
(self.unet).load_weights(model_path)
p, t = model_output_no_resize(self.unet, img_vol,
self.views[i], self.channel)
prob_view += p
self.pred_time += t
prob += prob_view
avg = prob / len(self.models)
self.pred = predictVol(avg, np.zeros(1))
return
def main(args):
modality = args.modality
data_folder = args.image_dir
fdr_postfix = args.folder_attr
output_data_folder = args.output
aug_num = args.aug_num
# TO-DO: NEED TO CHECK/FIX THIS BEFORE NEXT RUNS
if fdr_postfix == '_train':
ids = range(0, 17)
elif fdr_postfix == '_val':
#ids = [0, 19]
ids = range(17, 20)
#ids = range(0,4)
for m in modality:
try:
os.makedirs(os.path.join(output_data_folder, m + fdr_postfix))
except Exception as e:
print(e)
try:
os.makedirs(
os.path.join(output_data_folder, m + fdr_postfix + '_masks'))
except Exception as e:
print(e)
filenames_dic = {}
for m in modality:
x_train_filenames, y_train_filenames = getTrainNLabelNames(
data_folder, m)
print("Number of training volumes %d" % len(x_train_filenames))
print("Number of mask volumes %d" % len(y_train_filenames))
filenames_dic[m + '_x'] = x_train_filenames
filenames_dic[m + '_y'] = y_train_filenames
#find the blank spaces for image volumes
ratios = {}
ratios['ct'], ratios['mr'] = blankSpaces(filenames_dic['ct_y'],
filenames_dic['mr_y'])
#debug, show scale plots:
#import pandas as pd
#df = pd.DataFrame.from_dict(ratios)
#plt.figure()
#df.boxplot()
#plt.show()
#import sys
#sys.exit()
for m in modality:
ratios[m] = ratios[m].reshape(int(len(ratios[m]) / 3), 3)
#assume uniform distribution, find the range of blank space ratios
dim = 3
range_adjust = [None] * dim
for i in range(dim):
range_adjust[i] = [
np.min([np.min(ratios[m][:, i]) for m in modality]),
np.max([np.max(ratios[m][:, i]) for m in modality])
]
print("The ratio of blank space found for ct and mr is: ", range_adjust)
# Apply data augmentation so that the scale, spacing, orientation of MR and CT volumes are consistant
def _augment(train_filenames,
range_adjust,
ratio,
scale_factor=None,
order=1):
if scale_factor is None:
scale_factor = np.mean([
sample_in_range(range_adjust[0]),
sample_in_range(range_adjust[1]),
sample_in_range(range_adjust[2])
] / ratio)
sitkIm, ref_img = resample_spacing(train_filenames, order=order)
image = resample_scale(sitkIm, ref_img, scale_factor, order)
return image, scale_factor
for m in modality:
num = len(filenames_dic[m + '_x'])
print(filenames_dic[m + '_x'])
for i in ids:
print("ID: ", i)
fn = os.path.join(output_data_folder, m + fdr_postfix,
os.path.basename(filenames_dic[m + '_x'][i]))
img, _ = resample_spacing(filenames_dic[m + '_x'][i], order=1)
sitk.WriteImage(img, fn)
fn = os.path.join(output_data_folder, m + fdr_postfix + '_masks',
os.path.basename(filenames_dic[m + '_y'][i]))
mask, _ = resample_spacing(filenames_dic[m + '_y'][i], order=0)
sitk.WriteImage(mask, fn)
for j in range(aug_num):
mask, scale_factor = _augment(filenames_dic[m + '_y'][i],
range_adjust,
ratios[m][i, :],
order=0)
fn = os.path.join(
output_data_folder, m + fdr_postfix + '_masks',
m + '_aug_' + str(i) + '_' + str(j) + '_label.nii.gz')
sitk.WriteImage(mask, fn)
img, _ = _augment(filenames_dic[m + '_x'][i],
range_adjust,
ratios[m][i, :],
scale_factor=scale_factor,
order=1)
#apply intensity augmentation
if args.intensity:
print("Applying intensity augmentation!")
py_img, _ = apply_intensity_map(
sitk.GetArrayFromImage(img),
sitk.GetArrayFromImage(mask))
img = sitk.GetImageFromArray(py_img)
img.SetOrigin(mask.GetOrigin())
img.SetDirection(mask.GetDirection())
img.SetSpacing(mask.GetSpacing())
fn = os.path.join(
output_data_folder, m + fdr_postfix,
m + '_aug_' + str(i) + '_' + str(j) + '_image.nii.gz')
sitk.WriteImage(img, fn)
def loadImages(self):
self.fixed = preProcess.resample_spacing(sitk.ReadImage(
self.fixed_fn))[0]
self.moving = preProcess.resample_spacing(
sitk.ReadImage(self.moving_fn))[0]
args = parser.parse_args()
fns = glob.glob(os.path.join(args.folder, '*.nii.gz')) + glob.glob(
os.path.join(args.folder, '*.nii'))
try:
os.makedirs(os.path.join(args.out_folder))
except:
pass
spacing = []
spacing_ori = []
for fn in fns:
img = sitk.ReadImage(fn)
spacing_ori.append(img.GetSpacing())
img = resample_spacing(img, order=1)[0]
spacing.append(img.GetSpacing())
fn_out = os.path.join(args.out_folder, os.path.basename(fn))
sitk.WriteImage(img, fn_out)
spacing = np.array(spacing)
spacing_ori = np.array(spacing_ori)
print("Axial: ", np.mean(spacing[:, 0]), np.min(spacing[:, 0]),
np.max(spacing[:, 0]))
print("Sagittal: ", np.mean(spacing[:, 1]), np.min(spacing[:, 1]),
np.max(spacing[:, 1]))
print("Coronal: ", np.mean(spacing[:, 2]), np.min(spacing[:, 2]),
np.max(spacing[:, 2]))
print("Axial: ", np.mean(spacing_ori[:, 0]), np.min(spacing_ori[:, 0]),
np.max(spacing_ori[:, 0]))
print("Sagittal: ", np.mean(spacing_ori[:, 1]), np.min(spacing_ori[:, 1]),
imgVol = sorted(glob.glob(os.path.join(args.folder, "*.nii.gz")))
print(imgVol)
try:
os.makedirs(args.out_folder)
except Exception as e:
print(e)
def resample_prediction(im, orig_im):
#resample prediction so it matches the original image
im_info = resample_spacing(orig_im, order=1)[0]
im.SetSpacing(im_info.GetSpacing())
im.SetOrigin(im_info.GetOrigin())
im.SetDirection(im_info.GetDirection())
return centering(im, orig_im, order=1)
for im_fn in imgVol:
print(im_fn)
im_ori = resample_spacing(im_fn, order=1)[0]
im = sitk.GetArrayFromImage(im_ori)
im = RescaleIntensity(im, args.modality, args.intensity)
print(np.max(im), np.min(im), im.shape, im_ori.GetSpacing())
out_im_fn = os.path.join(args.out_folder, os.path.basename(im_fn))
im = sitk.GetImageFromArray(im)
im_2 = resample_prediction(im, sitk.ReadImage(im_fn))
sitk.WriteImage(im, out_im_fn)
sitk.WriteImage(im_2, out_im_fn)
