def main():
place = set_device(FLAGS.device)
fluid.enable_dygraph(place) if FLAGS.dynamic else None
# Generators
g_AB = Generator()
g_BA = Generator()
g = GeneratorCombine(g_AB, g_BA, is_train=False)
im_shape = [-1, 3, 256, 256]
input_A = Input(im_shape, 'float32', 'input_A')
input_B = Input(im_shape, 'float32', 'input_B')
g.prepare(inputs=[input_A, input_B], device=FLAGS.device)
g.load(FLAGS.init_model, skip_mismatch=True, reset_optimizer=True)
if not os.path.exists(FLAGS.output):
os.makedirs(FLAGS.output)
test_data_A = data.TestDataA()
test_data_B = data.TestDataB()
for i in range(len(test_data_A)):
data_A, A_name = test_data_A[i]
data_B, B_name = test_data_B[i]
data_A = np.array(data_A).astype("float32")
data_B = np.array(data_B).astype("float32")
fake_A, fake_B, cyc_A, cyc_B = g.test_batch([data_A, data_B])
datas = [fake_A, fake_B, cyc_A, cyc_B, data_A, data_B]
odatas = []
for o in datas:
d = np.squeeze(o[0]).transpose([1, 2, 0])
im = ((d + 1) * 127.5).astype(np.uint8)
odatas.append(im)
imsave(FLAGS.output + "/fakeA_" + B_name, odatas[0])
imsave(FLAGS.output + "/fakeB_" + A_name, odatas[1])
imsave(FLAGS.output + "/cycA_" + A_name, odatas[2])
imsave(FLAGS.output + "/cycB_" + B_name, odatas[3])
imsave(FLAGS.output + "/inputA_" + A_name, odatas[4])
imsave(FLAGS.output + "/inputB_" + B_name, odatas[5])
def main():
place = set_device(FLAGS.device)
fluid.enable_dygraph(place) if FLAGS.dynamic else None
# Generators
g_AB = Generator()
g_BA = Generator()
g = GeneratorCombine(g_AB, g_BA, is_train=False)
im_shape = [-1, 3, 256, 256]
input_A = Input(im_shape, 'float32', 'input_A')
input_B = Input(im_shape, 'float32', 'input_B')
g.prepare(inputs=[input_A, input_B])
g.load(FLAGS.init_model, skip_mismatch=True, reset_optimizer=True)
out_path = FLAGS.output + "/single"
if not os.path.exists(out_path):
os.makedirs(out_path)
for f in glob.glob(FLAGS.input):
image_name = os.path.basename(f)
image = Image.open(f).convert('RGB')
image = image.resize((256, 256), Image.BICUBIC)
image = np.array(image) / 127.5 - 1
image = image[:, :, 0:3].astype("float32")
data = image.transpose([2, 0, 1])[np.newaxis, :]
if FLAGS.input_style == "A":
_, fake, _, _ = g.test([data, data])
if FLAGS.input_style == "B":
fake, _, _, _ = g.test([data, data])
fake = np.squeeze(fake[0]).transpose([1, 2, 0])
opath = "{}/fake{}{}".format(out_path, FLAGS.input_style, image_name)
imsave(opath, ((fake + 1) * 127.5).astype(np.uint8))
print("transfer {} to {}".format(f, opath))
def main():
place = set_device(FLAGS.device)
fluid.enable_dygraph(place) if FLAGS.dynamic else None
# Generators
g_AB = Generator()
g_BA = Generator()
# Discriminators
d_A = Discriminator()
d_B = Discriminator()
g = GeneratorCombine(g_AB, g_BA, d_A, d_B)
da_params = d_A.parameters()
db_params = d_B.parameters()
g_params = g_AB.parameters() + g_BA.parameters()
da_optimizer = opt(da_params)
db_optimizer = opt(db_params)
g_optimizer = opt(g_params)
im_shape = [None, 3, 256, 256]
input_A = Input(im_shape, 'float32', 'input_A')
input_B = Input(im_shape, 'float32', 'input_B')
fake_A = Input(im_shape, 'float32', 'fake_A')
fake_B = Input(im_shape, 'float32', 'fake_B')
g_AB.prepare(inputs=[input_A], device=FLAGS.device)
g_BA.prepare(inputs=[input_B], device=FLAGS.device)
g.prepare(g_optimizer,
GLoss(),
inputs=[input_A, input_B],
device=FLAGS.device)
d_A.prepare(da_optimizer,
DLoss(),
inputs=[input_B, fake_B],
device=FLAGS.device)
d_B.prepare(db_optimizer,
DLoss(),
inputs=[input_A, fake_A],
device=FLAGS.device)
if FLAGS.resume:
g.load(FLAGS.resume)
loader_A = paddle.io.DataLoader(data.DataA(),
places=place,
shuffle=True,
return_list=True,
batch_size=FLAGS.batch_size)
loader_B = paddle.io.DataLoader(data.DataB(),
places=place,
shuffle=True,
return_list=True,
batch_size=FLAGS.batch_size)
A_pool = data.ImagePool()
B_pool = data.ImagePool()
for epoch in range(FLAGS.epoch):
for i, (data_A, data_B) in enumerate(zip(loader_A, loader_B)):
data_A = data_A[0][0] if not FLAGS.dynamic else data_A[0]
data_B = data_B[0][0] if not FLAGS.dynamic else data_B[0]
start = time.time()
fake_B = g_AB.test_batch(data_A)[0]
fake_A = g_BA.test_batch(data_B)[0]
g_loss = g.train_batch([data_A, data_B])[0]
fake_pb = B_pool.get(fake_B)
da_loss = d_A.train_batch([data_B, fake_pb])[0]
fake_pa = A_pool.get(fake_A)
db_loss = d_B.train_batch([data_A, fake_pa])[0]
t = time.time() - start
if i % 20 == 0:
print("epoch: {} | step: {:3d} | g_loss: {:.4f} | " \
"da_loss: {:.4f} | db_loss: {:.4f} | s/step {:.4f}".
format(epoch, i, g_loss[0], da_loss[0], db_loss[0], t))
g.save('{}/{}'.format(FLAGS.checkpoint_path, epoch))