BEGAN.py

import os
import glob
import math
import random
import datetime
import pickle

import numpy as np
import matplotlib.pyplot as plt

from tensorflow.python import keras
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.models import Model,  Sequential
from tensorflow.python.keras.layers import *
from tensorflow.python.keras.preprocessing.image import load_img, img_to_array, array_to_img, ImageDataGenerator
from tensorflow.python.keras.optimizers import *
from tensorflow.python.keras.losses import *


DATA_DIR = ""
BATCH_SIZE = 16
IMG_SHAPE = (64, 64, 3)

data_gen = ImageDataGenerator(rescale=1/255.)
train_data_generator = data_gen.flow_from_directory(directory=DATA_DIR, classes=['faces'], class_mode=None, 
                                                    batch_size=BATCH_SIZE, target_size=IMG_SHAPE[:2])
                                                    
                                                    
def build_encoder(input_shape, z_size, n_filters, n_layers):
    inputs = Input(shape = input_shape)
    x = Conv2D(filters=n_filters, kernel_size=(3, 3), strides=(1, 1), activation='elu', padding='same')(inputs)
    x = Conv2D(filters=n_filters, kernel_size=(3, 3), strides=(1, 1), padding='same')(x)
    
    for i in range(2, n_layers + 1):
        x = Conv2D(filters=n_filters*i, kernel_size=(3, 3), strides=(1, 1), activation='elu', padding='same')(x)
        x = Conv2D(filters=n_filters*i, kernel_size=(3, 3), strides=(2, 2), padding='same')(x)
        
    x = Conv2D(filters=n_filters*n_layers, kernel_size=(3, 3), padding='same')(x)
    x = Flatten()(x)
    outputs = Dense(z_size)(x)

    model = Model(inputs, outputs)    
    
    return model
    
    
