title image superresolution using an efficient subpixel cnn xingyu longhttps github comxingyulong date created 20200728 last modified 20200827 description implementing superresolution using efficient subpixel model on bsds500 accelerator gpu converted to keras 3 by md awsfalur rahmanhttps awsaf49 github io introduction espcn efficient subpixel cnn proposed by shi 2016https arxiv orgabs1609 05158 is a model that reconstructs a highresolution version of an image given a lowresolution version it leverages efficient subpixel convolution layers which learns an array of image upscaling filters in this code example we will implement the model from the paper and train it on a small dataset bsds500https www2 eecs berkeley eduresearchprojectscsvisiongroupingresources html bsds500https www2 eecs berkeley eduresearchprojectscsvisiongroupingresources html setup load data bsds500 dataset download dataset we use the builtin keras utils getfile utility to retrieve the dataset we create training and validation datasets via imagedatasetfromdirectory we rescale the images to take values in the range 0 1 scale from 0 255 to 0 1 let s visualize a few sample images we prepare a dataset of test image paths that we will use for visual evaluation at the end of this example crop and resize images let s process image data first we convert our images from the rgb color space to the yuv colour spacehttps en wikipedia orgwikiyuv for the input data lowresolution images we crop the image retrieve the y channel luninance and resize it with the area method use bicubic if you use pil we only consider the luminance channel in the yuv color space because humans are more sensitive to luminance change for the target data highresolution images we just crop the image and retrieve the y channel use tf ops to process let s take a look at the input and target data build a model compared to the paper we add one more layer and we use the relu activation function instead of tanh it achieves better performance even though we train the model for fewer epochs define utility functions we need to define several utility functions to monitor our results plotresults to plot an save an image getlowresimage to convert an image to its lowresolution version upscaleimage to turn a lowresolution image to a highresolution version reconstructed by the model in this function we use the y channel from the yuv color space as input to the model and then combine the output with the other channels to obtain an rgb image plot the result with zoomin area imgarray imgtoarrayimg imgarray imgarray astypefloat32 255 0 create a new figure with a default 111 subplot fig ax plt subplots im ax imshowimgarray 1 originlower plt titletitle zoomfactor 2 0 location upperleft axins zoomedinsetaxesax 2 loc2 axins imshowimgarray 1 originlower specify the limits x1 x2 y1 y2 200 300 100 200 apply the xlimits axins setxlimx1 x2 apply the ylimits axins setylimy1 y2 plt yticksvisiblefalse plt xticksvisiblefalse make the line markinsetax axins loc11 loc23 fcnone ecblue plt savefigstrprefix title png plt show def getlowresimageimg upscalefactor predict the result based on input image and restore the image as rgb ycbcr img convertycbcr y cb cr ycbcr split y imgtoarrayy y y astypefloat32 255 0 input np expanddimsy axis0 out model predictinput outimgy out0 outimgy 255 0 restore the image in rgb color space outimgy outimgy clip0 255 outimgy outimgy reshapenp shapeoutimgy0 np shapeoutimgy1 outimgy pil image fromarraynp uint8outimgy model outimgcb cb resizeoutimgy size pil image bicubic outimgcr cr resizeoutimgy size pil image bicubic outimg pil image mergeycbcr outimgy outimgcb outimgcr convert rgb return outimg define callbacks to monitor training class espcncallbackkeras callbacks callback def initself super init self testimg getlowresimageloadimgtestimgpaths0 upscalefactor store psnr value in each epoch def onepochbeginself epoch logsnone self psnr def onepochendself epoch logsnone printmean psnr for epoch 2f np meanself psnr if epoch 20 0 prediction upscaleimageself model self testimg plotresultsprediction epoch strepoch prediction def ontestbatchendself batch logsnone self psnr append10 math log101 logsloss earlystoppingcallback keras callbacks earlystoppingmonitorloss patience10 checkpointfilepath tmpcheckpoint keras modelcheckpointcallback keras callbacks modelcheckpoint filepathcheckpointfilepath saveweightsonlyfalse monitorloss modemin savebestonlytrue model getmodelupscalefactorupscalefactor channels1 model summary callbacks espcncallback earlystoppingcallback modelcheckpointcallback lossfn keras losses meansquarederror optimizer keras optimizers adamlearningrate0 001 train the model epochs 100 model compile optimizeroptimizer losslossfn model fit trainds epochsepochs callbackscallbacks validationdatavalidds verbose2 the model weights that are considered the best are loaded into the model model loadweightscheckpointfilepath run model prediction and plot the results totalbicubicpsnr 0 0 totaltestpsnr 0 0 for index testimgpath in enumeratetestimgpaths50 60 img loadimgtestimgpath lowresinput getlowresimageimg upscalefactor w lowresinput size0 upscalefactor h lowresinput size1 upscalefactor highresimg img resizew h prediction upscaleimagemodel lowresinput lowresimg lowresinput resizew h lowresimgarr imgtoarraylowresimg highresimgarr imgtoarrayhighresimg predictimgarr imgtoarrayprediction bicubicpsnr tf image psnrlowresimgarr highresimgarr maxval255 testpsnr tf image psnrpredictimgarr highresimgarr maxval255 totalbicubicpsnr bicubicpsnr totaltestpsnr testpsnr print psnr of low resolution image and high resolution image is 4f bicubicpsnr printpsnr of predict and high resolution is 4f testpsnr plotresultslowresimg index lowres plotresultshighresimg index highres plotresultsprediction index prediction printavg psnr of lowres images is 4f totalbicubicpsnr 10 printavg psnr of reconstructions is 4f totaltestpsnr 10 introduction espcn efficient sub pixel cnn proposed by shi 2016 https arxiv org abs 1609 05158 is a model that reconstructs a high resolution version of an image given a low resolution version it leverages efficient sub pixel convolution layers which learns an array of image upscaling filters in this code example we will implement the model from the paper and train it on a small dataset bsds500 https www2 eecs berkeley edu research projects cs vision grouping resources html bsds500 https www2 eecs berkeley edu research projects cs vision grouping resources html setup only for data preprocessing load data bsds500 dataset download dataset we use the built in keras utils get_file utility to retrieve the dataset we create training and validation datasets via image_dataset_from_directory we rescale the images to take values in the range 0 1 scale from 0 255 to 0 1 let s visualize a few sample images we prepare a dataset of test image paths that we will use for visual evaluation at the end of this example crop and resize images let s process image data first we convert our images from the rgb color space to the yuv colour space https en wikipedia org wiki yuv for the input data low resolution images we crop the image retrieve the y channel luninance and resize it with the area method use bicubic if you use pil we only consider the luminance channel in the yuv color space because humans are more sensitive to luminance change for the target data high resolution images we just crop the image and retrieve the y channel use tf ops to process let s take a look at the input and target data build a model compared to the paper we add one more layer and we use the relu activation function instead of tanh it achieves better performance even though we train the model for fewer epochs define utility functions we need to define several utility functions to monitor our results plot_results to plot an save an image get_lowres_image to convert an image to its low resolution version upscale_image to turn a low resolution image to a high resolution version reconstructed by the model in this function we use the y channel from the yuv color space as input to the model and then combine the output with the other channels to obtain an rgb image plot the result with zoom in area create a new figure with a default 111 subplot zoom factor 2 0 location upper left specify the limits apply the x limits apply the y limits make the line return low resolution image to use as model input predict the result based on input image and restore the image as rgb restore the image in rgb color space define callbacks to monitor training the espcncallback object will compute and display the psnr https en wikipedia org wiki peak_signal to noise_ratio metric this is the main metric we use to evaluate super resolution performance store psnr value in each epoch define modelcheckpoint and earlystopping callbacks train the model the model weights that are considered the best are loaded into the model run model prediction and plot the results let s compute the reconstructed version of a few images and save the results
import keras
from keras import layers
from keras import ops
from keras.utils import load_img
from keras.utils import array_to_img
from keras.utils import img_to_array
from keras.preprocessing import image_dataset_from_directory
import tensorflow as tf  

