def evaluate_generator(model_path):
'''
Evaluate performance of CT slices using generator
:model_path: path of evaluated model
'''
model = UNet(input_shape=(512, 512, 1))
model.compile(optimizer=Adam(lr=1e-4),
loss=focal_tversky_loss,
metrics=[dice_coef])
model.load_weights(model_path)
image_path = 'generate_data/test.h5' # h5 file of evaluated patients
TOTAL_IMAGE = 1814 # total images in h5 file
BATCH_SIZE = 2
reader = HDF5DatasetGenerator(db_path=image_path, batch_size=BATCH_SIZE)
generator = reader.generator()
results = model.evaluate_generator(generator,
steps=TOTAL_IMAGE // BATCH_SIZE,
verbose=1)
for i, metric in enumerate(model.metrics_names):
print(metric, ':', results[i])
reader.close()
def predict(model_path, idx):
'''
Predict segmentation mask of single CT slice
:model_path: path of evaluated model
:idx: index of evaluated patient
'''
pred_dir = 'pred/'
if not os.path.exists(pred_dir):
os.mkdir(pred_dir)
model = UNet((512, 512, 1))
model.compile(optimizer=Adam(lr=1e-4),
loss=focal_tversky_loss,
metrics=[dice_coef])
model.load_weights(model_path)
# reading original images
reader = sitk.ImageSeriesReader()
dicom_names = reader.GetGDCMSeriesFileNames(data_path + 'volume/' +
str(idx) + '/')
reader.SetFileNames(dicom_names)
images = reader.Execute()
images = sitk.GetArrayFromImage(images) # images.shape: (slices, 512, 512)
# HU values are restricted to lung window settings [-1100, 300]
# namely: window width: 1400, window level: -400
images[images < -1100] = -1100
images[images > 300] = 300
# normalization
max_value, min_value = images.max(), images.min()
images = images / (max_value - min_value)
pred_inputs = np.expand_dims(
images, axis=-1) # pred_inputs.shape: (slices, 512, 512, 1)
pred_masks = model.predict(
pred_inputs, batch_size=1,
verbose=1) # pred_masks.shape: (slices, 512, 512, 1)
# values in pred_masks are probabilities in range [0, 1]
# converting to binary mask using a threshold of 0.5
pred_masks[pred_masks >= 0.5] = 1
pred_masks[pred_masks < 0.5] = 0
pred_masks = pred_masks.astype(int)
out = sitk.GetImageFromArray(pred_masks)
sitk.WriteImage(out, pred_dir + 'pred_threshold_' + str(idx) + '.nii.gz')
def train():
train_reader = HDF5DatasetGenerator(db_path=train_path,
batch_size=BATCH_SIZE)
train_iter = train_reader.generator()
val_reader = HDF5DatasetGenerator(db_path=val_path, batch_size=BATCH_SIZE)
val_iter = val_reader.generator()
model = UNet(input_shape=(512, 512, 1))
model.compile(optimizer=Adam(lr=1e-4),
loss=focal_tversky_loss,
metrics=[dice_coef])
if not os.path.exists(save_path):
os.mkdir(save_path)
os.mkdir(save_path + '/model')
os.mkdir(save_path + '/model/logs')
model_checkpoint = ModelCheckpoint(
save_path + '/model/weights.{epoch:02d}-{val_loss:.2f}.hdf5',
monitor='val_loss',
verbose=1,
save_best_only=True)
tensorboard = TensorBoard(log_dir=save_path + '/model/logs')
reduce_lr = ReduceLROnPlateau(monitor='val_loss',
factor=0.1,
patience=4,
mode='auto')
callbacks = [model_checkpoint, tensorboard, reduce_lr]
model.fit_generator(train_iter,
steps_per_epoch=TOTAL_TRAIN // BATCH_SIZE,
epochs=20,
validation_data=val_iter,
validation_steps=TOTAL_VAL // BATCH_SIZE,
verbose=1,
callbacks=callbacks)
train_reader.close()
val_reader.close()
model.save(save_path + '/model/model.h5')
print('Finished training ......')
def evaluate(model_path, idx):
'''
Evaluate performance of single CT slice
:model_path: path of evaluated model
:idx: index of evaluated patient
'''
model = UNet(input_shape=(512, 512, 1))
model.compile(optimizer=Adam(lr=1e-4),
loss=focal_tversky_loss,
metrics=[dice_coef])
model.load_weights(model_path)
# reading original images
reader = sitk.ImageSeriesReader()
dicom_names = reader.GetGDCMSeriesFileNames(data_path + 'volume/' +
str(idx) + '/')
reader.SetFileNames(dicom_names)
images = reader.Execute()
images = sitk.GetArrayFromImage(images) # images.shape: (slices, 512, 512)
# HU values are restricted to lung window settings [-1100, 300]
# namely: window width: 1400, window level: -400
images[images < -1100] = -1100
images[images > 300] = 300
# normalization
max_value, min_value = images.max(), images.min()
images = images / (max_value - min_value)
# reading ground truth labels
labels = sitk.ReadImage(data_path + 'segmentation/' + str(idx) + '.nii')
labels = sitk.GetArrayFromImage(
labels) # eval_labels.shape: (slices, 512, 512)
# (slices, 512, 512) => (slices, 512, 512, 1)
images = np.expand_dims(images, axis=-1)
labels = np.expand_dims(labels, axis=-1)
results = model.evaluate(images, labels, batch_size=1)
for i, metric in enumerate(model.metrics_names):
print(metric, ':', results[i])