DBSRCNN/DBSRCNN_keras/DBSRCNN.py

# -*- coding: utf-8 -*-
"""
Created on Thu Nov  8 12:35:32 2018

@author: fdrea
"""

from __future__ import print_function
import keras
from keras.layers import Input, Convolution2D, Conv2DTranspose, merge, Conv2D
from keras.models import Model
from keras.callbacks import LearningRateScheduler
from keras import backend as K
from keras.optimizers import Adam
import matplotlib.pyplot as plt
import h5py
import math
import sys

from math import sqrt

from keras.callbacks import ModelCheckpoint
import generate_data as pd
import pandas
import numpy
import cv2

#for save excel:
import xlwt
from tempfile import TemporaryFile


scale = 2
batch_size = 64
nb_epoch = 80

    
def PSNRLoss(y_true, y_pred):
    """
    PSNR is Peek Signal to Noise Ratio, which is similar to mean squared error.

    It can be calculated as
    PSNR = 20 * log10(MAXp) - 10 * log10(MSE)

    When providing an unscaled input, MAXp = 255. Therefore 20 * log10(255)== 48.1308036087.
    However, since we are scaling our input, MAXp = 1. Therefore 20 * log10(1) = 0.
    Thus we remove that component completely and only compute the remaining MSE component.
    
    """
    
    #psnr= k.reduce_mean(k.square(y_pred - y_true))
    
    return 10.0 * K.log(1.0 / (K.mean(K.square(y_pred - y_true)))) / K.log(10.0)

def step_decay(epoch):
	initial_lrate = 0.001
	drop = 0.1
	epochs_drop = 20
	lrate = initial_lrate * math.pow(drop, math.floor((1+epoch)/epochs_drop))
	return lrate

def SRCNN_model():
    '''
    _input = Input(shape=(None, None, 1), name='input')

    EES = Conv2D(filters=16, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='relu')(_input)
    EES = Conv2DTranspose(filters=32, kernel_size=(14, 14), strides=(2, 2), padding='same', activation='relu')(EES)
    out = Conv2D(filters=1, kernel_size=(5, 5), strides=(1, 1), activation='relu', padding='same')(EES)

    model = Model(input=_input, output=out)
    '''
    
    x = Input(shape = (None, None, 1), name='input')
    c1 = Convolution2D(32, (9, 9), padding="same", kernel_initializer="he_normal", activation="relu")(x)
    c2 = Convolution2D(32, (5, 5), padding="same", kernel_initializer="he_normal", activation="relu")(c1)
    m= keras.layers.concatenate([c1, c2])
    c3 = Convolution2D(32, (5, 5), padding="same", kernel_initializer="he_normal", activation="relu")(m)
    c4 = Convolution2D(32, (5, 5), padding="same", kernel_initializer="he_normal", activation="relu")(c3)
    c5 = Convolution2D(32, (5, 5), padding="same", kernel_initializer="he_normal", activation="relu")(c4)
    c6 = Convolution2D(32, (5, 5), padding="same", kernel_initializer="he_normal", activation="relu")(c5)
    c7 = Convolution2D(32, (5, 5), padding="same", kernel_initializer="he_normal", activation="relu")(c6)
    c8 = Convolution2D(32, (5, 5), padding="same", kernel_initializer="he_normal", activation="relu")(c7)
    c9 = Convolution2D(32, (5, 5), padding="same", kernel_initializer="he_normal", activation="relu")(c8)
    c10 = Convolution2D(1, (5, 5), padding="same", kernel_initializer="he_normal")(c9)
    model = Model(inputs = x, outputs = c10)


    return model