import os
import math
import numpy as np

from IPython.display import display



dataset_url = "http://www.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/BSR/BSR_bsds500.tgz"
data_dir = keras.utils.get_file(origin=dataset_url, fname="BSR", untar=True)
root_dir = os.path.join(data_dir, "BSDS500/data")



crop_size = 300
upscale_factor = 3
input_size = crop_size // upscale_factor
batch_size = 8

train_ds = image_dataset_from_directory(
    root_dir,
    batch_size=batch_size,
    image_size=(crop_size, crop_size),
    validation_split=0.2,
    subset="training",
    seed=1337,
    label_mode=None,
)

valid_ds = image_dataset_from_directory(
    root_dir,
    batch_size=batch_size,
    image_size=(crop_size, crop_size),
    validation_split=0.2,
    subset="validation",
    seed=1337,
    label_mode=None,
)




def scaling(input_image):
    input_image = input_image / 255.0
    return input_image



train_ds = train_ds.map(scaling)
valid_ds = valid_ds.map(scaling)



for batch in train_ds.take(1):
    for img in batch:
        display(array_to_img(img))



dataset = os.path.join(root_dir, "images")
test_path = os.path.join(dataset, "test")

test_img_paths = sorted(
    [
        os.path.join(test_path, fname)
        for fname in os.listdir(test_path)
        if fname.endswith(".jpg")
    ]
)




