infer.py

#copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
#
#Licensed under the Apache License, Version 2.0 (the "License");
#you may not use this file except in compliance with the License.
#You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#Unless required by applicable law or agreed to in writing, software
#distributed under the License is distributed on an "AS IS" BASIS,
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#See the License for the specific language governing permissions and
#limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import functools
import os
from PIL import Image
import paddle.fluid as fluid
import paddle
import numpy as np
import imageio
import glob
from util.config import add_arguments, print_arguments
from data_reader import celeba_reader_creator
import copy

parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
# yapf: disable
add_arg('model_net',         str,   'cgan',            "The model used")
add_arg('net_G',             str,   "resnet_9block",   "Choose the CycleGAN and Pix2pix generator's network, choose in [resnet_9block|resnet_6block|unet_128|unet_256]")
add_arg('init_model',        str,   None,              "The init model file of directory.")
add_arg('output',            str,   "./infer_result",  "The directory the infer result to be saved to.")
add_arg('input_style',       str,   "A",               "The style of the input, A or B")
add_arg('norm_type',         str,   "batch_norm",      "Which normalization to used")
add_arg('use_gpu',           bool,  True,              "Whether to use GPU to train.")
add_arg('dropout',           bool,  False,             "Whether to use dropout")
add_arg('g_base_dims',       int,   64,                "Base channels in CycleGAN generator")
add_arg('c_dim',             int,   13,                "the size of attrs")
add_arg('use_gru',           bool,  False,             "Whether to use GRU")
add_arg('crop_size',         int,   178,               "crop size")
add_arg('image_size',        int,   128,               "image size")
add_arg('selected_attrs',    str,
    "Bald,Bangs,Black_Hair,Blond_Hair,Brown_Hair,Bushy_Eyebrows,Eyeglasses,Male,Mouth_Slightly_Open,Mustache,No_Beard,Pale_Skin,Young",
"the attributes we selected to change")
add_arg('batch_size',        int,   16,                "batch size when test")
add_arg('test_list',       str,   "./data/celeba/test_list_attr_celeba.txt",                "the test list file")
add_arg('dataset_dir',       str,   "./data/celeba/",                "the dataset directory to be infered")
add_arg('n_layers',        int,     5,      "default layers in generotor")
add_arg('gru_n_layers',    int,     4,       "default layers of GRU in generotor")
# yapf: enable


