Exemplo n.º 1
0
Arquivo: test.py Projeto: zdqf/hapi
def main():
    place = paddle.set_device(FLAGS.device)
    paddle.disable_static(place) if FLAGS.dynamic else None

    im_shape = [-1, 3, 256, 256]
    input_A = Input(im_shape, 'float32', 'input_A')
    input_B = Input(im_shape, 'float32', 'input_B')

    # Generators
    g_AB = Generator()
    g_BA = Generator()
    g = paddle.Model(
        GeneratorCombine(
            g_AB, g_BA, is_train=False),
        inputs=[input_A, input_B])

    g.prepare()
    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])
Exemplo n.º 2
0
def main():
    place = paddle.set_device(FLAGS.device)
    paddle.disable_static(place) if FLAGS.dynamic else None

    im_shape = [-1, 3, 256, 256]
    input_A = Input(im_shape, 'float32', 'input_A')
    input_B = Input(im_shape, 'float32', 'input_B')

    # Generators
    g_AB = Generator()
    g_BA = Generator()

    g = paddle.Model(GeneratorCombine(g_AB, g_BA, is_train=False),
                     inputs=[input_A, input_B])
    g.prepare()

    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_batch([data, data])

        if FLAGS.input_style == "B":
            fake, _, _, _ = g.test_batch([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))
Exemplo n.º 3
0
Arquivo: train.py Projeto: wzzju/hapi
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))