def process_input(input, input_size, upscale_factor):
    input = tf.image.rgb_to_yuv(input)
    last_dimension_axis = len(input.shape) - 1
    y, u, v = tf.split(input, 3, axis=last_dimension_axis)
    return tf.image.resize(y, [input_size, input_size], method="area")


def process_target(input):
    input = tf.image.rgb_to_yuv(input)
    last_dimension_axis = len(input.shape) - 1
    y, u, v = tf.split(input, 3, axis=last_dimension_axis)
    return y


train_ds = train_ds.map(
    lambda x: (process_input(x, input_size, upscale_factor), process_target(x))
)
train_ds = train_ds.prefetch(buffer_size=32)

valid_ds = valid_ds.map(
    lambda x: (process_input(x, input_size, upscale_factor), process_target(x))
)
valid_ds = valid_ds.prefetch(buffer_size=32)



for batch in train_ds.take(1):
    for img in batch[0]:
        display(array_to_img(img))
    for img in batch[1]:
        display(array_to_img(img))




class DepthToSpace(layers.Layer):
    def __init__(self, block_size):
        super().__init__()
        self.block_size = block_size

    def call(self, input):
        batch, height, width, depth = ops.shape(input)
        depth = depth // (self.block_size**2)

        x = ops.reshape(
            input, [batch, height, width, self.block_size, self.block_size, depth]
        )
        x = ops.transpose(x, [0, 1, 3, 2, 4, 5])
        x = ops.reshape(
            x, [batch, height * self.block_size, width * self.block_size, depth]
        )
        return x


def get_model(upscale_factor=3, channels=1):
    conv_args = {
        "activation": "relu",
        "kernel_initializer": "orthogonal",
        "padding": "same",
    }
    inputs = keras.Input(shape=(None, None, channels))
    x = layers.Conv2D(64, 5, **conv_args)(inputs)
    x = layers.Conv2D(64, 3, **conv_args)(x)
    x = layers.Conv2D(32, 3, **conv_args)(x)
    x = layers.Conv2D(channels * (upscale_factor**2), 3, **conv_args)(x)
    outputs = DepthToSpace(upscale_factor)(x)

    return keras.Model(inputs, outputs)




import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes
from mpl_toolkits.axes_grid1.inset_locator import mark_inset
import PIL


def plot_results(img, prefix, title):
    
    img_array = img_to_array(img)
    img_array = img_array.astype("float32") / 255.0

    
    fig, ax = plt.subplots()
    im = ax.imshow(img_array[::-1], origin="lower")

    plt.title(title)
    
    axins = zoomed_inset_axes(ax, 2, loc=2)
    axins.imshow(img_array[::-1], origin="lower")

    
    x1, x2, y1, y2 = 200, 300, 100, 200
    
    axins.set_xlim(x1, x2)
    
    axins.set_ylim(y1, y2)

    plt.yticks(visible=False)
    plt.xticks(visible=False)

    
    mark_inset(ax, axins, loc1=1, loc2=3, fc="none", ec="blue")
    plt.savefig(str(prefix) + "-" + title + ".png")
    plt.show()


