def segment_patient(image_file_in, model_weights, liver_file_out,
lesion_file_out=None, latent_file_out=None,
mean=48, batch=1, input_cols=8):
'''
From an input image:
1. Apply preprocessing
2. Segment the liver
3. Segment the lesion
'''
# Input size as defined by the pretrained network from the paper
input_size = 512
if not os.path.exists(image_file_in):
raise KeyError(('Image {} does not exist!').format(image_file_in))
# Read mage and convert to array
image = sitk.ReadImage(image_file_in)
image = sitk.GetArrayFromImage(image)
image = np.transpose(image, [2, 1, 0])
# Apply preprocessing
image = preprocessing(image, input_size)
image -= mean
# Load the H-DenseUnet model
model, latent_model = dense_rnn_net(batch, input_size, input_cols)
model.load_weights(model_weights)
if latent_file_out is not None:
latent_model.load_weights(model_weights, by_name=True)
else:
latent_model = None
sgd = SGD(lr=1e-2, momentum=0.9, nesterov=True)
model.compile(optimizer=sgd, loss=[weighted_crossentropy])
# Predict the tumor segmentation
print('Predicting masks on test data...' + str(id))
if latent_file_out is None:
liver, lesions = predict_liver_and_tumor(model, image, batch=batch,
input_size=input_size,
input_cols=input_cols)
else:
liver, lesions, latent_features =\
predict_liver_and_tumor(model, image, batch=batch,
input_size=input_size,
input_cols=input_cols,
latent_model=latent_model)
# Save the output
sitk.WriteImage(sitk.GetImageFromArray(liver), liver_file_out)
if lesion_file_out is not None:
sitk.WriteImage(sitk.GetImageFromArray(lesions), lesion_file_out)
if latent_file_out is not None:
sitk.WriteImage(sitk.GetImageFromArray(latent_features), latent_file_out)
def predict(args):
if not Path(args.save_path).exists():
os.mkdir(args.save_path)
for id in range(1): #for id in range(70):
print('-' * 30)
print('Loading model and preprocessing test data...' + str(id))
print('-' * 30)
model = dense_rnn_net(args)
model.load_weights(args.model_weight)
sgd = SGD(lr=1e-2, momentum=0.9, nesterov=True)
model.compile(optimizer=sgd, loss=[weighted_crossentropy])
# load data
img_test, img_test_header = load(args.data + str(id) + '.nii')
img_test -= args.mean
# load liver mask
mask, mask_header = load(args.liver_path + str(id) + '-ori.nii')
mask[mask == 2] = 1
mask = ndimage.binary_dilation(mask, iterations=1).astype(mask.dtype)
index = np.where(mask == 1)
mini = np.min(index, axis=-1)
maxi = np.max(index, axis=-1)
print('-' * 30)
print('Predicting masks on test data...' + str(id))
print('-' * 30)
score1, score2 = predict_tumor_inwindow(model, img_test, 3, mini, maxi,
args)
K.clear_session()
result1 = score1
result2 = score2
result1[result1 >= args.thres_liver] = 1
result1[result1 < args.thres_liver] = 0
result2[result2 >= args.thres_tumor] = 1
result2[result2 < args.thres_tumor] = 0
result1[result2 == 1] = 1
print('-' * 30)
print('Postprocessing on mask ...' + str(id))
print('-' * 30)
# preserve the largest liver
Segmask = result2
box = []
[liver_res, num] = measure.label(result1, return_num=True)
region = measure.regionprops(liver_res)
for i in range(num):
box.append(region[i].area)
label_num = box.index(max(box)) + 1
liver_res[liver_res != label_num] = 0
liver_res[liver_res == label_num] = 1
# preserve the largest liver
mask = ndimage.binary_dilation(mask, iterations=1).astype(mask.dtype)
box = []
[liver_labels, num] = measure.label(mask, return_num=True)
region = measure.regionprops(liver_labels)
for i in range(num):
box.append(region[i].area)
label_num = box.index(max(box)) + 1
liver_labels[liver_labels != label_num] = 0
liver_labels[liver_labels == label_num] = 1
liver_labels = ndimage.binary_fill_holes(liver_labels).astype(int)
# preserve tumor within ' largest liver' only
Segmask = Segmask * liver_labels
Segmask = ndimage.binary_fill_holes(Segmask).astype(int)
Segmask = np.array(Segmask, dtype='uint8')
liver_res = np.array(liver_res, dtype='uint8')
liver_res = ndimage.binary_fill_holes(liver_res).astype(int)
liver_res[Segmask == 1] = 2
liver_res = np.array(liver_res, dtype='uint8')
save(liver_res,
args.save_path + 'test-segmentation-' + str(id) + '.nii',
img_test_header)
del Segmask, liver_labels, mask, region, label_num, liver_res
def train_and_predict(args):
print('-' * 30)
print('Creating and compiling model...')
print('-' * 30)
if args.arch == "3dpart":
model = denseunet_3d(args)
model_path = "/3dpart_model"
sgd = SGD(lr=1e-3, momentum=0.9, nesterov=True)
model.compile(optimizer=sgd, loss=[weighted_crossentropy])
model.load_weights(args.model_weight, by_name=True)
else:
model = dense_rnn_net(args)
model_path = "/hybrid_model"
sgd = SGD(lr=1e-3, momentum=0.9, nesterov=True)
model.compile(optimizer=sgd, loss=[weighted_crossentropy])
model.load_weights(args.model_weight)
# liver tumor LITS
trainidx = list(range(2))
img_list = []
tumor_list = []
minindex_list = []
maxindex_list = []
tumorlines = []
tumoridx = []
liveridx = []
liverlines = []
for idx in range(2):
img, img_header = load(args.data + '/myTrainingData/volume-' +
str(idx) + '.nii')
tumor, tumor_header = load(args.data +
'/myTrainingData/segmentation-' + str(idx) +
'.nii')
img_list.append(img)
tumor_list.append(tumor)
maxmin = np.loadtxt(args.data + '/myTrainingDataTxt/LiverBox/box_' +
str(idx) + '.txt',
delimiter=' ')
minindex = maxmin[0:3]
maxindex = maxmin[3:6]
minindex = np.array(minindex, dtype='int')
maxindex = np.array(maxindex, dtype='int')
minindex[0] = max(minindex[0] - 3, 0)
minindex[1] = max(minindex[1] - 3, 0)
minindex[2] = max(minindex[2] - 3, 0)
maxindex[0] = min(img.shape[0], maxindex[0] + 3)
maxindex[1] = min(img.shape[1], maxindex[1] + 3)
maxindex[2] = min(img.shape[2], maxindex[2] + 3)
minindex_list.append(minindex)
maxindex_list.append(maxindex)
f1 = open(
args.data + '/myTrainingDataTxt/TumorPixels/tumor_' + str(idx) +
'.txt', 'r')
tumorline = f1.readlines()
tumorlines.append(tumorline)
tumoridx.append(len(tumorline))
f1.close()
f2 = open(
args.data + '/myTrainingDataTxt/LiverPixels/liver_' + str(idx) +
'.txt', 'r')
liverline = f2.readlines()
liverlines.append(liverline)
liveridx.append(len(liverline))
f2.close()
if not os.path.exists(args.save_path + model_path):
os.mkdir(args.save_path + model_path)
if not os.path.exists(args.save_path + "/history"):
os.mkdir(args.save_path + '/history')
else:
if os.path.exists(args.save_path + "/history/lossbatch.txt"):
os.remove(args.save_path + '/history/lossbatch.txt')
if os.path.exists(args.save_path + "/history/lossepoch.txt"):
os.remove(args.save_path + '/history/lossepoch.txt')
model_checkpoint = ModelCheckpoint(args.save_path + model_path +
'/weights.{epoch:02d}-{loss:.2f}.hdf5',
monitor='loss',
verbose=1,
save_best_only=False,
save_weights_only=False,
mode='min',
period=1)
print('-' * 30)
print('Fitting model......')
print('-' * 30)
steps = 27386 / (args.b * 6)
model.fit_generator(generate_arrays_from_file(args.b, trainidx, img_list,
tumor_list, tumorlines,
liverlines, tumoridx,
liveridx, minindex_list,
maxindex_list),
steps_per_epoch=steps,
epochs=1,
verbose=1,
callbacks=[model_checkpoint],
max_queue_size=10,
workers=1,
use_multiprocessing=True)
print('Finised Training .......')