class vggGan(): def __init__(self): self.height = height self.width = width self.channels = channels self.batch_size = batch_size self.epochs = epochs self.line = line self.n_show_image = n_show_image self.vgg = vgg self.optimizerD = optimizerD self.optimizerC = optimizerC self.DG = DataGenerator(X_train, Y_train, batch_size=batch_size) self.DGP = DataGenerator_predict(X_predict, batch_size=batch_size) self.number = number patch = int(self.height / 2**1) self.disc_patch = (patch, patch, 3) self.discriminator = self.build_discriminator() self.discriminator.compile(loss='binary_crossentropy', optimizer=self.optimizerD) self.generator = self.build_generator() self.discriminator.trainable = False side = Input(shape=(self.height, self.width, self.channels)) front = Input(shape=(self.height, self.width, self.channels)) image = self.generator(side) valid = self.discriminator([front, image]) self.combined = Model([side, front], [image, valid]) self.combined.compile(loss=['mae', "mse"], loss_weights=[100, 1], optimizer=self.optimizerC) def conv2d_block(self, layers, filters, kernel_size=(4, 4), strides=2, momentum=0.8, alpha=0.2): input = layers layer = Conv2D(filters=filters, kernel_size=kernel_size, strides=strides, padding="same")(input) layer = BatchNormalization(momentum=momentum)(layer) output = LeakyReLU(alpha=alpha)(layer) return output def Conv2DTranspose_block(self, layers, filters, kernel_size=(4, 4), strides=2, momentum=0.8, alpha=0.2): input = layers layer = Conv2DTranspose(filters=filters, kernel_size=kernel_size, strides=strides, padding='same')(input) layer = BatchNormalization(momentum=momentum)(layer) # output = LeakyReLU(alpha = alpha)(layer) output = ReLU()(layer) return output def build_discriminator(self): input_A = Input(shape=(self.height, self.width, self.channels)) input_B = Input(shape=(self.height, self.width, self.channels)) input_layer = Concatenate(axis=-1)([input_A, input_B]) layer = self.conv2d_block(input_layer, 64) # layer = self.conv2d_block(layer, 128) # layer = self.conv2d_block(layer, 256) # layer = self.conv2d_block(layer, 512) # layer = self.conv2d_block(layer, 1024) output = Conv2D(filters=3, kernel_size=(4, 4), strides=1, padding="same")(layer) model = Model([input_A, input_B], output) model.summary() return model def build_generator(self): input = self.vgg.get_layer("pool5").output conv2 = self.vgg.get_layer("conv2_2").output conv3 = self.vgg.get_layer("conv3_3").output conv4 = self.vgg.get_layer("conv4_3").output conv5 = self.vgg.get_layer("conv5_3").output layers = self.Conv2DTranspose_block(input, 512) layers = Concatenate(axis=-1)([layers, conv5]) layers = self.Conv2DTranspose_block(layers, 512) layers = Concatenate(axis=-1)([layers, conv4]) layers = self.Conv2DTranspose_block(layers, 256) layers = Concatenate(axis=-1)([layers, conv3]) layers = self.Conv2DTranspose_block(layers, 128) layers = Concatenate(axis=-1)([layers, conv2]) output = Conv2DTranspose(filters=3, kernel_size=(4, 4), strides=2, activation='tanh', padding='same')(layers) model = Model(self.vgg.input, output) for layer in model.layers: layer.trainable = False if layer.name == "pool5": break # model.summary() return model def train(self, epochs, batch_size, save_interval): fake = np.zeros((batch_size, ) + self.disc_patch) real = np.ones((batch_size, ) + self.disc_patch) for epoch in range(epochs): for batch in range(self.DG.__len__()): side_images, front_images = self.DG.__getitem__(batch) generated_images = self.generator.predict(side_images) discriminator_fake_loss = self.discriminator.train_on_batch( [front_images, generated_images], fake) discriminator_real_loss = self.discriminator.train_on_batch( [front_images, front_images], real) discriminator_loss = np.add(discriminator_fake_loss, discriminator_real_loss) generator_loss = self.combined.train_on_batch( [side_images, front_images], [front_images, real]) print( '\nTraining epoch : %d \nTraining batch : %d / %d \nLoss of discriminator : %f \nLoss of generator : %f' % (epoch + 1, batch + 1, self.DG.__len__(), discriminator_loss, generator_loss[1])) # if batch % save_interval == 0: # save_path = 'D:/Generated Image/Training' + str(line) + '/' # self.save_image(epoch = epoch, batch = batch, front_image = front_images, side_image = side_images, save_path = save_path) # self.generator.save(save_path+"{1}_{0}.h5".format(str(batch), str(line))) if epoch % 1 == 0: # print(i) save_path = 'D:/Generated Image/Training' + str(line) + '/' self.save_image(epoch=epoch, batch=batch, front_image=front_images, side_image=side_images, save_path=save_path) predict_side_images = self.DGP.__getitem__(0) save_path = 'D:/Generated Image/Predict' + str(line) + '/' self.save_predict_image(epoch=epoch, batch=batch, side_image=predict_side_images, save_path=save_path) self.generator.save( "D:/Generated Image/Predict{2}/{1}_{0}.h5".format( str(epoch), str(line), str(line))) self.DG.on_epoch_end() self.DGP.on_epoch_end() def save_image(self, epoch, batch, front_image, side_image, save_path): generated_image = (0.5 * self.generator.predict(side_image) + 0.5) front_image = (255 * ((front_image) + 1) / 2).astype(np.uint8) side_image = (255 * ((side_image) + 1) / 2).astype(np.uint8) for i in range(self.batch_size): plt.figure(figsize=(8, 2)) plt.subplots_adjust(wspace=0.6) for m in range(n_show_image): generated_image_plot = plt.subplot(1, 3, m + 1 + (2 * n_show_image)) generated_image_plot.set_title('Generated image (front image)') plt.imshow(generated_image[i]) original_front_face_image_plot = plt.subplot( 1, 3, m + 1 + n_show_image) original_front_face_image_plot.set_title( 'Origninal front image') plt.imshow(front_image[i]) original_side_face_image_plot = plt.subplot(1, 3, m + 1) original_side_face_image_plot.set_title('Origninal side image') plt.imshow(side_image[i]) # Don't show axis of x and y generated_image_plot.axis('off') original_front_face_image_plot.axis('off') original_side_face_image_plot.axis('off') self.number += 1 # plt.show() save_path = save_path # Check folder presence if not os.path.isdir(save_path): os.makedirs(save_path) save_name = '%d-%d-%d.png' % (epoch, batch, i) save_name = os.path.join(save_path, save_name) plt.savefig(save_name) plt.close() def save_predict_image(self, epoch, batch, side_image, save_path): generated_image = 0.5 * self.generator.predict(side_image) + 0.5 side_image = (255 * ((side_image) + 1) / 2).astype(np.uint8) for i in range(self.batch_size): plt.figure(figsize=(8, 2)) plt.subplots_adjust(wspace=0.6) for m in range(n_show_image): generated_image_plot = plt.subplot(1, 2, m + 1 + n_show_image) generated_image_plot.set_title('Generated image (front image)') plt.imshow(generated_image[i]) original_side_face_image_plot = plt.subplot(1, 2, m + 1) original_side_face_image_plot.set_title('Origninal side image') plt.imshow(side_image[i]) # Don't show axis of x and y generated_image_plot.axis('off') original_side_face_image_plot.axis('off') self.number += 1 # plt.show() save_path = save_path # Check folder presence if not os.path.isdir(save_path): os.makedirs(save_path) save_name = '%d-%d-%d.png' % (epoch, batch, i) save_name = os.path.join(save_path, save_name) plt.savefig(save_name) plt.close()
class TestModel(): def __init__(self): self.height = height self.width = width self.channels = channels self.batch_size = batch_size self.epochs = epochs self.line = line self.n_show_image = n_show_image self.vgg = vgg self.optimizer = optimizer self.DG = DataGenerator(X_train, Y_train, batch_size=batch_size, dim=(128, 128)) self.DGP = DataGenerator_predict(X_predict, batch_size=batch_size, dim=(128, 128)) self.number = number self.AE = self.build_AE() self.AE.compile(loss='mse', optimizer=self.optimizer) # self.AE.fit_generator(self.DG) # self.AE.save("TESTAE.H5") def conv2d_block(self, layers, filters, kernel_size=(4, 4), strides=2, momentum=0.8, alpha=0.2): input = layers layer = Conv2D(filters=filters, kernel_size=kernel_size, strides=strides, padding="same")(input) layer = BatchNormalization(momentum=momentum)(layer) output = LeakyReLU(alpha=alpha)(layer) return output def Conv2DTranspose_block(self, layers, filters, kernel_size=(4, 4), strides=2, momentum=0.8, alpha=0.2): input = layers layer = Conv2DTranspose(filters=filters, kernel_size=kernel_size, strides=strides, padding='same')(input) layer = BatchNormalization(momentum=momentum)(layer) output = LeakyReLU(alpha=alpha)(layer) return output def build_AE(self): last_layer = vgg.get_layer('pool5').output # vgg_model.summary() conv2 = vgg.get_layer("conv2_2").output conv3 = vgg.get_layer("conv3_3").output conv4 = vgg.get_layer("conv4_3").output conv5 = vgg.get_layer("conv5_3").output x = Conv2D(512, (3, 3), padding='same')(last_layer) x = BatchNormalization()(x) x = LeakyReLU()(x) x = UpSampling2D(2)(x) layers = Concatenate(axis=-1)([x, conv5]) x = Conv2D(512, (3, 3), activation='relu', padding='same')(layers) x = BatchNormalization()(x) x = LeakyReLU()(x) x = UpSampling2D(2)(x) layers = Concatenate(axis=-1)([x, conv4]) x = Conv2D(512, (3, 3), activation='relu', padding='same')(layers) x = BatchNormalization()(x) x = LeakyReLU()(x) x = UpSampling2D(2)(x) layers = Concatenate(axis=-1)([x, conv3]) x = Conv2D(256, (3, 3), activation='relu', padding='same')(layers) x = BatchNormalization()(x) x = LeakyReLU()(x) x = UpSampling2D(2)(x) layers = Concatenate(axis=-1)([x, conv2]) x = Conv2D(128, (3, 3), activation='relu', padding='same')(layers) x = LeakyReLU()(x) x = UpSampling2D(2)(x) decoder = Conv2D(3, (3, 3), padding='same')(x) model = Model(vgg.input, decoder) for layer in model.layers: layer.trainable = False # print(layer.get_weights()) if layer.name == "pool5": break model.compile(loss='mse', optimizer='adam') model.summary() return model def train(self, epochs, batch_size, save_interval): for i in range(epochs): for j in range(self.DG.__len__()): side_images, front_images = self.DG.__getitem__(j) AE_loss = self.AE.train_on_batch(side_images, front_images) print( '\nTraining epoch : %d \nTraining batch : %d / %d \nLoss of generator : %f' % (i + 1, j + 1, self.DG.__len__(), AE_loss)) # if j % save_interval == 0: # # save_path = 'D:/Generated Image/Training' + str(line) + '/' # # self.save_image(epoch = i, batch = j, front_image = front_images, side_image = side_images, save_path = save_path) # predict_side_images = self.DGP.__getitem__(j) # predict_front_images = self.AE.predict_on_batch(predict_side_images) # save_path = 'D:/Generated Image/Predict' + str(line) + '/' # self.save_image(epoch = i, batch = j, front_image = predict_front_images, side_image = predict_side_images, save_path = save_path) if i % 1 == 0: # print(i) # save_path = 'D:/Generated Image/Training' + str(line) + '/' # self.save_image(epoch = i, batch = j, front_image = front_images, side_image = side_images, save_path = save_path) predict_side_images = self.DGP.__getitem__(0) save_path = 'D:/Generated Image/Predict' + str(line) + '/' self.save_predict_image(epoch=i, batch=j, side_image=predict_side_images, save_path=save_path) self.AE.save("D:/Generated Image/Predict{2}/{1}_{0}.h5".format( str(i), str(line), str(line))) self.DG.on_epoch_end() self.DGP.on_epoch_end() def save_image(self, epoch, batch, front_image, side_image, save_path): generated_image = 0.5 * self.AE.predict(side_image) + 0.5 front_image = (255 * ((front_image) + 1) / 2).astype(np.uint8) side_image = (255 * ((side_image) + 1) / 2).astype(np.uint8) for i in range(self.batch_size): plt.figure(figsize=(8, 2)) plt.subplots_adjust(wspace=0.6) for m in range(n_show_image): generated_image_plot = plt.subplot(1, 3, m + 1 + (2 * n_show_image)) generated_image_plot.set_title('Generated image (front image)') plt.imshow(generated_image[i]) original_front_face_image_plot = plt.subplot( 1, 3, m + 1 + n_show_image) original_front_face_image_plot.set_title( 'Origninal front image') plt.imshow(front_image[i]) original_side_face_image_plot = plt.subplot(1, 3, m + 1) original_side_face_image_plot.set_title('Origninal side image') plt.imshow(side_image[i]) # Don't show axis of x and y generated_image_plot.axis('off') original_front_face_image_plot.axis('off') original_side_face_image_plot.axis('off') self.number += 1 # plt.show() save_path = save_path # Check folder presence if not os.path.isdir(save_path): os.makedirs(save_path) save_name = '%d-%d-%d.png' % (epoch, batch, i) save_name = os.path.join(save_path, save_name) plt.savefig(save_name) plt.close() def save_predict_image(self, epoch, batch, side_image, save_path): generated_image = 0.5 * self.AE.predict(side_image) + 0.5 side_image = (255 * ((side_image) + 1) / 2).astype(np.uint8) for i in range(self.batch_size): plt.figure(figsize=(8, 2)) plt.subplots_adjust(wspace=0.6) for m in range(n_show_image): generated_image_plot = plt.subplot(1, 2, m + 1 + n_show_image) generated_image_plot.set_title('Generated image (front image)') plt.imshow(generated_image[i]) original_side_face_image_plot = plt.subplot(1, 2, m + 1) original_side_face_image_plot.set_title('Origninal side image') plt.imshow(side_image[i]) # Don't show axis of x and y generated_image_plot.axis('off') original_side_face_image_plot.axis('off') self.number += 1 # plt.show() save_path = save_path # Check folder presence if not os.path.isdir(save_path): os.makedirs(save_path) save_name = '%d-%d-%d.png' % (epoch, batch, i) save_name = os.path.join(save_path, save_name) plt.savefig(save_name) plt.close()