def test_model():
input_shape = (256, 256, 3)
local_shape = (128, 128, 3)
batch_size = 4
test_datagen = Data(test_data, input_shape[:2], local_shape[:2])
completion_model = load_model("model/completion.h5", compile=False)
#completion_model.summary()
cnt = 0
for inputs, points, masks in test_datagen.flow(batch_size):
completion_image = completion_model.predict([inputs, masks])
num_img = min(5, batch_size)
fig, axs = plt.subplots(num_img, 3)
for i in range(num_img):
axs[i, 0].imshow(inputs[i] * (1 - masks[i]))
axs[i, 0].axis('off')
axs[i, 0].set_title('Input')
axs[i, 1].imshow(completion_image[i])
axs[i, 1].axis('off')
axs[i, 1].set_title('Output')
axs[i, 2].imshow(inputs[i])
axs[i, 2].axis('off')
axs[i, 2].set_title('Ground Truth')
fig.savefig(os.path.join("result", "result_test_%d.png" % cnt))
plt.close()
cnt += 1
def train_model(result_dir="result", data=data_train):
train_datagen = Data(data_train, input_shape[:2], local_shape[:2])
for n in range(n_epoch):
progbar = generic_utils.Progbar(len(train_datagen))
for inputs, points, masks in train_datagen.flow(batch_size):
completion_image = completion_model.predict([inputs, masks])
valid = np.ones((batch_size, 1))
fake = np.zeros((batch_size, 1))
g_loss = 0.0
d_loss = 0.0
if n < tc:
g_loss = completion_model.train_on_batch([inputs, masks],
inputs)
else:
d_loss_real = d_model.train_on_batch([inputs, points], valid)
d_loss_fake = d_model.train_on_batch(
[completion_image, points], fake)
d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)
if n >= tc + td:
g_loss = all_model.train_on_batch([inputs, masks, points],
[inputs, valid])
g_loss = g_loss[0] + alpha * g_loss[1]
progbar.add(inputs.shape[0],
values=[("D loss", d_loss), ("G mse", g_loss)])
num_img = min(5, batch_size)
fig, axs = plt.subplots(num_img, 3)
for i in range(num_img):
axs[i, 0].imshow(inputs[i] * (1 - masks[i]))
axs[i, 0].axis('off')
axs[i, 0].set_title('Input')
axs[i, 1].imshow(completion_image[i])
axs[i, 1].axis('off')
axs[i, 1].set_title('Output')
axs[i, 2].imshow(inputs[i])
axs[i, 2].axis('off')
axs[i, 2].set_title('Ground Truth')
fig.savefig(os.path.join(result_dir, "result_%d.png" % n))
plt.close()
# Save the checkpoints every 10 epochs
if (n + 1) % 10 == 0:
checkpoint.save(file_prefix=checkpoint_prefix)
# save model
generator.save(os.path.join("model", "generator.h5"))
completion_model.save(os.path.join("model", "completion.h5"))
discriminator.save(os.path.join("model", "discriminator.h5"))