def get_lowres_image(img, upscale_factor):
    
    return img.resize(
        (img.size[0] // upscale_factor, img.size[1] // upscale_factor),
        PIL.Image.BICUBIC,
    )


def upscale_image(model, img):
    
    ycbcr = img.convert("YCbCr")
    y, cb, cr = ycbcr.split()
    y = img_to_array(y)
    y = y.astype("float32") / 255.0

    input = np.expand_dims(y, axis=0)
    out = model.predict(input)

    out_img_y = out[0]
    out_img_y *= 255.0

    
    out_img_y = out_img_y.clip(0, 255)
    out_img_y = out_img_y.reshape((np.shape(out_img_y)[0], np.shape(out_img_y)[1]))
    out_img_y = PIL.Image.fromarray(np.uint8(out_img_y), mode="L")
    out_img_cb = cb.resize(out_img_y.size, PIL.Image.BICUBIC)
    out_img_cr = cr.resize(out_img_y.size, PIL.Image.BICUBIC)
    out_img = PIL.Image.merge("YCbCr", (out_img_y, out_img_cb, out_img_cr)).convert(
        "RGB"
    )
    return out_img





class ESPCNCallback(keras.callbacks.Callback):
    def __init__(self):
        super().__init__()
        self.test_img = get_lowres_image(load_img(test_img_paths[0]), upscale_factor)

    
    def on_epoch_begin(self, epoch, logs=None):
        self.psnr = []

    def on_epoch_end(self, epoch, logs=None):
        print("Mean PSNR for epoch: %.2f" % (np.mean(self.psnr)))
        if epoch % 20 == 0:
            prediction = upscale_image(self.model, self.test_img)
            plot_results(prediction, "epoch-" + str(epoch), "prediction")

    def on_test_batch_end(self, batch, logs=None):
        self.psnr.append(10 * math.log10(1 / logs["loss"]))




early_stopping_callback = keras.callbacks.EarlyStopping(monitor="loss", patience=10)

checkpoint_filepath = "/tmp/checkpoint.keras"

model_checkpoint_callback = keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_filepath,
    save_weights_only=False,
    monitor="loss",
    mode="min",
    save_best_only=True,
)

model = get_model(upscale_factor=upscale_factor, channels=1)
model.summary()

callbacks = [ESPCNCallback(), early_stopping_callback, model_checkpoint_callback]
loss_fn = keras.losses.MeanSquaredError()
optimizer = keras.optimizers.Adam(learning_rate=0.001)



epochs = 100

model.compile(
    optimizer=optimizer,
    loss=loss_fn,
)

model.fit(
    train_ds, epochs=epochs, callbacks=callbacks, validation_data=valid_ds, verbose=2
)


model.load_weights(checkpoint_filepath)



total_bicubic_psnr = 0.0
total_test_psnr = 0.0

for index, test_img_path in enumerate(test_img_paths[50:60]):
    img = load_img(test_img_path)
    lowres_input = get_lowres_image(img, upscale_factor)
    w = lowres_input.size[0] * upscale_factor
    h = lowres_input.size[1] * upscale_factor
    highres_img = img.resize((w, h))
    prediction = upscale_image(model, lowres_input)
    lowres_img = lowres_input.resize((w, h))
    lowres_img_arr = img_to_array(lowres_img)
    highres_img_arr = img_to_array(highres_img)
    predict_img_arr = img_to_array(prediction)
    bicubic_psnr = tf.image.psnr(lowres_img_arr, highres_img_arr, max_val=255)
    test_psnr = tf.image.psnr(predict_img_arr, highres_img_arr, max_val=255)

    total_bicubic_psnr += bicubic_psnr
    total_test_psnr += test_psnr

    print(
        "PSNR of low resolution image and high resolution image is %.4f" % bicubic_psnr
    )
    print("PSNR of predict and high resolution is %.4f" % test_psnr)
    plot_results(lowres_img, index, "lowres")
    plot_results(highres_img, index, "highres")
    plot_results(prediction, index, "prediction")

print("Avg. PSNR of lowres images is %.4f" % (total_bicubic_psnr / 10))
print("Avg. PSNR of reconstructions is %.4f" % (total_test_psnr / 10))