Overview

Continuing from the previous session. It will be part3, the last article. Last time: I tried to build a super-resolution method / SRCNN ① Last time: I tried to build a super-resolution method / SRCNN②

1.First of all 2. PC environment 3. Code description 4. At the end

1.First of all

Super-resolution is a technology that improves the resolution of low-resolution images and moving images, and SRCNN uses deep learning to measure results with higher accuracy than conventional methods. It is the method that was done. (Third time)

The full code is also posted on GitHub, so please check there. https://github.com/morisumori/srcnn_keras

2. PC environment

cpu : intel corei7 8th Gen gpu : NVIDIA GeForce RTX 1080ti os : ubuntu 20.04

3. Code description

As you can see from GitHub, it mainly consists of three codes. ・ Datacreate.py → Data set generation program ・ Model.py → SRCNN program ・ Main.py → Execution program I have created a function with datacreate.py and model.py and executed it with main.py.

__ This time, I will explain main.py. __

Description of model.py

`model.py`



import model
import data_create
import argparse
import os
import cv2

import numpy as np
import tensorflow as tf

if __name__ == "__main__":
    
    def psnr(y_true, y_pred):
        return tf.image.psnr(y_true, y_pred, 1, name=None)

    train_height = 33
    train_width = 33
    test_height = 700
    test_width = 700

    mag = 3.0
    cut_traindata_num = 10
    cut_testdata_num = 1

    train_file_path = "./train_data"
    test_file_path = "./test_data"

    BATSH_SIZE = 240
    EPOCHS = 1000
    opt = tf.keras.optimizers.Adam(learning_rate=0.0001)

    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', type=str, default='srcnn', help='srcnn, evaluate')

    args = parser.parse_args()

    if args.mode == "srcnn":
        train_x, train_y = data_create.save_frame(train_file_path,   #Path of the file containing the image to be cropped
                                                cut_traindata_num,  #Number of datasets generated
                                                train_height, #Storage size
                                                train_width,
                                                mag)   #magnification
                                                
        model = model.SRCNN() 
        model.compile(loss = "mean_squared_error",
                        optimizer = opt,
                        metrics = [psnr])
#https://keras.io/ja/getting-started/faq/
        model.fit(train_x,
                    train_y,
                    epochs = EPOCHS)

        model.save("srcnn_model.h5")

    elif args.mode == "evaluate":
        path = "srcnn_model"
        exp = ".h5"
        new_model = tf.keras.models.load_model(path + exp, custom_objects={'psnr':psnr})

        new_model.summary()

        test_x, test_y = data_create.save_frame(test_file_path,   #Path of the file containing the image to be cropped
                                                cut_testdata_num,  #Number of datasets generated
                                                test_height, #Storage size
                                                test_width,
                                                mag)   #magnification

        pred = new_model.predict(test_x)
        path = "resurt_" + path
        os.makedirs(path, exist_ok = True)
        path = path + "/"

        ps = psnr(tf.reshape(test_y[0], [test_height, test_width, 1]), pred[0])
        print("psnr:{}".format(ps))

        before_res = tf.keras.preprocessing.image.array_to_img(tf.reshape(test_x[0], [test_height, test_width, 1]))
        change_res = tf.keras.preprocessing.image.array_to_img(tf.reshape(test_y[0], [test_height, test_width, 1]))
        y_pred = tf.keras.preprocessing.image.array_to_img(pred[0])

        before_res.save(path + "low_" + str(0) + ".jpg ")
        change_res.save(path + "high_" + str(0) + ".jpg ")
        y_pred.save(path + "pred_" + str(0) + ".jpg ")

    else:
        raise Exception("Unknow --mode")

The main is quite long, but my impression is that if I can shorten it, I can do more. Below, I will explain the contents.

import model
import data_create
import argparse
import os
import cv2

import numpy as np
import tensorflow as tf

Here we are loading a function or another file in the same directory. datacreate.py, model.py and main.py should be in the same directory.

    def psnr(y_true, y_pred):
        return tf.image.psnr(y_true, y_pred, 1, name=None)

