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')
if command == "--Train":
if architecture == '--UNet':
UNet = UNet()
UNet.train()
elif architecture == '--UNet++':
UNetPP = UNetPP()
UNetPP.train()
else:
raise Exception(
"You have passed an invalid argument.\nCheck the documentation for the allowed arguments."
)
elif command == "--Predict":
if architecture == '--UNet':
UNet = UNet()
UNet.predict()
elif architecture == '--UNet++':
UNetPP = UNetPP()
UNetPP.predict()
else:
raise Exception(
"You have passed an invalid argument.Check the documentation for the allowed arguments."
)
else:
raise Exception(
"You have passed an invalid argument. Check the documentation for the allowed arguments."
)
elif command == "--Summary":
architecture = sys.argv[2]