def run_model(self,
writer,
set_path,
preload_epoch,
epoch,
training=False,
shuffle=False,
pretrain=False):
if not training:
print('VALIDATION')
image_loader = ImageLoader(batch_size=2, image_dir=set_path)
if shuffle:
image_loader.shuffle_data()
batch_gen = image_loader.getImages()
for i, (batch_x, batch_y) in enumerate(batch_gen):
subbatch = 100
feed = {
'tf_x:0': batch_x,
'tf_y:0': batch_y,
'tf_training:0': training
}
if (i % subbatch == 0):
print('batch ' + str(i) + '/' + str(image_loader.batch_count))
loss = self.sess.run(self.merged, feed_dict=feed)
writer.add_summary(loss)
if (training and i % 1000 == 0):
self.save(epoch=preload_epoch + epoch)
if (not pretrain):
_ = self.sess.run('train_op_disc', feed_dict=feed)
_ = self.sess.run('train_op', feed_dict=feed)
else:
loss, _ = self.sess.run([self.mse_loss_summ, 'train_mse_op'],
feed_dict=feed)
writer.add_summary(loss)
srgan.train(preload_epoch=preload_epoch,
initialize=not preload_model,
validation_set_path="./ImageNet/TestImages",
pretrain=pretrain)
del srgan
if not train:
srgan = SrGan(epochs=epoch)
srgan.load(epoch=preload_epoch, path='./mse-vgg-gen-model/')
if not demo:
il = ImageLoader(batch_size=10, image_dir="./ImageNet/TestImages")
else:
il = ImageLoader(batch_size=1,
image_dir="./ImageNet/DemoImages",
shrink=False)
for i, (input_images, target_images) in enumerate(il.getImages()):
preds = srgan.predict(input_images)
for i, img in enumerate(preds):
pictures = [
cv2.resize(input_images[i],
dsize=(input_images[i].shape[1] * 4,
input_images[i].shape[0] * 4),
interpolation=cv2.INTER_NEAREST),
cv2.resize(input_images[i],
dsize=(input_images[i].shape[1] * 4,
input_images[i].shape[0] * 4),
interpolation=cv2.INTER_CUBIC), (img + 1) / 2,
(target_images[i] + 1) / 2,
((img + 1) / 2 - (target_images[i] + 1) / 2)**2
]