This time, I used psnr as a criterion for judging the quality of the generated image, so that is the definition. psnr is called the peak signal-to-noise ratio, and in simple terms it is like calculating the difference between the pixel values of the images you want to compare. I will omit the detailed explanation here, but this article is relatively detailed, and multiple evaluation methods are described.

    train_height = 33 #size of train data
    train_width = 33
    test_height = 700 #test data size
    test_width = 700

    mag = 3.0 #I don't use it, but I included it in the function.
    cut_traindata_num = 10 #How many photos are generated from one photo in train data generation.
    cut_testdata_num = 1 #How many photos are generated from one photo in test data generation?

    train_file_path = "./train_data" #Path of the file containing the image to be cropped
    test_file_path = "./test_data"

    BATSH_SIZE = 240 #batchsize
    EPOCHS = 1000 #epoch number
    opt = tf.keras.optimizers.Adam(learning_rate=0.0001) #optimizer

Here, the value used this time is set. It's okay if you look at github separately as config.py, but since it is not a large-scale program, it is summarized.

As for the size of the training data, the train data was adopted because the paper stated that it was 33 * 33. The test is just oversized for easy viewing. The number of data is 10 times the number of images contained in the file. (If 800 sheets, the number of data is 8,000)

This time, I used DIV2K Dataset, which is often used for super-resolution, for the data. Since the quality of the data is good, it is said that a certain amount of accuracy can be obtained with a small amount of data.

    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', type=str, default='srcnn', help='srcnn, evaluate')

    args = parser.parse_args()

I wanted to separate the learning and evaluation of the model here, so I made it like this so that I can select it with --mode. I will not explain in detail, so I will post the official python documentation. https://docs.python.org/ja/3/library/argparse.html

if args.mode == "srcnn":
        train_x, train_y = data_create.save_frame(train_file_path,   #Path of the file containing the image to be cropped
                                                cut_traindata_num,  #Number of datasets generated
                                                train_height, #Storage size
                                                train_width,
                                                mag)   #magnification
                                                
        model = model.SRCNN() 
        model.compile(loss = "mean_squared_error",
                        optimizer = opt,
                        metrics = [psnr])
#https://keras.io/ja/getting-started/faq/
        model.fit(train_x,
                    train_y,
                    epochs = EPOCHS)

        model.save("srcnn_model.h5")

I am learning here. If you select srcnn (the method will be described later), this program will work.

In data_create.save_frame, the function called save_frame of data_create.py is read and made available. Now that the data is in train_x and train_y, load the model in the same way and compile and fit.

See keras documentation for more information on compile and more. We use the same papers as the papers.

Finally, save the model and you're done.

elif args.mode == "evaluate":
        path = "srcnn_model"
        exp = ".h5"
        new_model = tf.keras.models.load_model(path + exp, custom_objects={'psnr':psnr})

        new_model.summary()

        test_x, test_y = data_create.save_frame(test_file_path,   #Path of the file containing the image to be cropped
                                                cut_testdata_num,  #Number of datasets generated
                                                test_height, #Storage size
                                                test_width,
                                                mag)   #magnification

        pred = new_model.predict(test_x)
        path = "resurt_" + path
        os.makedirs(path, exist_ok = True)
        path = path + "/"

        ps = psnr(tf.reshape(test_y[0], [test_height, test_width, 1]), pred[0])
        print("psnr:{}".format(ps))

        before_res = tf.keras.preprocessing.image.array_to_img(tf.reshape(test_x[0], [test_height, test_width, 1]))
        change_res = tf.keras.preprocessing.image.array_to_img(tf.reshape(test_y[0], [test_height, test_width, 1]))
        y_pred = tf.keras.preprocessing.image.array_to_img(pred[0])

        before_res.save(path + "low_" + str(0) + ".jpg ")
        change_res.save(path + "high_" + str(0) + ".jpg ")
        y_pred.save(path + "pred_" + str(0) + ".jpg ")

    else:
        raise Exception("Unknow --mode")

It is finally the explanation of the last. First, load the model you saved earlier so that you can use psnr. Next, generate a dataset for test and generate an image with predict.

I wanted to know the psnr value on the spot, so I calculated it. I wanted to save the image, so I converted it from a tensor to a numpy array, saved it, and finally it's done! This is a high quality image. (The original image) This is a low quality image with a Gaussian filter. This is the image generated by the model. It's getting pretty. Was the psnr value about 34?

4. At the end

It's been a long time since I divided it into three articles, but thank you for reading. I will continue to work on various things in the future. If you have any questions or comments, please do not hesitate to contact us!

[PYTHON] I tried to build a super-resolution method / SRCNN ③