def train(epochs, batch_size, input_dir, output_dir, model_save_dir, number_of_images, train_test_ratio, image_extension): # Loading images x_train_lr, x_train_hr, x_test_lr, x_test_hr = \ Utils.load_training_data(input_dir, image_extension, image_shape, number_of_images, train_test_ratio) print('======= Loading VGG_loss ========') # Loading VGG loss loss = VGG_LOSS(image_shape) loss2 = VGG_LOSS(image_shape) print('====== VGG_LOSS =======', loss) batch_count = int(x_train_hr.shape[0] / batch_size) print('====== Batch_count =======', batch_count) shape = (image_shape[0] // downscale_factor, image_shape[1] // downscale_factor, image_shape[2]) print('====== Shape =======', shape) # Generator description generator = Generator(shape).generator() complex_generator = complex_Generator(shape).generator() # Discriminator description discriminator = Discriminator(image_shape).discriminator() discriminator2 = Discriminator(image_shape).discriminator() optimizer = Utils_model.get_optimizer() generator.compile(loss=loss.vgg_loss, optimizer=optimizer) complex_generator.compile(loss=loss2.vgg_loss, optimizer=optimizer) discriminator.compile(loss="binary_crossentropy", optimizer=optimizer) discriminator2.compile(loss="binary_crossentropy", optimizer=optimizer) gan = get_gan_network(discriminator, shape, generator, optimizer, loss.vgg_loss) complex_gan = get_gan_network(discriminator2, shape, complex_generator, optimizer, loss2.vgg_loss) loss_file = open(model_save_dir + 'losses.txt', 'w+') loss_file.close() for e in range(1, epochs + 1): print('-' * 15, 'Epoch %d' % e, '-' * 15) for _ in tqdm(range(batch_count)): rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] generated_images_sr = generator.predict(image_batch_lr) generated_images_csr = complex_generator.predict(image_batch_lr) real_data_Y = np.ones(batch_size) - np.random.random_sample(batch_size) * 0.2 fake_data_Y = np.random.random_sample(batch_size) * 0.2 discriminator.trainable = True discriminator2.trainable = True d_loss_real = discriminator.train_on_batch(image_batch_hr, real_data_Y) d_loss_fake = discriminator.train_on_batch(generated_images_sr, fake_data_Y) discriminator_loss = 0.5 * np.add(d_loss_fake, d_loss_real) d_loss_creal = discriminator2.train_on_batch(image_batch_hr, real_data_Y) d_loss_cfake = discriminator2.train_on_batch(generated_images_csr, fake_data_Y) discriminator_c_loss = 0.5 * np.add(d_loss_cfake, d_loss_creal) ######## rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] gan_Y = np.ones(batch_size) - np.random.random_sample(batch_size) * 0.2 discriminator.trainable = False discriminator2.trainable = False gan_loss = gan.train_on_batch(image_batch_lr, [image_batch_hr, gan_Y]) gan_c_loss = complex_gan.train_on_batch(image_batch_lr, [image_batch_hr, gan_Y]) print("discriminator_loss : %f" % discriminator_loss) print("gan_loss :", gan_loss) print("gan_c_loss :", gan_c_loss) gan_loss = str(gan_loss) loss_file = open(model_save_dir + 'losses.txt', 'a') loss_file.write('epoch%d : gan_loss = %s ; discriminator_loss = %f\n' % (e, gan_loss, discriminator_loss)) loss_file.close() if e % 1 == 0: Utils.plot_generated_images(output_dir, e, generator,complex_generator, x_test_hr, x_test_lr) if e % 50 == 0: generator.save(model_save_dir + 'gen_model%d.h5' % e) discriminator.save(model_save_dir + 'dis_model%d.h5' % e)
def train(epochs, batch_size, input_dir, output_dir, model_save_dir, number_of_images, train_test_ratio): x_train_lr, x_train_hr, x_test_lr, x_test_hr = Utils.load_training_data( input_dir, '.png', number_of_images, train_test_ratio) loss = VGG_LOSS(image_shape) batch_count = int(x_train_hr.shape[0] / batch_size) shape = (image_shape[0] // downscale_factor, image_shape[1] // downscale_factor, image_shape[2]) # Not good generator, _ = Generator(shape).generator() discriminator = Discriminator(image_shape).discriminator() optimizer = Utils_model.get_optimizer() generator.compile(loss=loss.vgg_loss, optimizer=optimizer) discriminator.compile(loss="binary_crossentropy", optimizer=optimizer) gan = get_gan_network(discriminator, shape, generator, optimizer, loss.vgg_loss) loss_file = open(model_save_dir + 'losses.txt', 'w+') loss_file.close() for e in range(1, epochs + 1): print('-' * 15, 'Epoch %d' % e, '-' * 15) for _ in tqdm(range(batch_count)): rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] generated_images_sr = generator.predict(image_batch_lr) real_data_Y = np.ones( batch_size) - np.random.random_sample(batch_size) * 0.2 fake_data_Y = np.random.random_sample(batch_size) * 0.2 discriminator.trainable = True d_loss_real = discriminator.train_on_batch(image_batch_hr, real_data_Y) d_loss_fake = discriminator.train_on_batch(generated_images_sr, fake_data_Y) discriminator_loss = 0.5 * np.add(d_loss_fake, d_loss_real) rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] gan_Y = np.ones( batch_size) - np.random.random_sample(batch_size) * 0.2 discriminator.trainable = False gan_loss = gan.train_on_batch(image_batch_lr, [image_batch_hr, gan_Y]) print("discriminator_loss : %f" % discriminator_loss) print("gan_loss :", gan_loss) gan_loss = str(gan_loss) loss_file = open(model_save_dir + 'losses.txt', 'a') loss_file.write('epoch%d : gan_loss = %s ; discriminator_loss = %f\n' % (e, gan_loss, discriminator_loss)) loss_file.close() if e == 1 or e % 5 == 0: Utils.plot_generated_images(output_dir, e, generator, x_test_hr, x_test_lr) # generator.save(model_save_dir + 'gen_model%d.h5' % e) # generator.save_weights(model_save_dir + 'gen_w%d.h5' % e) # exit() if e % 500 == 0: # generator.save(model_save_dir + 'gen_model%d.h5' % e) generator.save_weights(model_save_dir + 'gen_w%d.h5' % e)
def train(epochs, batch_size, input_dir, output_dir, model_save_dir, number_of_images, train_test_ratio, image_extension): # Loading images x_train_lr, x_train_hr, x_test_lr, x_test_hr = \ Utils.load_training_data(input_dir, image_extension, image_shape, number_of_images, train_test_ratio) # convert to loading PATCHES #num_samples = dataset_info['num_samples'][1] print('======= Loading VGG_loss ========') # Loading VGG loss # convert to 3 channels #img_input = Input(shape=original_image_shape) #image_shape_gray = Concatenate()([img_input, img_input, img_input]) #image_shape_gray = Concatenate()([original_image_shape, original_image_shape]) #image_shape_gray = Concatenate()([image_shape_gray,original_image_shape]) #image_shape = patch_shape #experimental_run_tf_function=False loss = VGG_LOSS(image_shape) # was image_shape print('====== VGG_LOSS =======', loss) # 1 channel #image_shape= original_image_shape batch_count = int(x_train_hr.shape[0] / batch_size) #batch_count = int(x_train_hr_patch.shape[0] / batch_size) # for patch print('====== Batch_count =======', batch_count) shape = (image_shape[0] // downscale_factor, image_shape[1] // downscale_factor, image_shape[2] ) # commented by Furat #shape = (image_shape[0] // downscale_factor, image_shape[1] // downscale_factor) print('====== Shape =======', shape) # Generator description generator = Generator(shape).generator() # Discriminator description discriminator = Discriminator(image_shape).discriminator() optimizer = Utils_model.get_optimizer() generator.compile(loss=loss.vgg_loss, optimizer=optimizer) discriminator.compile(loss="binary_crossentropy", optimizer=optimizer) gan = get_gan_network(discriminator, shape, generator, optimizer, loss.vgg_loss) loss_file = open(model_save_dir + 'losses.txt', 'w+') loss_file.close() ## restore the patches into 1 image: # x_train_hr should have a whole image insted of patches? ###### # input_data= x_train_hr # patch_shape = train_conf['patch_shape'] # output_shape = train_conf['output_shape'] # num_chs = num_modalities * dataset_info['RGBch'] # number of total channels # if input_data.ndim == 6: # augmentation case # num_samples = dataset_info['num_samples'][1] #num_samples = 3 ###### #lr_data= x_train_lr #hr_data= x_train_hr #for patch_idx in range (num_patches): #this_input_data = np.reshape(input_data[:,:,patch_idx], input_data.shape[:2]+input_data.shape[3:]) #this_hr_patch, this_lr_patch = overlap_patching(gen_conf, train_conf, x_train_hr) #this_output_patch, out = overlap_patching(gen_conf, train_conf, output_data) # take patches: this_hr_patch, = extract_2d(x_train_hr, (32, 32)) this_lr_patch = extract_2d(x_train_lr, (32, 32)) x_train_lr = this_lr_patch x_train_hr = this_hr_patch #convert to grayscale #x_train_hr= tf.image.rgb_to_grayscale(x_train_hr) #x_train_hr= rgb2gray(x_train_hr) #x_train_hr= np.concatenate(x_train_hr,1) #x_train_hr= np.array(x_train_hr) #x_train_lr= tf.image.rgb_to_grayscale(x_train_lr) #x_train_lr= np.array(x_train_lr) #x_train_lr= rgb2gray(x_train_lr) #x_train_lr= np.concatenate(x_train_lr,1) for e in range(1, epochs + 1): print('-' * 15, 'Epoch %d' % e, '-' * 15) for _ in tqdm(range(batch_count)): rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] generated_images_sr = generator.predict(image_batch_lr) real_data_Y = np.ones( batch_size) - np.random.random_sample(batch_size) * 0.2 fake_data_Y = np.random.random_sample(batch_size) * 0.2 discriminator.trainable = True d_loss_real = discriminator.train_on_batch(image_batch_hr, real_data_Y) d_loss_fake = discriminator.train_on_batch(generated_images_sr, fake_data_Y) discriminator_loss = 0.5 * np.add(d_loss_fake, d_loss_real) rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] gan_Y = np.ones( batch_size) - np.random.random_sample(batch_size) * 0.2 discriminator.trainable = False gan_loss = gan.train_on_batch(image_batch_lr, [image_batch_hr, gan_Y]) print("discriminator_loss : %f" % discriminator_loss) print("gan_loss :", gan_loss) gan_loss = str(gan_loss) loss_file = open(model_save_dir + 'losses.txt', 'a') loss_file.write('epoch%d : gan_loss = %s ; discriminator_loss = %f\n' % (e, gan_loss, discriminator_loss)) loss_file.close() if e == 1 or e % 5 == 0: Utils.plot_generated_images(output_dir, e, generator, x_test_hr, x_test_lr) if e % 200 == 0: generator.save(model_save_dir + 'gen_model%d.h5' % e) discriminator.save(model_save_dir + 'dis_model%d.h5' % e)
def train(epochs, batch_size, input_dir, tgt_dir, output_dir, model_save_dir, number_of_images, train_test_ratio, saved_model): x_train_lr, x_train_hr, x_test_lr, x_test_hr = Utils.load_training_data( input_dir, tgt_dir, '.npy', number_of_images, train_test_ratio) # x_train_lr, x_train_hr, x_test_lr, x_test_hr = Utils.load_training_data(input_dir, '.jpg', number_of_images, train_test_ratio) # x_train_hr = np.expand_dims(x_train_hr, axis=3) # x_test_hr = np.expand_dims(x_test_hr, axis=3) # x_train_hr = np.reshape(x_train_hr,(x_train_hr[0], x_train_hr[1], x_train_hr[2], 1)) # x_test_hr = np.reshape(x_test_hr, (x_test_hr[0], x_test_hr[1], x_test_hr[2], 1)) loss = VGG_LOSS(image_shape) batch_count = int(x_train_hr.shape[0] / batch_size) shape = (image_shape[0], image_shape[1], image_shape[2], image_shape[3]) # # generator = Generator(shape).generator() # # model = squeeze(Activation('tanh')(model), 4) # # # discriminator = Discriminator(dis_shape).discriminator() generator = Generator(shape).generator() generator_old = load_model(saved_model, custom_objects={'vgg_loss': loss.vgg_loss}) x_tmp = generator_old.layers[-2].output # generator_old.layers.pop() # generator_old.layers.pop() # for layers in generator_old.layers: # layers.trainable = False # define new layrs # x1 = Conv2D(filters = 64, kernel_size = 3, strides = 1, padding = "same") # x2 = Conv2D(filters=1, kernel_size=5, strides=1, padding="same") # x3 = Activation('tanh') # generator = Concatenate()[generator_old, x1,x2,x3] # gan_input = Input(shape=shape) # x_tmp = mid_out#generator_old(gan_input) # x_tmp = Conv3D(filters=64, kernel_size=9, strides=1, padding="same", name="TransConv2d_1")(x_tmp) x_tmp = Conv3D(filters=128, kernel_size=5, strides=1, padding="same", name="TransConv2d_2")(x_tmp) x_tmp = Conv3D(filters=32, kernel_size=3, strides=1, padding="same", name="TransConv2d_3")(x_tmp) x_tmp = Conv3D(filters=1, kernel_size=1, strides=1, padding="same", name="TransConv2d_4")(x_tmp) # gan_output = Activation('tanh')(x_tmp) generator = Model(inputs=generator_old.input, outputs=x_tmp) # fine tune the layers. for layers in generator.layers[:-4]: layers.trainable = False optimizer = Utils_model.get_optimizer() generator.compile(loss=loss.vgg_loss, optimizer=optimizer) # discriminator.compile(loss="binary_crossentropy", optimizer=optimizer) # gan = get_gan_network(discriminator, shape, generator, optimizer, loss.vgg_loss) loss_file = open(model_save_dir + 'losses.txt', 'w+') loss_file.close() for e in range(1, epochs + 1): print('-' * 15, 'Epoch %d' % e, '-' * 15) for _ in tqdm(range(batch_count)): rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] gan_loss = generator.train_on_batch(image_batch_lr, image_batch_hr) # print("discriminator_loss : %f" % discriminator_loss) print("gan_loss :", gan_loss) gan_loss = str(gan_loss) loss_file = open(model_save_dir + 'losses.txt', 'a') loss_file.write('epoch%d : Resnet_loss = %s ; \n' % (e, gan_loss)) loss_file.close() if e == 1 or e % 5 == 0: rand_nums = np.random.randint(0, x_test_hr.shape[0], size=batch_size) image_batch_hr = x_test_hr[rand_nums] image_batch_lr = x_test_lr[rand_nums] test_loss = generator.test_on_batch(image_batch_lr, image_batch_hr) print("test_loss :", test_loss) test_loss = str(test_loss) loss_file = open(model_save_dir + 'test_losses.txt', 'a') loss_file.write('epoch%d : test_loss = %s ; \n' % (e, test_loss)) loss_file.close() if e % 50 == 0: generator.save(model_save_dir + 'Resnet_model%d.h5' % e)
def train(epochs, batch_size, input_dir, tgt_dir, output_dir, model_save_dir, number_of_images, train_test_ratio): x_train_lr, x_train_hr, x_test_lr, x_test_hr = Utils.load_training_data( input_dir, tgt_dir, '.npy', number_of_images, train_test_ratio) # x_train_lr, x_train_hr, x_test_lr, x_test_hr = Utils.load_training_data(input_dir, '.jpg', number_of_images, train_test_ratio) # x_train_hr = np.expand_dims(x_train_hr, axis=3) # x_test_hr = np.expand_dims(x_test_hr, axis=3) # x_train_hr = np.reshape(x_train_hr,(x_train_hr[0], x_train_hr[1], x_train_hr[2], 1)) # x_test_hr = np.reshape(x_test_hr, (x_test_hr[0], x_test_hr[1], x_test_hr[2], 1)) loss = VGG_LOSS(image_shape) batch_count = int(x_train_hr.shape[0] / batch_size) shape = (image_shape[0], image_shape[1], image_shape[2], image_shape[3]) generator = Generator(shape).generator() # model = squeeze(Activation('tanh')(model), 4) # discriminator = Discriminator(dis_shape).discriminator() optimizer = Utils_model.get_optimizer() generator.compile(loss=loss.vgg_loss, optimizer=optimizer) # discriminator.compile(loss="binary_crossentropy", optimizer=optimizer) # gan = get_gan_network(discriminator, shape, generator, optimizer, loss.vgg_loss) loss_file = open(model_save_dir + 'losses.txt', 'w+') loss_file.close() for e in range(1, epochs + 1): print('-' * 15, 'Epoch %d' % e, '-' * 15) for _ in tqdm(range(batch_count)): rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] generated_images_sr = generator.predict(image_batch_lr) # real_data_Y = np.ones(batch_size) - np.random.random_sample(batch_size)*0.2 # fake_data_Y = np.random.random_sample(batch_size)*0.2 # # discriminator.trainable = True # d_loss_real = discriminator.train_on_batch(image_batch_hr, real_data_Y) # d_loss_fake = discriminator.train_on_batch(generated_images_sr, fake_data_Y) # discriminator_loss = 0.5 * np.add(d_loss_fake, d_loss_real) # # rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) # image_batch_hr = x_train_hr[rand_nums] # image_batch_lr = x_train_lr[rand_nums] # # gan_Y = np.ones(batch_size) - np.random.random_sample(batch_size)*0.2 # discriminator.trainable = False gan_loss = generator.train_on_batch(image_batch_lr, image_batch_hr) # print("discriminator_loss : %f" % discriminator_loss) print("gan_loss :", gan_loss) gan_loss = str(gan_loss) loss_file = open(model_save_dir + 'losses.txt', 'a') loss_file.write('epoch%d : gan_loss = %s ; \n' % (e, gan_loss)) loss_file.close() if e == 1 or e % 5 == 0: flag = 1 # Utils.plot_generated_images(output_dir, e, generator, x_test_hr, x_test_lr) if e % 500 == 0: generator.save(model_save_dir + 'gen_model%d.h5' % e)
def train(epochs, batch_size, input_dir, output_dir, model_save_dir, number_of_images, train_test_ratio, resume_train, downscale_factor, arch): config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) set_session(sess) hr_images, hr_label, lr_images, lr_label = Utils.load_training_data( input_dir, number_of_images) print(hr_images) loss = VGG_LOSS(IMAGE_SHAPE) lr_shape = (IMAGE_SHAPE[0] // downscale_factor, IMAGE_SHAPE[1] // downscale_factor, IMAGE_SHAPE[2]) print(lr_shape) generator = Generator(lr_shape, downscale_factor, arch).generator() discriminator = Discriminator(IMAGE_SHAPE).discriminator() optimizer = Utils_model.get_optimizer() if (resume_train == True): last_epoch_number = Utils.get_last_epoch_number(model_save_dir + 'last_model_epoch.txt') gen_model = model_save_dir + arch + "_gen_model" + str( last_epoch_number) + ".h5" dis_model = model_save_dir + arch + "_dis_model" + str( last_epoch_number) + ".h5" generator.load_weights(gen_model) discriminator.load_weights(dis_model) else: last_epoch_number = 1 generator.compile(loss=loss.vgg_loss, optimizer=optimizer) discriminator.compile(loss="binary_crossentropy", optimizer=optimizer) gan = gan_network(discriminator, lr_shape, generator, optimizer, loss.vgg_loss) for e in range(last_epoch_number, last_epoch_number + epochs): print('-' * 15, 'Epoch %d' % e, '-' * 15) for _ in tqdm(range(1)): rand_nums = np.random.randint(0, hr_images.shape[0], size=batch_size) image_batch_hr = hr_images[rand_nums] image_batch_lr = lr_images[rand_nums] # video_images = lr_images[0] generated_images = generator.predict( image_batch_lr) #array of generated images real_data = np.ones( batch_size) - np.random.random_sample(batch_size) * 0.2 fake_data = np.random.random_sample(batch_size) * 0.2 discriminator.trainable = True discriminator_loss_real = discriminator.train_on_batch( image_batch_hr, real_data) discriminator_loss_fake = discriminator.train_on_batch( generated_images, fake_data) discriminator_loss = 0.5 * np.add( discriminator_loss_fake, discriminator_loss_real) #Mean Of Discriminator Loss rand_nums = np.random.randint(0, hr_images.shape[0], size=batch_size) discriminator.trainable = False gan_loss = gan.train_on_batch(image_batch_lr, [image_batch_hr, real_data]) print("discriminator_loss : %f" % discriminator_loss) print("gan_loss : ", gan_loss) gan_loss = str(gan_loss) # generated_video_image = generator.predict(np.expand_dims(video_images,axis=0)) Utils.save_losses_file(model_save_dir, e, gan_loss, discriminator_loss, arch + '_losses.txt') # image_array.append(cv2.cvtColor(denormalize(generated_video_image[0]),cv2.COLOR_BGR2RGB)) if e % EPOCHS_CHECKPOINT == 0: Utils.save_losses_file(model_save_dir, e, gan_loss, discriminator_loss, 'last_model_epoch.txt') generator.save(model_save_dir + arch + '_gen_model%d.h5' % e) discriminator.save(model_save_dir + arch + '_dis_model%d.h5' % e)
def train(img_shape, epochs, batch_size, rescaling_factor, input_dirs, output_dir, model_save_dir, train_test_ratio, gpu=1): lr_shape = (img_shape[0] // rescaling_factor, img_shape[1] // rescaling_factor, img_shape[2]) img_train_gen, img_test_gen = create_data_generator( input_dirs[1], input_dirs[0], target_size_lr=(lr_shape[0], lr_shape[1]), target_size_hr=(img_shape[0], img_shape[1]), preproc_lr=rescale_imgs_to_neg1_1, preproc_hr=rescale_imgs_to_neg1_1, validation_split=train_test_ratio, batch_size=batch_size) batch_count = int( (len(os.listdir(os.path.join(input_dirs[1], 'ignore'))) / batch_size) * (1 - train_test_ratio)) test_image = [] for img in sorted(os.listdir(os.path.join(input_dirs[1], 'ignore'))): if 'niklas_city_0009' in img: test_image.append( rescale_imgs_to_neg1_1( cv2.imread(os.path.join(input_dirs[1], 'ignore', img)))) print("test length: ", len(test_image)) loss = VGG_LOSS(image_shape) generator = Generator(lr_shape, rescaling_factor).generator() print('memory usage generator: ', get_model_memory_usage(batch_size, generator)) optimizer = Utils_model.get_optimizer() if gpu > 1: try: print("multi_gpu_model generator") par_generator = multi_gpu_model(generator, gpus=2) except: par_generator = generator print("single_gpu_model generator") else: par_generator = generator print("single_gpu_model generator") par_generator.compile(loss=loss.loss, optimizer=optimizer) par_generator.summary() for e in range(1, epochs + 1): print('-' * 15, 'Epoch %d' % e, '-' * 15) for i in tqdm(range(batch_count)): batch = next(img_train_gen) image_batch_hr = batch[1] image_batch_lr = batch[0] if image_batch_hr.shape[0] == batch_size and image_batch_lr.shape[ 0] == batch_size: g_loss = par_generator.train_on_batch(image_batch_lr, image_batch_hr) else: print("weird multi_gpu_model batch error dis: ") print("hr batch shape: ", image_batch_hr.shape) print("lr batch shape: ", image_batch_lr.shape) #if e == 1 or e % 5 == 0: #Utils.generate_test_image(output_dir, e, generator, test_image) if e % 5 == 0: generator.save(os.path.join(model_save_dir, 'srresnet%d.h5' % e)) generator.save(os.path.join(model_save_dir, 'srresnet.h5' % e))
import keras import cv2 from keras.models import load_model from skimage import data, io import Utils_model as Utils_model from Utils_model import VGG_LOSS import numpy as np import matplotlib.pyplot as plt import os image_shape = (120, 120, 3) loss = VGG_LOSS(image_shape) optimizer = Utils_model.get_optimizer() path = os.getcwd() def normalize(input_data): return (input_data.astype(np.float32) - 127.5) / 127.5 def denormalize(input_data): input_data = (input_data + 1) * 127.5 return input_data.astype(np.uint8) def gen(file_name): image = data.imread(path + "/upload/%s" % (file_name)) cv2.imwrite(path + "/gen_output/before1.png", image)
def train(img_shape, epochs, batch_size, rescaling_factor, input_dirs, output_dir, model_save_dir, train_test_ratio): lr_shape = (img_shape[0] // rescaling_factor, img_shape[1] // rescaling_factor, img_shape[2]) img_train_gen, img_test_gen = create_data_generator( input_dirs[1], input_dirs[0], target_size_lr=(lr_shape[0], lr_shape[1]), target_size_hr=(img_shape[0], img_shape[1]), preproc_lr=rescale_imgs_to_neg1_1, preproc_hr=rescale_imgs_to_neg1_1, validation_split=train_test_ratio, batch_size=batch_size) loss = VGG_LOSS(image_shape) batch_count = int( (len(os.listdir(os.path.join(input_dirs[1], 'ignore'))) / batch_size) * (1 - train_test_ratio)) test_image = [] for img in sorted(os.listdir(os.path.join(input_dirs[1], 'ignore'))): if 'niklas_city_0009' in img: test_image.append( rescale_imgs_to_neg1_1( cv2.imread(os.path.join(input_dirs[1], 'ignore', img)))) print("test length: ", len(test_image)) generator = Generator(lr_shape, rescaling_factor).generator() discriminator = Discriminator(img_shape).discriminator() print('memory usage generator: ', get_model_memory_usage(batch_size, generator)) print('memory usage discriminator: ', get_model_memory_usage(batch_size, discriminator)) optimizer = Utils_model.get_optimizer() try: print("multi_gpu_model generator") par_generator = multi_gpu_model(generator, gpus=2) except: par_generator = generator print("single_gpu_model generator") try: print("multi_gpu_model discriminator") par_discriminator = multi_gpu_model(discriminator, gpus=2) except: par_discriminator = discriminator print("single_gpu_model discriminator") par_generator.compile(loss=loss.loss, optimizer=optimizer) par_discriminator.compile(loss="binary_crossentropy", optimizer=optimizer) gan, par_gan = get_gan_network(par_discriminator, lr_shape, par_generator, optimizer, loss.loss, batch_size) par_discriminator.summary() par_generator.summary() par_gan.summary() loss_file = open(model_save_dir + 'losses.txt', 'w+') loss_file.close() for e in range(1, epochs + 1): print('-' * 15, 'Epoch %d' % e, '-' * 15) if e == 100: optimizer.lr = 1e-5 for i in tqdm(range(batch_count)): batch = next(img_train_gen) image_batch_hr = batch[1] image_batch_lr = batch[0] generated_images_sr = generator.predict(image_batch_lr) real_data_Y = np.ones(batch_size) - \ np.random.random_sample(batch_size)*0.2 fake_data_Y = np.random.random_sample(batch_size) * 0.2 par_discriminator.trainable = True if image_batch_hr.shape[0] == batch_size and image_batch_lr.shape[ 0] == batch_size: d_loss_real = par_discriminator.train_on_batch( image_batch_hr, real_data_Y) d_loss_fake = par_discriminator.train_on_batch( generated_images_sr, fake_data_Y) discriminator_loss = 0.5 * np.add(d_loss_fake, d_loss_real) else: print("weird multi_gpu_model batch error dis: ") print("hr batch shape: ", image_batch_hr.shape) print("lr batch shape: ", image_batch_lr.shape) print("gan y shape: ", gan_Y.shape) batch = next(img_train_gen) image_batch_hr = batch[1] image_batch_lr = batch[0] gan_Y = np.ones(batch_size) - \ np.random.random_sample(batch_size)*0.2 discriminator.trainable = False if image_batch_hr.shape[0] == batch_size and image_batch_lr.shape[ 0] == batch_size: gan_loss = par_gan.train_on_batch(image_batch_lr, [image_batch_hr, gan_Y]) else: print("weird multi_gpu_model batch error gan: ") print("hr batch shape: ", image_batch_hr.shape) print("lr batch shape: ", image_batch_lr.shape) print("gan y shape: ", gan_Y.shape) print("discriminator_loss : %f" % discriminator_loss) print("gan_loss :", gan_loss) gan_loss = str(gan_loss) loss_file = open(model_save_dir + '_losses.txt', 'a') loss_file.write('epoch%d : gan_loss = %s ; discriminator_loss = %f\n' % (e, gan_loss, discriminator_loss)) loss_file.close() if e == 1 or e % 5 == 0: Utils.generate_test_image(output_dir, e, generator, test_image) if e % 5 == 0: generator.save(os.path.join(model_save_dir, 'gen_model%d.h5' % e)) discriminator.save( os.path.join(model_save_dir, 'dis_model%d.h5' % e))