def main():
args = build_parser().parse_args()
assert args.checkpoint_path
result_dir = args.result_dir
checkpoint_path = args.checkpoint_path
test_dir = args.test_dir
n_imgs = args.n_images
image_size = [args.img_height, args.img_width]
sess = tf.Session()
unet = Unet(input_shape=image_size,
sess=sess,
filter_num=args.filter_num,
batch_norm=args.batch_norm)
unet.build_net(is_train=False)
unet.load_weights(checkpoint_path)
img_names = os.listdir(test_dir)
img_names.sort()
mask_names = None
total_dice = None
if args.mask_dir:
mask_names = os.listdir(args.mask_dir)
mask_names.sort()
total_dice = 0
if n_imgs <= 0:
n_imgs = len(img_names)
for i in range(n_imgs):
print('%s %d/%d' % (img_names[i], i, n_imgs))
img_mat = read_car_img(os.path.join(test_dir, img_names[i]),
image_size=image_size)
img_mat = np.expand_dims(img_mat, axis=0)
if mask_names:
mask_mat = read_mask_img(os.path.join(args.mask_dir,
mask_names[i]),
image_size=image_size)
mask_mat = np.expand_dims(mask_mat, axis=0)
res, dice = unet.predict_test(img_mat, mask_mat)
dice = np.mean(dice)
print('Dice coefficient:%.6f' % dice)
total_dice += dice
else:
res = unet.predict(img_mat)
if args.result_dir:
res = res.reshape(image_size)
misc.imsave(os.path.join(result_dir, img_names[i]), res)
if total_dice:
print('Average Dice coefficient:%.6f' % (total_dice / n_imgs))
log_dir = "log_dir"
log_path = os.path.join(sys.path[0], log_dir)
#show_one_example()
nx, ny = 572, 572
generator = image_gen.GrayScaleDataProvider(nx, ny, cnt=20)
channels = generator.channels
n_class = generator.n_class
net = Unet(channels=channels, n_class=n_class, features_root=16, layers=3)
trainer = Trainer(net, optimizer="momentum", opt_kwargs=dict(momentum=0.2, learning_rate=0.2))
path = trainer.train(generator, training_iters=20, epochs=10, display_step=2, restore=False, write_graph=False)
x_test, y_test = generator(1)
prediction = net.predict(os.path.join(log_path, "model.cpkt"), x_test)
fig, ax = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(12,5))
ax[0].imshow(x_test[0,...,0], aspect="auto")
ax[1].imshow(y_test[0,...,1], aspect="auto")
mask = prediction[0,...,1] > 0.9
ax[2].imshow(mask, aspect="auto")
ax[0].set_title("Input")
ax[1].set_title("Ground truth")
ax[2].set_title("Prediction")
fig.tight_layout()
fig.savefig(os.path.join(sys.path[0], "toy_problem.png"))
