StarGAN.py

from __future__ import print_function, division
import numpy as np
import os
import cv2
from PIL import Image
import random
from functools import partial
from utils import *

import tensorflow as tf
from keras.models import Model, Sequential, load_model
from keras.layers.merge import _Merge
from keras.layers import Input, Conv2D, MaxPooling2D, ZeroPadding2D, Conv2D, BatchNormalization, LeakyReLU, ReLU, UpSampling2D
from keras.layers import Reshape, Dropout, Concatenate, Lambda, Multiply, Add, Flatten, Dense
from keras_contrib.layers.normalization.instancenormalization import InstanceNormalization
from keras.optimizers import Adam
from keras import backend as K

class RandomWeightedAverage(_Merge):
    """Provides a (random) weighted average between real and generated image samples"""
    def define_batch_size(self, bs):
        self.bs = bs
    def _merge_function(self, inputs):
        alpha = K.random_uniform((self.bs, 1, 1, 1))
        return (alpha * inputs[0]) + ((1 - alpha) * inputs[1])

class StarGAN(object):
    def __init__(self, args):

        # Model configuration.
        self.c_dim = args.c_dim
        self.image_size = args.image_size
        self.g_conv_dim = args.g_conv_dim
        self.d_conv_dim = args.d_conv_dim
        self.g_repeat_num = args.g_repeat_num
        self.d_repeat_num = args.d_repeat_num
        self.lambda_cls = args.lambda_cls
        self.lambda_rec = args.lambda_rec
        self.lambda_gp = args.lambda_gp

        # Training configuration.
        self.dataset = args.dataset
        self.batch_size = args.batch_size
        self.num_iters = args.num_iters
        self.num_iters_decay = args.num_iters_decay
        self.g_lr = args.g_lr
        self.d_lr = args.d_lr
        self.n_critic = args.n_critic
        self.beta1 = args.beta1
        self.beta2 = args.beta2
        self.selected_attrs = args.selected_attrs

        # Test configurations.
        self.test_iters = args.test_iters

        # Miscellaneous.
        self.mode = args.mode 

        # Directories.
        self.data_dir = args.data_dir
        self.sample_dir = args.sample_dir
        self.model_save_dir = args.model_save_dir
        self.result_dir = args.result_dir

        # Step size.
        self.log_step = args.log_step
        self.sample_step = args.sample_step
        self.model_save_step = args.model_save_step
        self.lr_update_step = args.lr_update_step

        # Custom image
        self.custom_image_name = args.custom_image_name
        self.custom_image_label = args.custom_image_label

    def ResidualBlock(self, inp, dim_out):
        """Residual Block with instance normalization."""
        x = ZeroPadding2D(padding = 1)(inp)
        x = Conv2D(filters = dim_out, kernel_size=3, strides=1, padding='valid', use_bias = False)(x)
        x = InstanceNormalization(axis = -1)(x)
        x = ReLU()(x)
        x = ZeroPadding2D(padding = 1)(x)
        x = Conv2D(filters = dim_out, kernel_size=3, strides=1, padding='valid', use_bias = False)(x)
        x = InstanceNormalization(axis = -1)(x)
        return Add()([inp, x])

    def build_generator(self):
        """Generator network."""
        # Input tensors
        inp_c = Input(shape = (self.c_dim, ))
        inp_img = Input(shape = (self.image_size, self.image_size, 3))
    
        # Replicate spatially and concatenate domain information
        c = Lambda(lambda x: K.repeat(x, self.image_size**2))(inp_c)
        c = Reshape((self.image_size, self.image_size, self.c_dim))(c)
        x = Concatenate()([inp_img, c])
    
        # First Conv2D
        x = Conv2D(filters = self.g_conv_dim, kernel_size = 7, strides = 1, padding = 'same', use_bias = False)(x)
        x = InstanceNormalization(axis = -1)(x)
        x = ReLU()(x)
    
        # Down-sampling layers
        curr_dim = self.g_conv_dim
        for i in range(2):
            x = ZeroPadding2D(padding = 1)(x)
            x = Conv2D(filters = curr_dim*2, kernel_size = 4, strides = 2, padding = 'valid', use_bias = False)(x)
            x = InstanceNormalization(axis = -1)(x)
            x = ReLU()(x)
            curr_dim = curr_dim * 2
        
        # Bottleneck layers.
        for i in range(self.g_repeat_num):
            x = self.ResidualBlock(x, curr_dim)
        
        # Up-sampling layers
        for i in range(2):
            x = UpSampling2D(size = 2)(x)       
            x = Conv2D(filters = curr_dim // 2, kernel_size = 4, strides = 1, padding = 'same', use_bias = False)(x)
            x = InstanceNormalization(axis = -1)(x)
            x = ReLU()(x)        
            curr_dim = curr_dim // 2
    
        # Last Conv2D
        x = ZeroPadding2D(padding = 3)(x)
        out = Conv2D(filters = 3, kernel_size = 7, strides = 1, padding = 'valid', activation = 'tanh', use_bias = False)(x)
    
        return Model(inputs = [inp_img, inp_c], outputs = out)        

    def build_discriminator(self):
        """Discriminator network with PatchGAN."""
        inp_img = Input(shape = (self.image_size, self.image_size, 3))
        x = ZeroPadding2D(padding = 1)(inp_img)
        x = Conv2D(filters = self.d_conv_dim, kernel_size = 4, strides = 2, padding = 'valid', use_bias = False)(x)
        x = LeakyReLU(0.01)(x)
    
        curr_dim = self.d_conv_dim
        for i in range(1, self.d_repeat_num):
            x = ZeroPadding2D(padding = 1)(x)
            x = Conv2D(filters = curr_dim*2, kernel_size = 4, strides = 2, padding = 'valid')(x)
            x = LeakyReLU(0.01)(x)
            curr_dim = curr_dim * 2
    
        kernel_size = int(self.image_size / np.power(2, self.d_repeat_num))
    
        out_src = ZeroPadding2D(padding = 1)(x)
        out_src = Conv2D(filters = 1, kernel_size = 3, strides = 1, padding = 'valid', use_bias = False)(out_src)
    
        out_cls = Conv2D(filters = self.c_dim, kernel_size = kernel_size, strides = 1, padding = 'valid', use_bias = False)(x)
        out_cls = Reshape((self.c_dim, ))(out_cls)
    
        return Model(inp_img, [out_src, out_cls])

    def classification_loss(self, Y_true, Y_pred) :
        return tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=Y_true, logits=Y_pred))

    def wasserstein_loss(self, Y_true, Y_pred):
        return K.mean(Y_true*Y_pred)

    def reconstruction_loss(self, Y_true, Y_pred):
        return K.mean(K.abs(Y_true - Y_pred))

    def gradient_penalty_loss(self, y_true, y_pred, averaged_samples):
        """
        Computes gradient penalty based on prediction and weighted real / fake samples
        """
        gradients = K.gradients(y_pred, averaged_samples)[0]
        # compute the euclidean norm by squaring ...
        gradients_sqr = K.square(gradients)
        #   ... summing over the rows ...
        gradients_sqr_sum = K.sum(gradients_sqr, axis=np.arange(1, len(gradients_sqr.shape)))
        #   ... and sqrt
        gradient_l2_norm = K.sqrt(gradients_sqr_sum)
        # compute lambda * (1 - ||grad||)^2 still for each single sample
        gradient_penalty = K.square(1 - gradient_l2_norm)
        # return the mean as loss over all the batch samples
        return K.mean(gradient_penalty)

    def build_model(self):
        self.G = self.build_generator()
        self.D = self.build_discriminator()

        # First don't update weights of Generator block
        self.G.trainable = False

        # Compute output with real images.
        x_real = Input(shape = (self.image_size, self.image_size, 3))
        out_src_real, out_cls_real = self.D(x_real)

        # Compute output with fake images.
        label_trg = Input(shape = (self.c_dim, ))
        x_fake = self.G([x_real, label_trg])
        out_src_fake, out_cls_fake = self.D(x_fake)

        # Compute output for gradient penalty.
        rd_avg = RandomWeightedAverage()
        rd_avg.define_batch_size(self.batch_size)
        x_hat = rd_avg([x_real, x_fake])
        out_src, _ = self.D(x_hat)
            
        # Use Python partial to provide loss function with additional 'averaged_samples' argument
        partial_gp_loss = partial(self.gradient_penalty_loss, averaged_samples=x_hat)
        partial_gp_loss.__name__ = 'gradient_penalty' # Keras requires function names

        # Define training model D
        self.train_D = Model([x_real, label_trg], [out_src_real, out_cls_real, out_src_fake, out_src])

        # Setup loss for train_D
        self.train_D.compile(loss = [self.wasserstein_loss, self.classification_loss, self.wasserstein_loss, partial_gp_loss], 
                             optimizer=Adam(lr = self.d_lr, beta_1 = self.beta1, beta_2 = self.beta2), loss_weights = [1, self.lambda_cls, 1, self.lambda_gp])

        # Update G and not update D
        self.G.trainable = True
        self.D.trainable = False

        # All inputs
        real_x = Input(shape = (self.image_size, self.image_size, 3))
        org_label = Input(shape = (self.c_dim, ))
        trg_label = Input(shape = (self.c_dim, ))

        # Compute output of fake image
        fake_x = self.G([real_x, trg_label])
        fake_out_src, fake_out_cls = self.D(fake_x)

        # Target-to-original domain.
        x_reconst = self.G([fake_x, org_label])

        # Define traning model G
        self.train_G = Model([real_x, org_label, trg_label], [fake_out_src, fake_out_cls, x_reconst])

        # Setup loss for train_G
        self.train_G.compile(loss = [self.wasserstein_loss, self.classification_loss, self.reconstruction_loss], 
                             optimizer=Adam(lr = self.g_lr, beta_1 = self.beta1, beta_2 = self.beta2), loss_weights = [1, self.lambda_cls, self.lambda_rec])
        
        """ Input Image"""
        self.Image_data_class = ImageData(data_dir=self.data_dir, selected_attrs=self.selected_attrs)
        self.Image_data_class.preprocess()

    def train(self):
        data_iter = get_loader(self.Image_data_class.train_dataset, self.Image_data_class.train_dataset_label, self.Image_data_class.train_dataset_fix_label, 
                               image_size=self.image_size, batch_size=self.batch_size, mode=self.mode)

        # Training
        valid = -np.ones((self.batch_size, 2, 2, 1))
        fake =  np.ones((self.batch_size, 2, 2, 1))
        dummy = np.zeros((self.batch_size, 2, 2, 1)) # Dummy gt for gradient penalty
        for epoch in range(self.num_iters):
            imgs, orig_labels, target_labels, fix_labels, _ = next(data_iter)
    
            # Setting learning rate (linear decay)
            if epoch > (self.num_iters - self.num_iters_decay):
                K.set_value(self.train_D.optimizer.lr, self.d_lr*(self.num_iters - epoch)/(self.num_iters - self.num_iters_decay))
                K.set_value(self.train_G.optimizer.lr, self.g_lr*(self.num_iters - epoch)/(self.num_iters - self.num_iters_decay))
    
            # Training Discriminators        
            D_loss = self.train_D.train_on_batch(x = [imgs, target_labels], y = [valid, orig_labels, fake, dummy])
        
            # Training Generators
            if (epoch + 1) % self.n_critic == 0:
                G_loss = self.train_G.train_on_batch(x = [imgs, orig_labels, target_labels], y = [valid, target_labels, imgs])
        
            if (epoch + 1) % self.log_step == 0:
                print(f"Iteration: [{epoch + 1}/{self.num_iters}]")
                print(f"\tD/loss_real = [{D_loss[1]:.4f}], D/loss_fake = [{D_loss[3]:.4f}], D/loss_cls =  [{D_loss[2]:.4f}], D/loss_gp = [{D_loss[4]:.4f}]")
                print(f"\tG/loss_fake = [{G_loss[1]:.4f}], G/loss_rec = [{G_loss[3]:.4f}], G/loss_cls = [{G_loss[2]:.4f}]") 

            if (epoch + 1) % self.model_save_step == 0:  
                self.G.save_weights(os.path.join(self.model_save_dir, 'G_weights.hdf5'))
                self.D.save_weights(os.path.join(self.model_save_dir, 'D_weights.hdf5'))
                self.train_D.save_weights(os.path.join(self.model_save_dir, 'train_D_weights.hdf5'))
                self.train_G.save_weights(os.path.join(self.model_save_dir, 'train_G_weights.hdf5'))    

    def test(self):
        G_weights_dir = os.path.join(self.model_save_dir, 'G_weights.hdf5')
        if not os.path.isfile(G_weights_dir):
            print("Don't find weight's generator model")
        else:
            self.G.load_weights(G_weights_dir)

        data_iter = get_loader(self.Image_data_class.test_dataset, self.Image_data_class.test_dataset_label, self.Image_data_class.test_dataset_fix_label, 
                               image_size=self.image_size, batch_size=self.batch_size, mode=self.mode)        
        n_batches = int(len(self.sample_step) / self.batch_size)
        total_samples = n_batches * self.batch_size

        for i in range(n_batches):
            imgs, orig_labels, target_labels, fix_labels, names = next(data_iter)
            for j in range(self.batch_size):
                preds = self.G.predict([np.repeat(np.expand_dims(imgs[j], axis = 0), len(self.selected_attrs), axis = 0), fix_labels[j]])
                for k in range(len(self.selected_attrs)):                    
                    Image.fromarray((preds[k]*127.5 + 127.5).astype(np.uint8)).save(os.path.join(self.result_dir, names[j].split(os.path.sep)[-1].split('.')[0] + f'_{k + 1}.png'))

    def custom(self):
        G_weights_dir = os.path.join(self.model_save_dir, 'G_weights.hdf5')
        if not os.path.isfile(G_weights_dir):
            print("Don't find weight's generator model")
        else:
            self.G.load_weights(G_weights_dir)                        

        path = os.path.join(self.sample_dir, self.custom_image_name)
        target_list = create_labels([self.custom_image_label], selected_attrs=self.selected_attrs)[0]
        image = cv2.imread(path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = resize_keep_aspect_ratio(image, width = self.image_size, height = self.image_size)
        image = np.array([image])/127.5 - 1
        preds = self.G.predict([np.repeat(image, len(self.selected_attrs), axis = 0), target_list])
        for k in range(len(self.selected_attrs)):                    
            Image.fromarray((preds[k]*127.5 + 127.5).astype(np.uint8)).save(os.path.join(self.sample_dir, self.custom_image_name.split('.')[0] + f'_{k + 1}.png'))