def build_decoder(output_shape, z_size, n_filters, n_layers):
    # upsampling param
    scale = 2**(n_layers - 1)
    fc_shape = (output_shape[0]//scale, output_shape[1]//scale, n_filters)
    fc_size = fc_shape[0]*fc_shape[1]*fc_shape[2]
    
    inputs = Input(shape=(z_size, ))
    x = Dense(fc_size)(inputs)
    x = Reshape(fc_shape)(x)
    for i in range(n_layers - 1):
        x = Conv2D(filters=n_filters, kernel_size=(3, 3), strides=(1, 1), activation='elu', padding='same')(x)
        x = Conv2D(filters=n_filters, kernel_size=(3, 3), strides=(1, 1), activation='elu', padding='same')(x)
        x = UpSampling2D()(x)
    
    x = Conv2D(filters=n_filters, kernel_size=(3, 3), strides=(1, 1), activation='elu', padding='same')(x)
    x = Conv2D(filters=n_filters, kernel_size=(3, 3), strides=(1, 1), activation='elu', padding='same')(x)
    outputs = Conv2D(filters=3, kernel_size=(3, 3), strides=(1, 1), padding='same')(x)

    model = Model(inputs, outputs)    
    
    return model
    
    
def build_generator(img_shape, z_size, n_filters, n_layers):
    decoder = build_decoder(img_shape, z_size, n_filters, n_layers)
    return decoder


def build_discriminator(img_shape, z_size, n_filters, n_layers):
    encoder = build_encoder(img_shape, z_size, n_filters, n_layers)
    decoder = build_decoder(img_shape, z_size, n_filters, n_layers)
    merge = decoder(encoder.output)
    return Model(inputs=encoder.input, outputs=merge)    


def build_discriminator_trainer(discriminator):
    img_shape = discriminator.input_shape[1:]
    real_inputs = Input(img_shape)
    fake_inputs = Input(img_shape)
    real_outputs = discriminator(real_inputs)
    fake_outputs = discriminator(fake_inputs)
    
    return Model(inputs=[real_inputs, fake_inputs], outputs=[real_outputs, fake_outputs])
    
    
n_filters = 64
n_layers = 4
z_size = 32

generator = build_generator(IMG_SHAPE, z_size, n_filters, n_layers)
discriminator = build_discriminator(IMG_SHAPE, z_size, n_filters, n_layers)
discriminator_trainer = build_discriminator_trainer(discriminator)

generator.summary()
discriminator.summary()


def build_generator_loss(discriminator):
    def loss(y_true, y_pred):
        # y_ture is dummy
        reconst = discriminator(y_pred)
        return mean_absolute_error(reconst, y_pred)
    return loss
    
    
g_lr = 0.0001

generator_loss = build_generator_loss(discriminator)
generator.compile(loss=generator_loss, optimizer=Adam(g_lr))


d_lr = 0.0001

k_var = 0.0
k = K.variable(k_var)
discriminator_trainer.compile(loss=[mean_absolute_error, mean_absolute_error], loss_weights=[1., -k], optimizer=Adam(d_lr))


def measure(real_loss, fake_loss, gamma):
    return real_loss + np.abs(gamma*real_loss - fake_loss)
    
    
GAMMA = 0.5
LR_K = 0.001
TOTAL_STEPS = 100000
MODEL_SAVE_DIR = 'began_s/models'
IMG_SAVE_DIR = 'began_s/imgs'
IMG_SAMPLE_SHAPE = (5, 5)
N_IMG_SAMPLES = np.prod(IMG_SAMPLE_SHAPE)


os.makedirs(MODEL_SAVE_DIR, exist_ok=True)
os.makedirs(IMG_SAVE_DIR, exist_ok=True)


def save_imgs(path, imgs, rows, cols):
    """画像をタイル状にならべて保存する
    
    Arguments:
        path (str): 保存先のファイルパス
        imgs (np.array): 保存する画像のリスト
        rows (int): タイルの縦のサイズ
        cols (int): タイルの横のサイズ
        
    Returns:
        None
    """
    base_width = imgs.shape[1]
    base_height = imgs.shape[2]
    channels = imgs.shape[3]
    output_shape = (
        base_height*rows,
        base_width*cols,
        channels
    )
    buffer = np.zeros(output_shape)
    for row in range(rows):
        for col in range(cols):
            img = imgs[row*cols + col]
            buffer[
                row*base_height:(row + 1)*base_height,
                col*base_width:(col + 1)*base_width
            ] = img
    array_to_img(buffer).save(path)
    
    
sample_seeds = np.random.uniform(-1, 1, (N_IMG_SAMPLES, z_size))

history = []
logs = []

for step, batch in enumerate(train_data_generator):
    if len(batch) < BATCH_SIZE:
        continue
        
    if step > TOTAL_STEPS:
        break
        
    z_g = np.random.uniform(-1, 1, (BATCH_SIZE, z_size))
    z_d = np.random.uniform(-1, 1, (BATCH_SIZE, z_size))
    
    g_pred = generator.predict(z_d)
    generator.train_on_batch(z_g, batch)
    _, real_loss, fake_loss = discriminator_trainer.train_on_batch([batch, g_pred], [batch, g_pred])
    
    k_var += LR_K*(GAMMA*real_loss - fake_loss)
    K.set_value(k, k_var)
    
    history.append({'real_loss': real_loss,
                    'fake_loss': fake_loss})
    
    if step%1000 == 0:
        measurement = np.mean([
            measure(loss['real_loss'], loss['fake_loss'], GAMMA) for loss in history[-1000:]
        ])
        
        logs.append({
            'k': K.get_value(k),
            'measure': measurement,
            'real_liss': real_loss,
            'fake_loss': fake_loss
        })
        print(logs[-1])
        
        img_path = '{}/generated_{}.png'.format(IMG_SAVE_DIR, step)
        save_imgs(img_path, generator.predict(sample_seeds), rows=IMG_SAMPLE_SHAPE[0], cols=IMG_SAMPLE_SHAPE[1])
        generator.save('{}/generator_{}.hd5'.format(MODEL_SAVE_DIR, step))
        discriminator.save('{}/discriminator_{}.hd5'.format(MODEL_SAVE_DIR, step))
        
        
array_to_img(generator.predict(np.random.uniform(-1, 1, (BATCH_SIZE, z_size)))[0])