def test_3d(self, data_unit):
""" Test Thread 3d segmentation
Assumes images are already prepared
"""
""" Get Data from data_unit
"""
[im_path, gt_path, im, gt] = data_unit
print im_path, gt_path
filename, ext = osp.splitext(osp.split(im_path)[1])
filename = '{}_pred{}'.format(filename, ext)
""" Do Forward
"""
prediction_map = self.do_forward_3d(im, self.params.CHUNK_SHAPE,
self.params.STRIDE)
# transform the datatype to unint8
prediction_map = prediction_map.astype(np.uint8)
# fig, axes = plt.subplots(1, 5)
# for slices in xrange(0,prediction_map.shape[2]):
# vis_seg(axes, [im[:,:,slices], gt[:,:,slices], prediction_map[:,:,slices]])
""" Save
"""
# construct output_path
output_path = osp.join(self.output_dir, filename)
print('Save to {}'.format(output_path))
save_data(prediction_map, output_path)
""" return status
"""
return [im_path, gt_path, 'Success']
def model_combination(infolders, outfolder):
models = []
for infolder in infolders:
# load each models' images
model = load_segmentations_test(infolder,
prefix='test-segmentation-',
suffix='.nii')
if models:
assert len(models[len(models) - 1]) == len(model), 'model mismatch'
print 'model has {} images\n'.format(len(model))
models.append(model)
if not os.path.exists(outfolder):
os.makedirs(outfolder)
# iterate images
for ind_image in xrange(len(models[0])):
multi_images_d = []
multi_images_f = []
# iterate load models
for ind_model in xrange(len(models)):
model = models[ind_model]
image_f = model[ind_image]
# image_path
image_path = os.path.join(infolders[ind_model], image_f)
print 'Model {} image_path: {}'.format(ind_model, image_path)
# load image data
image_metadata = load_data(image_path)
assert image_metadata is not None, 'image open failed'
image_data = image_metadata['image_data']
multi_images_d.append(image_data)
multi_images_f.append(image_f)
# iterate check models' images
for ind_f in xrange(1, len(multi_images_f)):
assert multi_images_f[0] == multi_images_f[ind_f], 'index mismatch'
assert multi_images_d[0].shape == multi_images_d[
ind_f].shape, 'image shape mismatch'
### iterate combine models
out_image = np.zeros(multi_images_d[0].shape,
dtype=multi_images_d[0].dtype)
for ind_model in xrange(len(multi_images_d)):
image_data = multi_images_d[ind_model]
# print image_data.dtype, np.sum(image_data == 0), np.sum(image_data == 1), np.sum(image_data == 2)
image_data[image_data > 1] = 10
# print image_data.dtype, np.sum(image_data == 0), np.sum(image_data == 1), np.sum(image_data == 10)
out_image += image_data
### process the label
out_image[out_image < 1] = 0
out_image[(out_image >= 1) & (out_image < 10)] = 1
out_image[out_image >= 10] = 2
# print out_image.dtype, np.sum(out_image == 0), np.sum(out_image == 1), np.sum(out_image == 2)
### save merge results
outpath = os.path.join(outfolder, multi_images_f[0])
print 'Output file will save to: {}\n'.format(outpath)
save_data(out_image, outpath)
print '=== DONE ==='
def combine_liver_liver(infolder1, infolder2, outfolder):
livers1 = load_segmentations_test(infolder1,
prefix='test-segmentation-',
suffix='.nii')
livers2 = load_segmentations_test(infolder2,
prefix='test-segmentation-',
suffix='.nii')
#print livers1, livers2
assert len(livers1) == len(
livers2), 'liver1 number must equal to livers2 number'
print 'Total number of livers1: {}\n'.format(len(livers1))
if not os.path.exists(outfolder):
os.makedirs(outfolder)
for ind in range(len(livers1)):
liver1_f = livers1[ind]
liver2_f = livers2[ind]
liver1_index = os.path.splitext(liver1_f)[0].split('-')[2]
liver2_index = os.path.splitext(liver2_f)[0].split('-')[2]
assert liver1_index == liver2_index, 'index mismatch'
liver1_path = os.path.join(infolder1, liver1_f)
liver2_path = os.path.join(infolder2, liver2_f)
print 'liver1_path: {}'.format(liver1_path)
print 'liver2_path: {}'.format(liver2_path)
# load Image
liver1_metadata = load_data(liver1_path)
liver2_metadata = load_data(liver2_path)
assert liver1_metadata is not None, 'liver1 open failed'
assert liver2_metadata is not None, 'liver2 open failed'
liver1_data = liver1_metadata['image_data']
liver2_data = liver2_metadata['image_data']
###
# keep both liver1_data's liver label and liver2_data's liver label
###
print liver1_data.dtype, liver2_data.dtype
print np.sum(liver1_data == 0), np.sum(liver1_data == 1), np.sum(
liver1_data == 2)
print np.sum(liver2_data == 0), np.sum(liver2_data == 1), np.sum(
liver2_data == 2)
liver1_data += liver2_data
liver1_data[liver1_data > 0] = 1
print np.sum(liver1_data == 0), np.sum(liver1_data == 1), np.sum(
liver1_data == 2)
assert np.sum(liver1_data == 0) == np.sum(
liver1_data < 1), 'liver1_data == 0, error'
assert np.sum(liver1_data == 1) == np.sum(
(liver1_data > 0) & (liver1_data < 2)), 'liver1_data == 1, error'
assert np.sum(liver1_data == 2) == np.sum(
liver1_data > 1), 'liver1_data == 2, error'
print liver1_data.dtype
### save merge results
outpath = os.path.join(outfolder, liver1_f)
print 'Output file will save to: {}\n'.format(outpath)
save_data(liver1_data, outpath)
print '=== DONE ==='
def merge_liver_lesion(infolder1, infolder2, outfolder):
###
# keep liver_data's liver label and lesion_data's lesion label
###
livers = load_segmentations_test(infolder1,
prefix='test-segmentation-',
suffix='.nii')
lesions = load_segmentations_test(infolder2,
prefix='test-segmentation-',
suffix='.nii')
#print livers, lesions
assert len(livers) == len(
lesions), 'liver number must equal to lesions number'
print 'Total number of livers: {}\n'.format(len(livers))
if not os.path.exists(outfolder):
os.makedirs(outfolder)
for ind in range(len(livers)):
liver_f = livers[ind]
lesion_f = lesions[ind]
liver_index = os.path.splitext(liver_f)[0].split('-')[2]
lesion_index = os.path.splitext(lesion_f)[0].split('-')[2]
assert liver_index == lesion_index, 'index mismatch'
liver_path = os.path.join(infolder1, liver_f)
lesion_path = os.path.join(infolder2, lesion_f)
print 'liver_path: {}'.format(liver_path)
print 'lesion_path: {}'.format(lesion_path)
# load Image
liver_metadata = load_data(liver_path)
lesion_metadata = load_data(lesion_path)
assert liver_metadata is not None, 'liver open failed'
assert lesion_metadata is not None, 'lesion open failed'
liver_data = liver_metadata['image_data']
lesion_data = lesion_metadata['image_data']
###
# keep liver_data's liver label and lesion_data's lesion label
###
print liver_data.dtype, lesion_data.dtype
# get liver only from liver_data, set others to zeros
liver_data[liver_data > 1] = 0
# get lesion only from lesion_data, set others to zeros
lesion_data[lesion_data < 2] = 0
# move lesion to liver_data, liver will overrie
liver_data += lesion_data
liver_data[liver_data > 1] = 2
# print np.sum(liver_data==0), np.sum(liver_data == 1), np.sum(liver_data==2)
print liver_data.dtype
### save merge results
outpath = os.path.join(outfolder, liver_f)
print 'Output file will save to: {}\n'.format(outpath)
save_data(liver_data, outpath)
#print 'Output saved to: {}\n'.format(foutpath)
print '=== DONE ==='
def refine_liver(infolder1, infolder2, outfolder):
livers1 = load_segmentations_test(infolder1,
prefix='test-segmentation-',
suffix='.nii')
livers2 = load_segmentations_test(infolder2,
prefix='test-segmentation-',
suffix='.nii')
#print livers1, livers2
assert len(livers1) == len(
livers2), 'liver1 number must equal to livers2 number'
print 'Total number of livers1: {}\n'.format(len(livers1))
if not os.path.exists(outfolder):
os.makedirs(outfolder)
for ind in range(len(livers1)):
liver1_f = livers1[ind]
liver2_f = livers2[ind]
liver1_index = os.path.splitext(liver1_f)[0].split('-')[2]
liver2_index = os.path.splitext(liver2_f)[0].split('-')[2]
assert liver1_index == liver2_index, 'index mismatch'
liver1_path = os.path.join(infolder1, liver1_f)
liver2_path = os.path.join(infolder2, liver2_f)
print 'liver1_path: {}'.format(liver1_path)
print 'liver2_path: {}'.format(liver2_path)
# load Image
liver1_metadata = load_data(liver1_path)
liver2_metadata = load_data(liver2_path)
assert liver1_metadata is not None, 'liver1 open failed'
assert liver2_metadata is not None, 'liver2 open failed'
liver1_data = liver1_metadata['image_data']
liver2_data = liver2_metadata['image_data']
###
# keep liver1_data's liver label and liver1_data's lesions label
# and
# add liver2_data's liver label
###
print liver1_data.dtype, liver2_data.dtype
print np.sum(liver1_data == 0), np.sum(liver1_data == 1), np.sum(
liver1_data == 2)
print np.sum(liver2_data == 0), np.sum(liver2_data == 1), np.sum(
liver2_data == 2)
# set liver1_data's lesion to label 10 in order to keep it
# set liver2_data's lesion to label 0, in order to remove it
liver1_data[liver1_data > 1] = 10
liver2_data[liver2_data > 1] = 0
print np.sum(liver1_data == 0), np.sum(liver1_data == 1), np.sum(
liver1_data == 10)
print np.sum(liver2_data == 0), np.sum(liver2_data == 1), np.sum(
liver2_data == 2)
# add liver2_data's liver to liver1_data
liver1_data += liver2_data
# set label
liver1_data[liver1_data <= 0] = 0
liver1_data[(liver1_data > 0) & (liver1_data < 10)] = 1
liver1_data[liver1_data >= 10] = 2
print liver1_data.dtype
print np.sum(liver1_data == 0), np.sum(liver1_data == 1), np.sum(
liver1_data == 2)
### save merge results
outpath = os.path.join(outfolder, liver1_f)
print 'Output file will save to: {}\n'.format(outpath)
save_data(liver1_data, outpath)
print '=== DONE ==='
def model_average(infolders, outfolder):
'''
Label_Path = infolder + 'label'
Prob_Path = infolder + 'prob'
MA_Path = infolder + 'model_average'
'''
models = []
for infolder in infolders:
# load each models' images
model = load_segmentations_test(os.path.join(infolder, 'label'),
prefix='test-segmentation-',
suffix='.nii')
if models:
assert len(models[len(models) - 1]) == len(model), 'model mismatch'
print 'model has {} images\n'.format(len(model))
models.append(model)
if not os.path.exists(outfolder):
os.makedirs(outfolder)
# iterate images
for ind_image in xrange(len(models[0])):
multi_images_d = []
multi_images_f = []
# iterate load models
for ind_model in xrange(len(models)):
model = models[ind_model]
image_f = model[ind_image]
prob_c0_f = '{}-class-0{}'.format(
os.path.splitext(image_f)[0],
os.path.splitext(image_f)[1])
prob_c1_f = '{}-class-1{}'.format(
os.path.splitext(image_f)[0],
os.path.splitext(image_f)[1])
prob_c2_f = '{}-class-2{}'.format(
os.path.splitext(image_f)[0],
os.path.splitext(image_f)[1])
# path
label_path = os.path.join(infolders[ind_model], 'label', image_f)
prob_c0_path = os.path.join(infolders[ind_model], 'prob',
prob_c0_f)
prob_c1_path = os.path.join(infolders[ind_model], 'prob',
prob_c1_f)
prob_c2_path = os.path.join(infolders[ind_model], 'prob',
prob_c2_f)
print 'model {}:\n'.format(ind_model)
print 'label_path: {}\n'.format(label_path)
print 'prob_c0_path: {}\n'.format(prob_c0_path)
print 'prob_c1_path: {}\n'.format(prob_c1_path)
print 'prob_c2_path: {}\n'.format(prob_c2_path)
# load data
prob_c0_metadata = load_data(prob_c0_path)
prob_c1_metadata = load_data(prob_c1_path)
prob_c2_metadata = load_data(prob_c2_path)
assert prob_c0_metadata is not None, 'prob_c0 open failed'
assert prob_c1_metadata is not None, 'prob_c1 open failed'
assert prob_c2_metadata is not None, 'prob_c2 open failed'
prob_c0_data = prob_c0_metadata['image_data']
prob_c1_data = prob_c1_metadata['image_data']
prob_c2_data = prob_c2_metadata['image_data']
prob_data = np.concatenate(
(prob_c0_data[np.newaxis, ...], prob_c1_data[np.newaxis, ...],
prob_c2_data[np.newaxis, ...]),
axis=0)
multi_images_d.append(prob_data)
multi_images_f.append(image_f)
# iterate check models' images
for ind_f in xrange(1, len(multi_images_f)):
assert multi_images_f[0] == multi_images_f[ind_f], 'index mismatch'
assert multi_images_d[0].shape == multi_images_d[
ind_f].shape, 'image shape mismatch'
### iterate add models
out_image = np.zeros(multi_images_d[0].shape,
dtype=multi_images_d[0].dtype)
for ind_model in xrange(len(multi_images_d)):
out_image += multi_images_d[ind_model]
### average models
out_image /= float(len(multi_images_d))
### argmax and transform the datatype to unint8
out_label = np.argmax(out_image, axis=0)
out_label = out_label.astype(np.uint8)
### save out_label
outpath = os.path.join(outfolder, multi_images_f[0])
print out_label.dtype
print 'Output file will save to: {}\n'.format(outpath)
save_data(out_label, outpath)
### process the label
# out_image[out_image<1] = 0
# out_image[(out_image>=1)&(out_image<10)] = 1
# out_image[out_image>=10] = 2
# print out_image.dtype, np.sum(out_image == 0), np.sum(out_image == 1), np.sum(out_image == 2)
print '=== DONE ==='
def test_2d(self, data_unit):
"""Test Thread for 2d segmentation
Assumes images are already prepared
"""
""" Get data from data_unit """
[im_path, gt_path, im, gt] = data_unit
print im_path, gt_path
filename, ext = osp.splitext(osp.split(im_path)[1])
if filename.split('-')[0] == 'test':
'''TEST'''
fileindex = filename.split('-')[2]
filename = 'test-segmentation-{}{}'.format(fileindex, ext)
else:
filename = '{}_pred{}'.format(filename, ext)
""" Do Forward
"""
prediction_map_shape = [
self.params.CLASS_NUM,
]
prediction_map_shape.extend(im.shape)
prediction_map = np.zeros(prediction_map_shape, dtype=im.dtype)
# initial set background prob 1.0
prediction_map[0] = 1.0
if ext in ('.jpg', '.png', '.npy'):
prediction_map[:, :, :, 0] = self.get_output_map(im, gt)
if self.params.DEBUG:
fig, axes = plt.subplots(1, 3)
vis_seg(axes, [
im[:, :, im.shape[2] / 2], gt[:, :, gt.shape[2] / 2],
np.argmax(prediction_map[:, :, :, 0], axis=0)
])
exit()
# set outdir
outdir = osp.join(self.output_dir, filename.split('_slice_')[0])
if not os.path.exists(outdir):
os.makedirs(outdir)
elif ext in ('.nii'):
if self.params.DEBUG:
plt.ion()
if self.params.FAST_INFERENCE:
""" Apply mask to remove irrelavant slices for fast inference """
# [0:row_col_page_b[2]] and row_col_page_e[2]:] no need inference
assert gt is not None, 'gt must be not None'
row_col_page_b, new_size = trim_volume_square(
gt, min_size=[1, 1, 1], pad=[0, 0, 1])
row_col_page_e = row_col_page_b + new_size
ind_begin = row_col_page_b[2]
ind_end = row_col_page_e[2]
else:
ind_begin = 0
ind_end = im.shape[2]
for ind_slice in xrange(ind_begin, ind_end):
### Get input data
if self.params.ADJACENT:
ind_slices_merged = [
max(0, ind_slice - 2),
max(0, ind_slice - 1), ind_slice,
min(im.shape[2] - 1, ind_slice + 1),
min(im.shape[2] - 1, ind_slice + 2)
]
each_im = im[:, :, ind_slices_merged]
each_gt = gt[:, :, ind_slices_merged]
else:
each_im = im[:, :, ind_slice]
each_im = each_im[:, :, np.newaxis]
each_gt = gt[:, :, ind_slice]
each_gt = each_gt[:, :, np.newaxis]
prediction_map[:, :, :, ind_slice] = self.get_output_map(
each_im, each_gt)
if self.params.DEBUG:
print prediction_map.shape, type(
prediction_map), prediction_map.dtype
fig, axes = plt.subplots(2, 3)
for ind_slice in xrange(ind_begin, ind_end):
img_list = [im[:, :, ind_slice], gt[:, :, ind_slice]]
for ind_class in xrange(prediction_map.shape[0]):
img_list.append(prediction_map[ind_class, :, :,
ind_slice])
img_list.append(
np.argmax(prediction_map[:, :, :, ind_slice], axis=0))
vis_seg(axes, img_list)
exit()
# set outdir
outdir = self.output_dir
else:
print 'error'
""" Save """
# construct output_path
pred_label_dir = osp.join(outdir, 'label')
pred_prob_dir = osp.join(outdir, 'prob')
if not os.path.exists(pred_label_dir):
os.makedirs(pred_label_dir)
if not os.path.exists(pred_prob_dir):
os.makedirs(pred_prob_dir)
### save prediction label
pred_label_path = osp.join(pred_label_dir, filename)
# argmax and transform the datatype to unint8
prediction_label = np.argmax(prediction_map, axis=0)
prediction_label = prediction_label.astype(np.uint8)
print('Save label to {}'.format(pred_label_path))
save_data(prediction_label, pred_label_path)
### save prediction prob maps
for ind_class in xrange(prediction_map.shape[0]):
f, ext = osp.splitext(filename)
pred_prob_name = '{}-class-{}{}'.format(
osp.splitext(filename)[0], ind_class,
osp.splitext(filename)[1])
pred_prob_path = osp.join(pred_prob_dir, pred_prob_name)
print('Save class {} to {}'.format(ind_class, pred_prob_path))
save_data(prediction_map[ind_class], pred_prob_path)
if self.params.DEBUG:
print prediction_map.shape, type(
prediction_map), prediction_map.dtype
fig, axes = plt.subplots(2, 3)
plt.ion()
for ind_slice in xrange(prediction_map.shape[-1]):
img_list = [im[:, :, ind_slice], gt[:, :, ind_slice]]
for ind_class in xrange(prediction_map.shape[0]):
img_list.append(prediction_map[ind_class, :, :, ind_slice])
img_list.append(prediction_label[:, :, ind_slice])
vis_seg(axes, img_list)
exit()
""" return status """
return [im_path, gt_path, 'Success']
