<reponame>SebastianPartarrieu/ThroughArtistsEyes
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

from os import listdir

from numpy import asarray
from numpy import vstack
from numpy import savez_compressed
from numpy import load

from keras.preprocessing.image import img_to_array
from keras.preprocessing.image import load_img

from keras.layers import Conv2D, Dense, LeakyReLU, Conv2DTranspose, Activation, Concatenate
from keras.models import Model, Input
from keras.initializers import RandomNormal
from keras.optimizers import Adam
from keras_contrib.layers.normalization.instancenormalization import InstanceNormalization

#from tensorflow.keras import backend as K

# config = tf.compat.v1.ConfigProto(intra_op_parallelism_threads=8,
#                         inter_op_parallelism_threads=8, 
#                         allow_soft_placement=True,
#                         device_count = {'CPU' : 1,
#                                         'GPU' : 1}
#                        )

# session = tf.compat.v1.Session(config=config)
# K.set_session(session)

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
LOADED = True
N_EPOCHS = 5

if not LOADED:
	# load all images in a directory into memory
	def load_images(path, size=(256,256)):
		data_list = list()
		# enumerate filenames in directory, assume all are images
		for filename in listdir(path):
			# load and resize the image
			pixels = load_img(path + filename, target_size=size)
			# convert to numpy array
			pixels = img_to_array(pixels)
			# store
			data_list.append(pixels)
		return asarray(data_list)

	#dataset path
	LOCAL_PATH = 'C:/Users/VR/Documents/MOVIE_local/datasets/monet_2_360/'   
	N_EPOCHS = 5

	#load dataset A - 360° photos
	################################################### BEWARE OF SIZE REMEMBER TO CHANGE SIZE IN LOAD IMAGES FUNCTION
	data_train_A = load_images(LOCAL_PATH + 'trainA/')
	data_test_A = load_images(LOCAL_PATH + 'testA/')
	dataA = vstack((data_train_A, data_test_A))
	print(f'Loaded data set A : {dataA.shape}')

	#load dataset B - monet photos
	data_train_B = load_images(LOCAL_PATH + 'trainB/')
	data_test_B = load_images(LOCAL_PATH + 'testB/')
	dataB = vstack((data_train_B, data_test_B))
	print(f'Loaded data set B : {dataB.shape}')

	#save as compressed numpy array
	filename = 'monet_2_360.npz'
	savez_compressed(filename, dataA, dataB)
	print('Saved dataset: ', filename)

	#Let's get plotting
	data = load('monet_2_360.npz')
	dataA, dataB = data['arr_0'], data['arr_1']
	print(f'Loaded from compressed NumPy format: dataset A {dataA.shape} dataset B {dataB.shape}') 

	#optional plotting - WE COULD ADD ARGUMENT PARSER FOR THIS 
	n_samples = 3
	# SHOW pictures
	for i in range(n_samples):
		plt.subplot(2, n_samples, i + 1)
		plt.axis('off')
		plt.imshow(dataA[i].astype('uint8'))
	plt.show()

	#SHOW monet paintings (or other chosen artists)
	for i in range(n_samples):
		plt.subplot(2, n_samples, i+1)
		plt.axis('off')
		plt.imshow(dataB[i].astype('uint8'))
	plt.show()

	#Starting to actually build model from scratch