def SRCNN_train():
    '''
    EES = model_EES16()
    EES.compile(optimizer=adam(lr=0.0003), loss='mse')
    print (EES.summary())

    data, label = pd.read_training_data("./train1.h5")
    val_data, val_label = pd.read_training_data("./val.h5")

    checkpoint = ModelCheckpoint("EES2_check.h5", monitor='val_loss', verbose=1, save_best_only=True,
                                 save_weights_only=False, mode='min')
    callbacks_list = [checkpoint]

    history_callback = EES.fit(data, label, batch_size=64, validation_data=(val_data, val_label),
                               callbacks=callbacks_list, shuffle=True, nb_epoch=10, verbose=1)
    pandas.DataFrame(history_callback.history).to_csv("history.csv")
    EES.save_weights("EES2_final.h5")
    '''
    
    model = SRCNN_model()
    
    ##compile
    adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-8) 
    model.compile(loss='mean_squared_error', metrics=[PSNRLoss], optimizer = adam) #'mse'
    print (model.summary())


    #data, label = pd.read_training_data("./data/train91.h5")
    #data, label = pd.read_training_data("./data/train_aug.h5")
    data, label = pd.read_training_data("./data/train291.h5")
    val_data, val_label = pd.read_training_data("./data/val.h5")
    
    # save the best model weights
    ModelCheckpoint("SRCNN_check.h5", monitor='val_loss', verbose=0, save_best_only=True, 
                                    save_weights_only=False, mode='min')
    
    # learning schedule callback
    lrate = LearningRateScheduler(step_decay)
    print('lrate=',lrate)
    #callbacks_list = [lrate]
    history = model.fit(data, label, batch_size=batch_size, epochs=nb_epoch, callbacks = [lrate],
          verbose=1, validation_data=(val_data, val_label))            
    print(history.history.keys())
    
    # to save history
    pandas.DataFrame(history.history).to_csv("history.csv")


    #save model and weights
    json_string = model.to_json()  
    open('DBSRCNN_model.json','w').write(json_string)  
    model.save_weights('DBSRCNN_model_weights.h5')
    model.save('dbsrcnn_model.h5')

    # summarize history for Peak signal to noise ratio (PSNR)
    plt.figure()
    plt.plot(history.history['PSNRLoss'])
    plt.plot(history.history['val_PSNRLoss'])
    plt.title('Peak Signal to Noise Ratio')
    plt.ylabel('PSNR/dB')
    plt.xlabel('Epoch')
    plt.legend(['Train', 'Test'], loc='lower right')
    #plt.show()
    plt.grid(True,which="both",ls="-")
    plt.savefig('Epoch and PSNR SR.png')

    # summarize history for loss
    plt.figure()
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('Model Loss')
    plt.ylabel('Loss')
    plt.xlabel('Epoch')
    plt.legend(['Train', 'Test'], loc='upper right')
    #plt.show()
    plt.grid(True,which="both",ls="-")
    plt.savefig('Epoch and Loss SR.png')
    
    '''
    # for save excel file:
    book = xlwt.Workbook()
    sheet1 = book.add_sheet('sheet1')
    for i,e in enumerate(history.history['val_PSNRLoss']):
        sheet1.write(i,1,e)

    name = "PSNR_Validation.xls"
    book.save(name)
    book.save(TemporaryFile())
    '''
    
