[PYTHON] Super-resolution technology-SRCNN-Implemented (Tensorflow 2.0) Prediction phase

Overview

Continuing from Last time, this is the prediction phase of SRCNN. Actually, I tried to explain the code, but I didn't explain it that much. .. .. So this time I will omit the explanation of the code. It ends after showing the result of super-resolution. .. ..

Postscript

After receiving advice in the comments and partially modifying the code [Caution (solved)](https://qiita.com/hima_zin331/items/ebb6046a2a8d860254e1#%E6%B3%A8%E6%84%8F%E8 Problems such as% A7% A3% E6% B1% BA% E6% B8% 88% E3% 81% BF) have been resolved. The image of the execution result has also been replaced. Similarly, the programs in the Github repository have been reworked.

environment

-Software- Windows 10 Home Anaconda3 64-bit(Python3.7) Spyder -Library- Tensorflow 2.1.0 opencv-python 4.1.2.30 -Hardware- CPU: Intel core i9 9900K GPU: NVIDIA GeForce RTX2080ti RAM: 16GB 3200MHz

program

I will post it on Github. https://github.com/himazin331/Super-resolution-CNN The repository contains a learning phase, a prediction phase, and a dataset (General-100).

Source code

I will put only the source code. ** Please note that the code is dirty ... **

srcnn_pre.py

import argparse as arg
import os
import sys

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' #Hide TF message

import tensorflow as tf
import tensorflow.keras.layers as kl

import cv2
from PIL import Image
import matplotlib.pyplot as plt

import numpy as np

# SRCNN
class SRCNN(tf.keras.Model):

    def __init__(self, h, w):
        super(SRCNN, self).__init__()

        self.conv1 = kl.Conv2D(64, 3, padding='same', activation='relu', input_shape=(None, h, w, 3))
        self.conv2 = kl.Conv2D(32, 3, padding='same', activation='relu')
        self.conv3 = kl.Conv2D(3, 3, padding='same', activation='relu')

    def call(self, x):
        
        h1 = self.conv1(x)
        h2 = self.conv2(h1)
        h3 = self.conv3(h2)

        return h3
    
def main():
    
    #Command line option creation
    parser = arg.ArgumentParser(description='Super-resolution CNN prediction')
    parser.add_argument('--param', '-p', type=str, default=None,
                        help='Specifying learned parameters(Error if not specified)')    
    parser.add_argument('--data_img', '-d', type=str, default=None,
                        help='Specifying an image file(Error if not specified)')
    parser.add_argument('--out', '-o', type=str, default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "result"),
                        help='Specify save destination(Default value=./result)')                        
    parser.add_argument('--he', '-he', type=int, default=256,
                        help='Resize height specification(Default value=256)')      
    parser.add_argument('--wi', '-wi', type=int, default=256,
                        help='Specify resizing(Default value=256)')
    parser.add_argument('--mag', '-m', type=int, default=2,
                        help='Specifying the reduction magnification(Default value=2)')                           
    args = parser.parse_args()

    #When parameter file is not specified->exception
    if args.param == None:
        print("\nException: Trained Parameter-File not specified.\n")
        sys.exit()
    #When a parameter file that does not exist is specified->exception
    if os.path.exists(args.param) != True:
        print("\nException: Trained Parameter-File {} is not found.\n".format(args.param))
        sys.exit()
    #When no image file is specified->exception
    if args.data_img == None:
        print("\nException: Image not specified.\n")
        sys.exit()
    #When a non-existent image file is specified->exception
    if os.path.exists(args.data_img) != True:
        print("\nException: Image {} is not found.\n".format(args.data_img))
        sys.exit()
    #When 0 is entered in either the width height or the reduction magnification->exception
    if args.he == 0 or args.wi == 0 or args.mag == 0:
        print("\nException: Invalid value has been entered.\n")
        sys.exit()

    #Setting information output
    print("=== Setting information ===")
    print("# Trained Prameter-File: {}".format(os.path.abspath(args.param)))
    print("# Image: {}".format(args.data_img))
    print("# Output folder: {}".format(args.out))
    print("")
    print("# Height: {}".format(args.he))
    print("# Width: {}".format(args.wi))
    print("# Magnification: {}".format(args.mag))
    print("===========================")

    #Create output folder(Do not create if the folder exists)
    os.makedirs(args.out, exist_ok=True)

    #Model building
    model = SRCNN(args.he, args.wi)
    model.build((None, args.he, args.wi, 3))
    model.load_weights(args.param)

    #Input image processing(High resolution image)
    img = cv2.imread(args.data_img)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    hr_img = cv2.resize(img, (args.he, args.wi))

    #Low resolution image creation
    lr_img = cv2.resize(hr_img, (int(args.he/args.mag), int(args.wi/args.mag)))
    lr_img = cv2.resize(lr_img, (args.he, args.wi))
    lr_img_s = lr_img

    #Normalization
    lr_img = tf.convert_to_tensor(lr_img, np.float32)
    lr_img /= 255.0
    lr_img = lr_img[np.newaxis, :, :, :]

    #Super resolution
    re = model.predict(lr_img)

    #Data processing
    re = np.reshape(re, (args.he, args.wi, 3))
    re *= 255
    re = np.clip(re, 0.0, 255.0) #clipping(0~Round to 255)

    #Low resolution image storage
    lr_img = Image.fromarray(np.uint8(lr_img_s))
    lr_img.show()
    lr_img.save(os.path.join(args.out, "Low-resolution Image(SRCNN).bmp"))

    #Super-resolution image storage
    sr_img = Image.fromarray(np.uint8(re))
    sr_img.show()
    sr_img.save(os.path.join(args.out, "Super-resolution Image(SRCNN).bmp"))

    #High resolution image storage
    hr_img = Image.fromarray(np.uint8(hr_img))
    hr_img.show()
    hr_img.save(os.path.join(args.out, "High-resolution Image(SRCNN).bmp"))

if __name__ == "__main__":
    main()

Attention (solved)

In the code

    re *= 240 #If it is 255, the output will be abnormal.

There is such a process. re is the result ofmodel.predict (), that is, a super-resolution image. On the other hand, I think the correct answer is to multiply the pixel value from 0 to 1 to 0 to 255 by multiplying it by 255, but when it is multiplied by 255, the color information is broken (?) As shown in the image below. I will. I don't know why this happens, but if anyone knows, please let me know in the comments. Since it is multiplied by 240 instead of 255, the color is a little darker than the original.

Execution result

This is the result of 3000 Epoch numbers. I used General-100 for training data.

** Result of magnification 2 **

The resolution is so high that it is almost indistinguishable from high resolution images.

** Result of magnification 4 **

If the resolution is so low, detailed information will be lost.

in conclusion

I tried to increase the resolution of low resolution images using SRCNN, a super-resolution technology. As I mentioned at the beginning of the last time, I implemented SRGAN in addition to SRCNN. I will summarize it someday.