# -*- coding: utf-8 -*-
"""
Created on Sun Apr 28 18:32:15 2019
@author: wmy
"""
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from PIL import Image
from keras import backend as K
from keras.losses import mean_absolute_error, mean_squared_error
from keras.models import load_model
from keras.optimizers import Adam
import random
from model import wdsr_a, wdsr_b
from optimizer import AdamWithWeightsNormalization
from utils import DataLoader
class SuperResolution(object):
def __init__(self, scale=4, num_res_blocks=32, pretrained_weights=None, name=None):
self.scale = scale
self.num_res_blocks = num_res_blocks
self.model = wdsr_b(scale=scale, num_res_blocks=num_res_blocks)
self.model.compile(optimizer=AdamWithWeightsNormalization(lr=0.001), \
loss=self.mae, metrics=[self.psnr])
if pretrained_weights != None:
self.model.load_weights(pretrained_weights)
print("[OK] weights loaded.")
pass
self.data_loader = DataLoader(scale=scale, crop_size=256)
self.pretrained_weights = pretrained_weights
self.default_weights_save_path = 'weights/wdsr-b-' + \
str(self.num_res_blocks) + '-x' + str(self.scale) + '.h5'
self.name = name
pass
def mae(self, hr, sr):
margin = (tf.shape(hr)[1] - tf.shape(sr)[1]) // 2
hr_crop = tf.cond(tf.equal(margin, 0), lambda: hr, lambda: hr[:, margin:-margin, margin:-margin, :])
hr = K.in_train_phase(hr_crop, hr)
hr.uses_learning_phase = True
return mean_absolute_error(hr, sr)
def psnr(self, hr, sr):
margin = (tf.shape(hr)[1] - tf.shape(sr)[1]) // 2
hr_crop = tf.cond(tf.equal(margin, 0), lambda: hr, lambda: hr[:, margin:-margin, margin:-margin, :])
hr = K.in_train_phase(hr_crop, hr)
hr.uses_learning_phase = True
return tf.image.psnr(hr, sr, max_val=255)
def train(self, epoches=10000, batch_size=8, weights_save_path=None):
if weights_save_path == None:
weights_save_path = self.default_weights_save_path
pass
for epoch in range(epoches):
for batch_i, (lrs, hrs) in enumerate(self.data_loader.batches(batch_size=batch_size)):
temp_loss, temp_psnr = self.model.train_on_batch(lrs, hrs)
print("[epoch: {}/{}][batch: {}/{}][loss: {}][psnr: {}]".format(epoch+1, epoches, \
batch_i+1, self.data_loader.n_batches, temp_loss, temp_psnr))
if (batch_i+1) % 25 == 0:
self.sample(epoch=epoch+1, batch=batch_i+1)
pass
pass
self.model.save_weights(weights_save_path)
print("[OK] weights saved.")
pass
pass
def sample(self, setpath='datasets/train', save_folder='samples', epoch=1, batch=1):
images = self.data_loader.search(setpath)
image = random.choice(images)
hr = self.data_loader.imread(image)
lr = self.data_loader.downsampling(hr)
lr_resize = lr.resize(hr.size)
lr = np.asarray(lr)
sr = self.model.predict(np.array([lr]))[0]
sr = np.clip(sr, 0, 255)
sr = sr.astype('uint8')
lr = Image.fromarray(lr)
sr = Image.fromarray(sr)
lr_resize.save(save_folder + "/" + "epoch_" + str(epoch) + "_batch_" + str(batch) + "_lr.jpg")
sr.save(save_folder + "/" + "epoch_" + str(epoch) + "_batch_" + str(batch) + "_sr.jpg")
hr.save(save_folder + "/" + "epoch_" + str(epoch) + "_batch_" + str(batch) + "_hr.jpg")
pass
pass
sr = SuperResolution(pretrained_weights='./weights/wdsr-b-32-x4.h5')
sr.train()