#################### 70x70 PatchGAN discriminator model, useful as can take different dimension sizes as input
# define the discriminator model
def define_discriminator(image_shape):
	# weight initialization
	init = RandomNormal(stddev=0.02)
	# source image input
	in_image = Input(shape=image_shape)
	# C64
	d = Conv2D(64, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(in_image)
	d = LeakyReLU(alpha=0.2)(d)
	# C128
	d = Conv2D(128, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d)
	d = InstanceNormalization(axis=-1)(d)
	d = LeakyReLU(alpha=0.2)(d)
	# C256
	d = Conv2D(256, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d)
	d = InstanceNormalization(axis=-1)(d)
	d = LeakyReLU(alpha=0.2)(d)
	# C512
	d = Conv2D(512, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d)
	d = InstanceNormalization(axis=-1)(d)
	d = LeakyReLU(alpha=0.2)(d)
	# second last output layer
	d = Conv2D(512, (4,4), padding='same', kernel_initializer=init)(d)
	d = InstanceNormalization(axis=-1)(d)
	d = LeakyReLU(alpha=0.2)(d)
	# patch output
	patch_out = Conv2D(1, (4,4), padding='same', kernel_initializer=init)(d)
	# define model
	model = Model(in_image, patch_out)
	# compile model
	model.compile(loss='mse', optimizer=Adam(lr=0.0002, beta_1=0.5), loss_weights=[0.5])
	return model


################### Generator models
#Encoder-decoder architecture. Downsampling, encoding down to bottleneck then upsampling and creating output image

def resnet_block(n_filters, input_layer):
	# weight initialization
	init = RandomNormal(stddev=0.02)
	# first layer convolutional layer
	g = Conv2D(n_filters, (3,3), padding='same', kernel_initializer=init)(input_layer)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# second convolutional layer
	g = Conv2D(n_filters, (3,3), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	# concatenate merge channel-wise with input layer
	g = Concatenate()([g, input_layer])
	return g

# define the standalone generator model
def define_generator(image_shape, n_resnet=9):
	# weight initialization
	init = RandomNormal(stddev=0.02)
	# image input
	in_image = Input(shape=image_shape)
	# c7s1-64
	g = Conv2D(64, (7,7), padding='same', kernel_initializer=init)(in_image)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# d128
	g = Conv2D(128, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# d256
	g = Conv2D(256, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# R256
	for _ in range(n_resnet):
		g = resnet_block(256, g)
	# u128
	g = Conv2DTranspose(128, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# u64
	g = Conv2DTranspose(64, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# c7s1-3
	g = Conv2D(3, (7,7), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	out_image = Activation('tanh')(g)
	# define model
	model = Model(in_image, out_image)
	return model

def define_composite_model(g_model_1, d_model, g_model_2, image_shape):
	# ensure the model we're updating is trainable
	g_model_1.trainable = True
	# mark discriminator as not trainable
	d_model.trainable = False
	# mark other generator model as not trainable
	g_model_2.trainable = False
	# discriminator element
	input_gen = Input(shape=image_shape)
	gen1_out = g_model_1(input_gen)
	output_d = d_model(gen1_out)
	# identity element
	input_id = Input(shape=image_shape)
	output_id = g_model_1(input_id)
	# forward cycle
	output_f = g_model_2(gen1_out)
	# backward cycle
	gen2_out = g_model_2(input_id)
	output_b = g_model_1(gen2_out)
	# define model graph
	model = Model([input_gen, input_id], [output_d, output_id, output_f, output_b])
	# define optimization algorithm configuration
	opt = Adam(lr=0.0002, beta_1=0.5)
	# compile model with weighting of least squares loss and L1 loss
	model.compile(loss=['mse', 'mae', 'mae', 'mae'], loss_weights=[1, 5, 10, 10], optimizer=opt)
	return model
	
# load and prepare training images
def load_real_samples(filename):
	# load the dataset
	data = load(filename)
	# unpack arrays
	X1, X2 = data['arr_0'], data['arr_1']
	# scale from [0,255] to [-1,1]
	X1 = (X1 - 127.5) / 127.5
	X2 = (X2 - 127.5) / 127.5
	return [X1, X2]

# select a batch of random samples, returns images and target
def generate_real_samples(dataset, n_samples, patch_shape):
	# choose random instances
	ix = np.random.randint(0, dataset.shape[0], n_samples)
	# retrieve selected images
	X = dataset[ix]
	# generate 'real' class labels (1)
	y = np.ones((n_samples, patch_shape, patch_shape, 1))
	return X, y

# generate a batch of images, returns images and targets
def generate_fake_samples(g_model, dataset, patch_shape):
	# generate fake instance
	X = g_model.predict(dataset)
	# create 'fake' class labels (0)
	y = np.zeros((len(X), patch_shape, patch_shape, 1))
	return X, y    


# save the generator models to file - need h5py installed to run this 
def save_models(step, g_model_AtoB, g_model_BtoA):
	# save the first generator model
	filename1 = 'g_model_AtoB_%06d.h5' % (step+1)
	g_model_AtoB.save(filename1)
	# save the second generator model
	filename2 = 'g_model_BtoA_%06d.h5' % (step+1)
	g_model_BtoA.save(filename2)
	print('>Saved: %s and %s' % (filename1, filename2))

	
# generate samples and save as a plot and save the model
def summarize_performance(step, g_model, trainX, name, n_samples=5):
	# select a sample of input images
	X_in, _ = generate_real_samples(trainX, n_samples, 0)
	# generate translated images
	X_out, _ = generate_fake_samples(g_model, X_in, 0)
	# scale all pixels from [-1,1] to [0,1]
	X_in = (X_in + 1) / 2.0
	X_out = (X_out + 1) / 2.0
	# plot real images
	for i in range(n_samples):
		plt.subplot(2, n_samples, 1 + i)
		plt.axis('off')
		plt.imshow(X_in[i])
	# plot translated image
	for i in range(n_samples):
		plt.subplot(2, n_samples, 1 + n_samples + i)
		plt.axis('off')
		plt.imshow(X_out[i])
	# save plot to file
	filename1 = '%s_generated_plot_%06d.png' % (name, (step+1))
	plt.savefig(filename1)
	plt.close()

# update image pool for fake images
def update_image_pool(pool, images, max_size=50):
	selected = list()
	for image in images:
		if len(pool) < max_size:
			# stock the pool
			pool.append(image)
			selected.append(image)
		elif np.random() < 0.5:
			# use image, but don't add it to the pool
			selected.append(image)
		else:
			# replace an existing image and use replaced image
			ix = randint(0, len(pool))
			selected.append(pool[ix])
			pool[ix] = image
	return asarray(selected)

#train cyclegan models
def train(d_model_A, d_model_B, g_model_AtoB, g_model_BtoA, c_model_AtoB, c_model_BtoA, dataset, n_epochs):
	# define properties of the training run
	n_batch = 1
	# determine the output square shape of the discriminator
	n_patch = d_model_A.output_shape[1]
	# unpack dataset
	trainA, trainB = dataset
	# prepare image pool for fakes
	poolA, poolB = list(), list()
	# calculate the number of batches per training epoch
	bat_per_epo = int(len(trainA) / n_batch)
	# calculate the number of training iterations
	n_steps = bat_per_epo * n_epochs
	# manually enumerate epochs
	for i in range(n_steps):
		# select a batch of real samples
		X_realA, y_realA = generate_real_samples(trainA, n_batch, n_patch)
		X_realB, y_realB = generate_real_samples(trainB, n_batch, n_patch)
		# generate a batch of fake samples
		X_fakeA, y_fakeA = generate_fake_samples(g_model_BtoA, X_realB, n_patch)
		X_fakeB, y_fakeB = generate_fake_samples(g_model_AtoB, X_realA, n_patch)
		# update fakes from pool
		X_fakeA = update_image_pool(poolA, X_fakeA)
		X_fakeB = update_image_pool(poolB, X_fakeB)
		# update generator B->A via adversarial and cycle loss
		g_loss2, _, _, _, _  = c_model_BtoA.train_on_batch([X_realB, X_realA], [y_realA, X_realA, X_realB, X_realA])
		# update discriminator for A -> [real/fake]
		dA_loss1 = d_model_A.train_on_batch(X_realA, y_realA)
		dA_loss2 = d_model_A.train_on_batch(X_fakeA, y_fakeA)
		# update generator A->B via adversarial and cycle loss
		g_loss1, _, _, _, _ = c_model_AtoB.train_on_batch([X_realA, X_realB], [y_realB, X_realB, X_realA, X_realB])
		# update discriminator for B -> [real/fake]
		dB_loss1 = d_model_B.train_on_batch(X_realB, y_realB)
		dB_loss2 = d_model_B.train_on_batch(X_fakeB, y_fakeB)
		# summarize performance
		print('>%d, dA[%.3f,%.3f] dB[%.3f,%.3f] g[%.3f,%.3f]' % (i+1, dA_loss1,dA_loss2, dB_loss1,dB_loss2, g_loss1,g_loss2))
		# evaluate the model performance every so often
		if (i+1) % (bat_per_epo * 1) == 0:
			# plot A->B translation
			summarize_performance(i, g_model_AtoB, trainA, 'AtoB')
			# plot B->A translation
			summarize_performance(i, g_model_BtoA, trainB, 'BtoA')
		if (i+1) % (bat_per_epo * 5) == 0:
			# save the models
			save_models(i, g_model_AtoB, g_model_BtoA)



########################## Making the whole thing run 

# with tf.device('/GPU:0'):
dataset = load_real_samples('monet_2_360.npz')
print('Loaded', dataset[0].shape, dataset[1].shape)
# define input shape based on the loaded dataset
image_shape = dataset[0].shape[1:]
# generator: A -> B
g_model_AtoB = define_generator(image_shape)
# generator: B -> A
g_model_BtoA = define_generator(image_shape)
# discriminator: A -> [real/fake]
d_model_A = define_discriminator(image_shape)
# discriminator: B -> [real/fake]
d_model_B = define_discriminator(image_shape)
# composite: A -> B -> [real/fake, A]
c_model_AtoB = define_composite_model(g_model_AtoB, d_model_B, g_model_BtoA, image_shape)
# composite: B -> A -> [real/fake, B]
c_model_BtoA = define_composite_model(g_model_BtoA, d_model_A, g_model_AtoB, image_shape)
# train models
train(d_model_A, d_model_B, g_model_AtoB, g_model_BtoA, c_model_AtoB, c_model_BtoA, dataset, N_EPOCHS)
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

from os import listdir

from numpy import asarray
from numpy import vstack
from numpy import savez_compressed
from numpy import load

from keras.preprocessing.image import img_to_array
from keras.preprocessing.image import load_img

from keras.layers import Conv2D, Dense, LeakyReLU, Conv2DTranspose, Activation, Concatenate
from keras.models import Model, Input
from keras.initializers import RandomNormal
from keras.optimizers import Adam
from keras_contrib.layers.normalization.instancenormalization import InstanceNormalization

#from tensorflow.keras import backend as K

# config = tf.compat.v1.ConfigProto(intra_op_parallelism_threads=8,
#                         inter_op_parallelism_threads=8, 
#                         allow_soft_placement=True,
#                         device_count = {'CPU' : 1,
#                                         'GPU' : 1}
#                        )

# session = tf.compat.v1.Session(config=config)
# K.set_session(session)

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
LOADED = True
N_EPOCHS = 5

if not LOADED:
	# load all images in a directory into memory
	def load_images(path, size=(256,256)):
		data_list = list()
		# enumerate filenames in directory, assume all are images
		for filename in listdir(path):
			# load and resize the image
			pixels = load_img(path + filename, target_size=size)
			# convert to numpy array
			pixels = img_to_array(pixels)
			# store
			data_list.append(pixels)
		return asarray(data_list)

	#dataset path
	LOCAL_PATH = 'C:/Users/VR/Documents/MOVIE_local/datasets/monet_2_360/'   
	N_EPOCHS = 5

	#load dataset A - 360° photos
	################################################### BEWARE OF SIZE REMEMBER TO CHANGE SIZE IN LOAD IMAGES FUNCTION
	data_train_A = load_images(LOCAL_PATH + 'trainA/')
	data_test_A = load_images(LOCAL_PATH + 'testA/')
	dataA = vstack((data_train_A, data_test_A))
	print(f'Loaded data set A : {dataA.shape}')

	#load dataset B - monet photos
	data_train_B = load_images(LOCAL_PATH + 'trainB/')
	data_test_B = load_images(LOCAL_PATH + 'testB/')
	dataB = vstack((data_train_B, data_test_B))
	print(f'Loaded data set B : {dataB.shape}')

	#save as compressed numpy array
	filename = 'monet_2_360.npz'
	savez_compressed(filename, dataA, dataB)
	print('Saved dataset: ', filename)

	#Let's get plotting
	data = load('monet_2_360.npz')
	dataA, dataB = data['arr_0'], data['arr_1']
	print(f'Loaded from compressed NumPy format: dataset A {dataA.shape} dataset B {dataB.shape}') 

	#optional plotting - WE COULD ADD ARGUMENT PARSER FOR THIS 
	n_samples = 3
	# SHOW pictures
	for i in range(n_samples):
		plt.subplot(2, n_samples, i + 1)
		plt.axis('off')
		plt.imshow(dataA[i].astype('uint8'))
	plt.show()

	#SHOW monet paintings (or other chosen artists)
	for i in range(n_samples):
		plt.subplot(2, n_samples, i+1)
		plt.axis('off')
		plt.imshow(dataB[i].astype('uint8'))
	plt.show()

	#Starting to actually build model from scratch


#################### 70x70 PatchGAN discriminator model, useful as can take different dimension sizes as input
# define the discriminator model
def define_discriminator(image_shape):
	# weight initialization
	init = RandomNormal(stddev=0.02)
	# source image input
	in_image = Input(shape=image_shape)
	# C64
	d = Conv2D(64, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(in_image)
	d = LeakyReLU(alpha=0.2)(d)
	# C128
	d = Conv2D(128, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d)
	d = InstanceNormalization(axis=-1)(d)
	d = LeakyReLU(alpha=0.2)(d)
	# C256
	d = Conv2D(256, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d)
	d = InstanceNormalization(axis=-1)(d)
	d = LeakyReLU(alpha=0.2)(d)
	# C512
	d = Conv2D(512, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d)
	d = InstanceNormalization(axis=-1)(d)
	d = LeakyReLU(alpha=0.2)(d)
	# second last output layer
	d = Conv2D(512, (4,4), padding='same', kernel_initializer=init)(d)
	d = InstanceNormalization(axis=-1)(d)
	d = LeakyReLU(alpha=0.2)(d)
	# patch output
	patch_out = Conv2D(1, (4,4), padding='same', kernel_initializer=init)(d)
	# define model
	model = Model(in_image, patch_out)
	# compile model
	model.compile(loss='mse', optimizer=Adam(lr=0.0002, beta_1=0.5), loss_weights=[0.5])
	return model


################### Generator models
#Encoder-decoder architecture. Downsampling, encoding down to bottleneck then upsampling and creating output image

def resnet_block(n_filters, input_layer):
	# weight initialization
	init = RandomNormal(stddev=0.02)
	# first layer convolutional layer
	g = Conv2D(n_filters, (3,3), padding='same', kernel_initializer=init)(input_layer)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# second convolutional layer
	g = Conv2D(n_filters, (3,3), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	# concatenate merge channel-wise with input layer
	g = Concatenate()([g, input_layer])
	return g

# define the standalone generator model
def define_generator(image_shape, n_resnet=9):
	# weight initialization
	init = RandomNormal(stddev=0.02)
	# image input
	in_image = Input(shape=image_shape)
	# c7s1-64
	g = Conv2D(64, (7,7), padding='same', kernel_initializer=init)(in_image)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# d128
	g = Conv2D(128, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# d256
	g = Conv2D(256, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# R256
	for _ in range(n_resnet):
		g = resnet_block(256, g)
	# u128
	g = Conv2DTranspose(128, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# u64
	g = Conv2DTranspose(64, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	g = Activation('relu')(g)
	# c7s1-3
	g = Conv2D(3, (7,7), padding='same', kernel_initializer=init)(g)
	g = InstanceNormalization(axis=-1)(g)
	out_image = Activation('tanh')(g)
	# define model
	model = Model(in_image, out_image)
	return model

def define_composite_model(g_model_1, d_model, g_model_2, image_shape):
	# ensure the model we're updating is trainable
	g_model_1.trainable = True
	# mark discriminator as not trainable
	d_model.trainable = False
	# mark other generator model as not trainable
	g_model_2.trainable = False
	# discriminator element
	input_gen = Input(shape=image_shape)
	gen1_out = g_model_1(input_gen)
	output_d = d_model(gen1_out)
	# identity element
	input_id = Input(shape=image_shape)
	output_id = g_model_1(input_id)
	# forward cycle
	output_f = g_model_2(gen1_out)
	# backward cycle
	gen2_out = g_model_2(input_id)
	output_b = g_model_1(gen2_out)
	# define model graph
	model = Model([input_gen, input_id], [output_d, output_id, output_f, output_b])
	# define optimization algorithm configuration
	opt = Adam(lr=0.0002, beta_1=0.5)
	# compile model with weighting of least squares loss and L1 loss
	model.compile(loss=['mse', 'mae', 'mae', 'mae'], loss_weights=[1, 5, 10, 10], optimizer=opt)
	return model
	
# load and prepare training images
def load_real_samples(filename):
	# load the dataset
	data = load(filename)
	# unpack arrays
	X1, X2 = data['arr_0'], data['arr_1']
	# scale from [0,255] to [-1,1]
	X1 = (X1 - 127.5) / 127.5
	X2 = (X2 - 127.5) / 127.5
	return [X1, X2]

# select a batch of random samples, returns images and target
def generate_real_samples(dataset, n_samples, patch_shape):
	# choose random instances
	ix = np.random.randint(0, dataset.shape[0], n_samples)
	# retrieve selected images
	X = dataset[ix]
	# generate 'real' class labels (1)
	y = np.ones((n_samples, patch_shape, patch_shape, 1))
	return X, y

# generate a batch of images, returns images and targets
def generate_fake_samples(g_model, dataset, patch_shape):
	# generate fake instance
	X = g_model.predict(dataset)
	# create 'fake' class labels (0)
	y = np.zeros((len(X), patch_shape, patch_shape, 1))
	return X, y    


# save the generator models to file - need h5py installed to run this 
def save_models(step, g_model_AtoB, g_model_BtoA):
	# save the first generator model
	filename1 = 'g_model_AtoB_%06d.h5' % (step+1)
	g_model_AtoB.save(filename1)
	# save the second generator model
	filename2 = 'g_model_BtoA_%06d.h5' % (step+1)
	g_model_BtoA.save(filename2)
	print('>Saved: %s and %s' % (filename1, filename2))

	
# generate samples and save as a plot and save the model
def summarize_performance(step, g_model, trainX, name, n_samples=5):
	# select a sample of input images
	X_in, _ = generate_real_samples(trainX, n_samples, 0)
	# generate translated images
	X_out, _ = generate_fake_samples(g_model, X_in, 0)
	# scale all pixels from [-1,1] to [0,1]
	X_in = (X_in + 1) / 2.0
	X_out = (X_out + 1) / 2.0
	# plot real images
	for i in range(n_samples):
		plt.subplot(2, n_samples, 1 + i)
		plt.axis('off')
		plt.imshow(X_in[i])
	# plot translated image
	for i in range(n_samples):
		plt.subplot(2, n_samples, 1 + n_samples + i)
		plt.axis('off')
		plt.imshow(X_out[i])
	# save plot to file
	filename1 = '%s_generated_plot_%06d.png' % (name, (step+1))
	plt.savefig(filename1)
	plt.close()

# update image pool for fake images
def update_image_pool(pool, images, max_size=50):
	selected = list()
	for image in images:
		if len(pool) < max_size:
			# stock the pool
			pool.append(image)
			selected.append(image)
		elif np.random() < 0.5:
			# use image, but don't add it to the pool
			selected.append(image)
		else:
			# replace an existing image and use replaced image
			ix = randint(0, len(pool))
			selected.append(pool[ix])
			pool[ix] = image
	return asarray(selected)

#train cyclegan models
def train(d_model_A, d_model_B, g_model_AtoB, g_model_BtoA, c_model_AtoB, c_model_BtoA, dataset, n_epochs):
	# define properties of the training run
	n_batch = 1
	# determine the output square shape of the discriminator
	n_patch = d_model_A.output_shape[1]
	# unpack dataset
	trainA, trainB = dataset
	# prepare image pool for fakes
	poolA, poolB = list(), list()
	# calculate the number of batches per training epoch
	bat_per_epo = int(len(trainA) / n_batch)
	# calculate the number of training iterations
	n_steps = bat_per_epo * n_epochs
	# manually enumerate epochs
	for i in range(n_steps):
		# select a batch of real samples
		X_realA, y_realA = generate_real_samples(trainA, n_batch, n_patch)
		X_realB, y_realB = generate_real_samples(trainB, n_batch, n_patch)
		# generate a batch of fake samples
		X_fakeA, y_fakeA = generate_fake_samples(g_model_BtoA, X_realB, n_patch)
		X_fakeB, y_fakeB = generate_fake_samples(g_model_AtoB, X_realA, n_patch)
		# update fakes from pool
		X_fakeA = update_image_pool(poolA, X_fakeA)
		X_fakeB = update_image_pool(poolB, X_fakeB)
		# update generator B->A via adversarial and cycle loss
		g_loss2, _, _, _, _  = c_model_BtoA.train_on_batch([X_realB, X_realA], [y_realA, X_realA, X_realB, X_realA])
		# update discriminator for A -> [real/fake]
		dA_loss1 = d_model_A.train_on_batch(X_realA, y_realA)
		dA_loss2 = d_model_A.train_on_batch(X_fakeA, y_fakeA)
		# update generator A->B via adversarial and cycle loss
		g_loss1, _, _, _, _ = c_model_AtoB.train_on_batch([X_realA, X_realB], [y_realB, X_realB, X_realA, X_realB])
		# update discriminator for B -> [real/fake]
		dB_loss1 = d_model_B.train_on_batch(X_realB, y_realB)
		dB_loss2 = d_model_B.train_on_batch(X_fakeB, y_fakeB)
		# summarize performance
		print('>%d, dA[%.3f,%.3f] dB[%.3f,%.3f] g[%.3f,%.3f]' % (i+1, dA_loss1,dA_loss2, dB_loss1,dB_loss2, g_loss1,g_loss2))
		# evaluate the model performance every so often
		if (i+1) % (bat_per_epo * 1) == 0:
			# plot A->B translation
			summarize_performance(i, g_model_AtoB, trainA, 'AtoB')
			# plot B->A translation
			summarize_performance(i, g_model_BtoA, trainB, 'BtoA')
		if (i+1) % (bat_per_epo * 5) == 0:
			# save the models
			save_models(i, g_model_AtoB, g_model_BtoA)



########################## Making the whole thing run 

# with tf.device('/GPU:0'):
dataset = load_real_samples('monet_2_360.npz')
print('Loaded', dataset[0].shape, dataset[1].shape)
# define input shape based on the loaded dataset
image_shape = dataset[0].shape[1:]
# generator: A -> B
g_model_AtoB = define_generator(image_shape)
# generator: B -> A
g_model_BtoA = define_generator(image_shape)
# discriminator: A -> [real/fake]
d_model_A = define_discriminator(image_shape)
# discriminator: B -> [real/fake]
d_model_B = define_discriminator(image_shape)
# composite: A -> B -> [real/fake, A]
c_model_AtoB = define_composite_model(g_model_AtoB, d_model_B, g_model_BtoA, image_shape)
# composite: B -> A -> [real/fake, B]
c_model_BtoA = define_composite_model(g_model_BtoA, d_model_A, g_model_AtoB, image_shape)
# train models
train(d_model_A, d_model_B, g_model_AtoB, g_model_BtoA, c_model_AtoB, c_model_BtoA, dataset, N_EPOCHS)
#from tensorflow.keras import backend as K # config = tf.compat.v1.ConfigProto(intra_op_parallelism_threads=8, #                         inter_op_parallelism_threads=8, #                         allow_soft_placement=True, #                         device_count = {'CPU' : 1, #                                         'GPU' : 1} #                        ) # session = tf.compat.v1.Session(config=config) # K.set_session(session) # load all images in a directory into memory # enumerate filenames in directory, assume all are images # load and resize the image # convert to numpy array # store #dataset path #load dataset A - 360° photos ################################################### BEWARE OF SIZE REMEMBER TO CHANGE SIZE IN LOAD IMAGES FUNCTION #load dataset B - monet photos #save as compressed numpy array #Let's get plotting #optional plotting - WE COULD ADD ARGUMENT PARSER FOR THIS # SHOW pictures #SHOW monet paintings (or other chosen artists) #Starting to actually build model from scratch #################### 70x70 PatchGAN discriminator model, useful as can take different dimension sizes as input # define the discriminator model # weight initialization # source image input # C64 # C128 # C256 # C512 # second last output layer # patch output # define model # compile model ################### Generator models #Encoder-decoder architecture. Downsampling, encoding down to bottleneck then upsampling and creating output image # weight initialization # first layer convolutional layer # second convolutional layer # concatenate merge channel-wise with input layer # define the standalone generator model # weight initialization # image input # c7s1-64 # d128 # d256 # R256 # u128 # u64 # c7s1-3 # define model # ensure the model we're updating is trainable # mark discriminator as not trainable # mark other generator model as not trainable # discriminator element # identity element # forward cycle # backward cycle # define model graph # define optimization algorithm configuration # compile model with weighting of least squares loss and L1 loss # load and prepare training images # load the dataset # unpack arrays # scale from [0,255] to [-1,1] # select a batch of random samples, returns images and target # choose random instances # retrieve selected images # generate 'real' class labels (1) # generate a batch of images, returns images and targets # generate fake instance # create 'fake' class labels (0) # save the generator models to file - need h5py installed to run this # save the first generator model # save the second generator model # generate samples and save as a plot and save the model # select a sample of input images # generate translated images # scale all pixels from [-1,1] to [0,1] # plot real images # plot translated image # save plot to file # update image pool for fake images # stock the pool # use image, but don't add it to the pool # replace an existing image and use replaced image #train cyclegan models # define properties of the training run # determine the output square shape of the discriminator # unpack dataset # prepare image pool for fakes # calculate the number of batches per training epoch # calculate the number of training iterations # manually enumerate epochs # select a batch of real samples # generate a batch of fake samples # update fakes from pool # update generator B->A via adversarial and cycle loss # update discriminator for A -> [real/fake] # update generator A->B via adversarial and cycle loss # update discriminator for B -> [real/fake] # summarize performance # evaluate the model performance every so often # plot A->B translation # plot B->A translation # save the models ########################## Making the whole thing run # with tf.device('/GPU:0'): # define input shape based on the loaded dataset # generator: A -> B # generator: B -> A # discriminator: A -> [real/fake] # discriminator: B -> [real/fake] # composite: A -> B -> [real/fake, A] # composite: B -> A -> [real/fake, B] # train models