###====================================================================================================
def SRCNN_predict():
    #IMG_NAME = "./butterfly_GT.bmp"
    #down_NAME = "down.jpg"
    #INPUT_NAME = "input.jpg"
    IMG_NAME = sys.argv[1]
    down_NAME = sys.argv[2]
    INPUT_NAME = sys.argv[3]
    OUTPUT_NAME = sys.argv[4]+'_output.jpg'
    
    b = 0.1 #blur sigma
      
    label = cv2.imread(IMG_NAME)
    shape = label.shape

    img = cv2.resize(label, (shape[1] // scale, shape[0] // scale), cv2.INTER_CUBIC)
    cv2.imwrite(down_NAME, img)
    
    img = cv2.resize(img, (img.shape[1] * scale, img.shape[0] * scale), cv2.INTER_CUBIC)
    img = cv2.GaussianBlur(img, (0,0), sigmaX = b)
    cv2.imwrite(INPUT_NAME, img)

    SRCNN = SRCNN_model()
    #SRCNN.load_weights("SRCNN_check.h5")
    SRCNN.load_weights("DBSRCNN_model_weights.h5")

    img = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
    Y = numpy.zeros((1, img.shape[0], img.shape[1], 1))
    Y[0, :, :, 0] = img[:, :, 0].astype(float) / 255.
    img = cv2.cvtColor(label, cv2.COLOR_BGR2YCrCb)

    pre = SRCNN.predict(Y, batch_size=1) * 255.
    pre[pre[:] > 255] = 255
    pre = numpy.uint8(pre)
    img[:, :, 0] = pre[0, :, :, 0]
    img = cv2.cvtColor(img, cv2.COLOR_YCrCb2BGR)
    cv2.imwrite(OUTPUT_NAME, img)

    # psnr calculation:
    im1 = cv2.imread(IMG_NAME, cv2.IMREAD_COLOR)
    im1 = cv2.cvtColor(im1, cv2.COLOR_BGR2YCrCb)
    im2 = cv2.imread(INPUT_NAME, cv2.IMREAD_COLOR)
    im2 = cv2.cvtColor(im2, cv2.COLOR_BGR2YCrCb)
    im2 = cv2.resize(im2, (img.shape[1], img.shape[0]))
    cv2.imwrite("Bicubic.jpg", cv2.cvtColor(im2, cv2.COLOR_YCrCb2BGR))
    im3 = cv2.imread(OUTPUT_NAME, cv2.IMREAD_COLOR)
    im3 = cv2.cvtColor(im3, cv2.COLOR_BGR2YCrCb)

    print ("Bicubic:")
    print (cv2.PSNR(im1[:, :, 0], im2[:, :, 0]))
    #m = PSNRLoss(im1.astype('float32')[:, :, 0], im2.astype('float32')[:, :, 0])
    #print (m)
    print ("SRCNN:")
    print (cv2.PSNR(im1[:, :, 0], im3[:, :, 0]))
    
    
###======================================================================================================
# Function by gcalmettes from http://stackoverflow.com/questions/11159436/multiple-figures-in-a-single-window
def plot_figures(figures, nrows=1, ncols=1, titles=False):
    """Plot a dictionary of figures.
    Parameters
    ----------
    figures : <title, figure> dictionary
    ncols : number of columns of subplots wanted in the display
    nrows : number of rows of subplots wanted in the figure
    """
    fig, axeslist = plt.subplots(ncols=ncols, nrows=nrows)
    for ind, title in enumerate(sorted(figures.keys(), key=lambda s: int(s[3:]))):
        axeslist.ravel()[ind].imshow(figures[title], cmap=plt.gray())
        if titles:
            axeslist.ravel()[ind].set_title(title)

    for ind in range(nrows*ncols):
        axeslist.ravel()[ind].set_axis_off()

    if titles:
        plt.tight_layout()
    plt.show()


def get_dim(num):
    """
    Simple function to get the dimensions of a square-ish shape for plotting
    num images
    """
    s = sqrt(num)
    if round(s) < s:
        return (int(s), int(s)+1)
    else:
        return (int(s)+1, int(s)+1)


def feature_map_visilization(model, _input):
    # Get the convolutional layers
    conv_layers = [layer for layer in model.layers if isinstance(layer, Conv2D)]

    # Use a keras function to extract the conv layer data
    convout_func = K.function([model.layers[0].input, K.learning_phase()], [layer.output for layer in conv_layers])
    conv_imgs_filts = convout_func([_input, 0])
    # Also get the prediction so we know what we predicted
    predictions = model.predict(_input)

    imshow = plt.imshow  # alias

    # Show the original image
    plt.title("Image used:")
    imshow(_input[0, :, :, 0], cmap='gray')
    plt.tight_layout()
    plt.show()

    # Plot the filter images
    for i, conv_imgs_filt in enumerate(conv_imgs_filts):
        conv_img_filt = conv_imgs_filt[0]
        print("Visualizing Convolutions Layer %d" % i)
        # Get it ready for the plot_figures function
        fig_dict = {'flt{0}'.format(i): conv_img_filt[:, :, i] for i in range(conv_img_filt.shape[-1])}
        plot_figures(fig_dict, *get_dim(len(fig_dict)))

    cv2.waitKey(0)


def vilization_and_show():
    model = SRCNN_model()
    #model.load_weights("SRCNN_check.h5")
    model.load_weights("DBSRCNN_model_weights.h5")
    #IMG_NAME = "comic.bmp"
    IMG_NAME = "./butterfly_GT.bmp"
    INPUT_NAME = "input.jpg"

    img = cv2.imread(IMG_NAME)
    shape = img.shape
    img = cv2.resize(img, (shape[1] // 2, shape[0] // 2), cv2.INTER_CUBIC)
    cv2.imwrite(INPUT_NAME, img)

    img = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
    Y = numpy.zeros((1, img.shape[0], img.shape[1], 1))
    Y[0, :, :, 0] = img[:, :, 0]

    feature_map_visilization(model, Y)

####===================================================================================================
if __name__ == "__main__":
    #SRCNN_train()
    SRCNN_predict()
    #vilization_and_show()