def infer(args):
    data_shape = [-1, 3, args.image_size, args.image_size]
    input = fluid.layers.data(name='input', shape=data_shape, dtype='float32')
    label_org_ = fluid.layers.data(
        name='label_org_', shape=[args.c_dim], dtype='float32')
    label_trg_ = fluid.layers.data(
        name='label_trg_', shape=[args.c_dim], dtype='float32')

    model_name = 'net_G'
    if args.model_net == 'CycleGAN':
        from network.CycleGAN_network import CycleGAN_model
        model = CycleGAN_model()
        if args.input_style == "A":
            fake = model.network_G(input, name="GA", cfg=args)
        elif args.input_style == "B":
            fake = model.network_G(input, name="GB", cfg=args)
        else:
            raise "Input with style [%s] is not supported." % args.input_style
    elif args.model_net == 'Pix2pix':
        from network.Pix2pix_network import Pix2pix_model
        model = Pix2pix_model()
        fake = model.network_G(input, "generator", cfg=args)
    elif args.model_net == 'StarGAN':
        from network.StarGAN_network import StarGAN_model
        model = StarGAN_model()
        fake = model.network_G(input, label_trg_, name="g_main", cfg=args)
    elif args.model_net == 'STGAN':
        from network.STGAN_network import STGAN_model
        model = STGAN_model()
        fake, _ = model.network_G(
            input, label_org_, label_trg_, cfg=args, name='net_G')
    elif args.model_net == 'AttGAN':
        from network.AttGAN_network import AttGAN_model
        model = AttGAN_model()
        fake, _ = model.network_G(
            input, label_org_, label_trg_, cfg=args, name='net_G')
    else:
        raise NotImplementedError("model_net {} is not support".format(
            args.model_net))

    # prepare environment
    place = fluid.CPUPlace()
    if args.use_gpu:
        place = fluid.CUDAPlace(0)
    exe = fluid.Executor(place)
    exe.run(fluid.default_startup_program())
    for var in fluid.default_main_program().global_block().all_parameters():
        print(var.name)
    print(args.init_model + '/' + model_name)
    fluid.io.load_persistables(exe, args.init_model + "/" + model_name)
    print('load params done')

    if not os.path.exists(args.output):
        os.makedirs(args.output)

    if args.model_net == 'AttGAN' or args.model_net == 'STGAN':
        test_reader = celeba_reader_creator(
            image_dir=args.dataset_dir,
            list_filename=args.test_list,
            batch_size=args.batch_size,
            drop_last=False,
            args=args)
        reader_test = test_reader.get_test_reader(
            args, shuffle=False, return_name=True)
        for data in zip(reader_test()):
            real_img, label_org, name = data[0]
            print("read {}".format(name))
            label_trg = copy.deepcopy(label_org)
            tensor_img = fluid.LoDTensor()
            tensor_label_org = fluid.LoDTensor()
            tensor_label_trg = fluid.LoDTensor()
            tensor_label_org_ = fluid.LoDTensor()
            tensor_label_trg_ = fluid.LoDTensor()
            tensor_img.set(real_img, place)
            tensor_label_org.set(label_org, place)
            real_img_temp = np.squeeze(real_img).transpose([0, 2, 3, 1])
            images = [real_img_temp]
            for i in range(args.c_dim):
                label_trg_tmp = copy.deepcopy(label_trg)
                for j in range(len(label_org)):
                    label_trg_tmp[j][i] = 1.0 - label_trg_tmp[j][i]
                label_trg_ = list(
                    map(lambda x: ((x * 2) - 1) * 0.5, label_trg_tmp))
                for j in range(len(label_org)):
                    label_trg_[j][i] = label_trg_[j][i] * 2.0
                tensor_label_org_.set(label_org, place)
                tensor_label_trg.set(label_trg, place)
                tensor_label_trg_.set(label_trg_, place)
                out = exe.run(feed={
                    "input": tensor_img,
                    "label_org_": tensor_label_org_,
                    "label_trg_": tensor_label_trg_
                },
                              fetch_list=[fake.name])
                fake_temp = np.squeeze(out[0]).transpose([0, 2, 3, 1])
                images.append(fake_temp)
            images_concat = np.concatenate(images, 1)
            images_concat = np.concatenate(images_concat, 1)
            imageio.imwrite(args.output + "/fake_img_" + name[0], (
                (images_concat + 1) * 127.5).astype(np.uint8))
    elif args.model_net == 'StarGAN':
        test_reader = celeba_reader_creator(
            image_dir=args.dataset_dir,
            list_filename=args.test_list,
            batch_size=args.batch_size,
            drop_last=False,
            args=args)
        reader_test = test_reader.get_test_reader(
            args, shuffle=False, return_name=True)
        for data in zip(reader_test()):
            real_img, label_org, name = data[0]
            print("read {}".format(name))
            tensor_img = fluid.LoDTensor()
            tensor_label_org = fluid.LoDTensor()
            tensor_img.set(real_img, place)
            tensor_label_org.set(label_org, place)
            real_img_temp = np.squeeze(real_img).transpose([0, 2, 3, 1])
            images = [real_img_temp]
            for i in range(args.c_dim):
                label_trg = np.zeros(
                    [len(label_org), args.c_dim]).astype("float32")
                for j in range(len(label_org)):
                    label_trg[j][i] = 1
                tensor_label_trg = fluid.LoDTensor()
                tensor_label_trg.set(label_trg, place)
                out = exe.run(
                    feed={"input": tensor_img,
                          "label_trg_": tensor_label_trg},
                    fetch_list=[fake.name])
                fake_temp = np.squeeze(out[0]).transpose([0, 2, 3, 1])
                images.append(fake_temp)
            images_concat = np.concatenate(images, 1)
            images_concat = np.concatenate(images_concat, 1)
            imageio.imwrite(args.output + "/fake_img_" + name[0], (
                (images_concat + 1) * 127.5).astype(np.uint8))

    elif args.model_net == 'Pix2pix' or args.model_net == 'CycleGAN':
        for file in glob.glob(args.dataset_dir):
            print("read {}".format(file))
            image_name = os.path.basename(file)
            image = Image.open(file).convert('RGB')
            image = image.resize((256, 256), Image.BICUBIC)
            image = np.array(image).transpose([2, 0, 1]).astype('float32')
            image = image / 255.0
            image = (image - 0.5) / 0.5
            data = image[np.newaxis, :]
            tensor = fluid.LoDTensor()
            tensor.set(data, place)

            fake_temp = exe.run(fetch_list=[fake.name], feed={"input": tensor})
            fake_temp = np.squeeze(fake_temp[0]).transpose([1, 2, 0])
            input_temp = np.squeeze(data).transpose([1, 2, 0])

            imageio.imwrite(args.output + "/fake_" + image_name, (
                (fake_temp + 1) * 127.5).astype(np.uint8))
    else:
        raise NotImplementedError("model_net {} is not support".format(
            args.model_net))


if __name__ == "__main__":
    args = parser.parse_args()
    print_arguments(args)
    infer(args)