[논문 구현] SRCNN(by Keras)을 활용한 Super Resolution

저번 시간에는 SRCNN 논문을 리뷰해 봤습니다. 이번 딥러닝 프레임 워크 중, Keras를 활용하여 SRCNN을 구현해 보도록 하겠습니다. 이 코드는 앞에서 제가 작성한 논문 리뷰 내용을 기반으로 작성했기 때문에 코드에 대해 더 잘 이해하고 싶으시다면, SRCNN 논문 리뷰 글을 먼저 읽고 오시는 것을 권장드립니다. :)

[논문리뷰] SRCNN: Image Super-Resolution Using Deep Convolutional Networks

[논문리뷰] SRCNN: Image Super-Resolution Using Deep Convolutional Networks

 

 

---

목차

모델의 구현 조건

SRCNN Code 구현

    > 필요한 라이브러리 선언

    > 파라미터 선언

    > 구글 드라이브 연동

    > 데이터셋 Load 및 확인

    > 저해상도 이미지 만들기

    > Crop된 (저해상도, 고해상도) Paired 이미지 확인

    > Model 선언 및 파라미터 초기화

    > Model Compile 및 Fit

    > 성능 확인을 위한 Test Model 선언

    > 이미지 비교

참고 문헌

---

 

 

모델의 구현 조건 (코드는 Google Colab 기반으로 작성했습니다.)

 

1. 논문 Conclusion에서 Color 이미지를 YCrCb로 변환하고 Y채널만 활용하는 방식이 아니라 RGB Channel을 그대로 활용해도 성능이 좋다고 되어 있습니다. 해당 결과를 참고하여 RGB Channel을 그대로 입력받는 모델을 구현합니다.

2. 논문의 Filter Numbers에 대한 실험 결과를 참고하여 n1 = 128, n2 = 64, n3=3으로 설정합니다. 이때 n은 각 Convolution Layer의 Filter 숫자입니다.

3. 논문의 Filter Size에 대한 실험 결과를 참고하여 f1 = 9, f2 = 3, f3= 5 구조를 사용합니다. 이때 f는 각 Convolution Layer의 Filter Size입니다.

4. 첫 번째, 두 번째 Convolution Layer의 활성함수는 ReLU, 마지막 Layer의 활성 함수는 Linear로 설정합니다.

5. 논문 조건과 유사하게 Weights는 Xaiver 정규 분포로, Bias는 Zero Vector(0)으로 초기화합니다.

6. 학습률(Learning Rate)은 구현 편의성을 고려하여 0.003으로 일괄 통일합니다.

7. SRCNN에 사용할 저해상도 이미지는 논문의 방법과 동일하게 고해상도 이미지를 Random Crop 후 Bicubic 보간법에 의해 Upscaling 된 이미지들을 활용합니다.

8. 고해상도 이미지는 논문과 동일하게 91 Images를 활용했으며, Kaggle에서 직접 다운로드하였습니다. 검증 데이터도 논문과 동일하게 set5를 사용했으며 마찬가지로 Kaggel에서 직접 다운로드했습니다. 데이터는 아래 링크 참고해 주세요.

https://www.kaggle.com/datasets/ll01dm/t91-image-dataset

T91 Image Dataset

https://www.kaggle.com/datasets/ll01dm/set-5-14-super-resolution-dataset

Set 5 & 14 Super Resolution Dataset

 

 

 

SRCNN Code 구현

 

필요한 라이브러리 선언

import numpy as np
import matplotlib.pyplot as plt
#import h5py
import tensorflow as tf
from keras.models import Sequential
from keras.layers import Conv2D, Input
from keras import initializers
from keras.optimizers import Adam
import random

import cv2
import os
from glob import glob

 

파라미터 선언

각각의 사용 환경에 맞춰서 변경해야 하는 부분은 path와 save_path입니다. path는 모델 Training에 사용하는 91 Images Dataset이 저장된 위치입니다. 참고로 아래 경로는 제가 Kaggle에서 직접 다운로드하여서 Google Drive에 저장한 경로입니다. save_path는 최종 모델의 가중치를 저장하는 경로입니다. 각자 환경에 맞춰서 변경하시면 됩니다.

n1,n2,n3 = 128, 64, 3
f1,f2,f3 = 9,3,5
upscale_factor = 3
#To avoid border effects during training, all the convolutional layers have no padding, and the network
#produces a smaller output ((fsub − f1 − f2 − f3 + 3)2 × c).

input_size = 33
output_size = input_size - f1 - f2 - f3 + 3
pad = abs(input_size - output_size) // 2 # 7
stride = 14

batch_size = 128
epochs = 200

path = "/content/drive/MyDrive/Blog/초해상도/SRCNN/T91"
save_path = "/content/drive/MyDrive/Blog/초해상도/SRCNN/SRCNN_200EPOCHS.h5"

 

구글 드라이브 연동

Google Colab에 Drive를 연동하는 코드입니다. Colab을 사용하지 않으시다면 아래 코드는 무시해 주세요.

from google.colab import drive
drive.mount('/content/drive')

 

데이터셋 Load 및 확인

glob은 데이터셋이 저장된 path에서 '.png'로 끝나는 파일의 목록을 반환합니다. 참고로 Colab이 아니라 Local PC에서 활용하시는 분들은 파일 경로를 표시할 때 '/'가 아니라 '역슬래시(\)'를 1개 또는 2개 사용해야 할 수도 있습니다.

img_paths = glob(path + '/' + '*.png')

img = cv2.imread(img_paths[0], cv2.COLOR_BGR2RGB)
print(img.shape)
plt.imshow(img)

# img_paths에 저장된 값은 아래와 같은 구조입니다.(참고용)
# /content/drive/MyDrive/Blog/초해상도/SRCNN/T91/t16.png' 등등

 

T91 Datasets의 이미지 중 하나

 

 

저해상도 이미지 만들기

SRCNN의 학습은 저해상도 이미지와 고해상도 이미지를 각각 Sub Images로 Crop 하고 이들 간의 관계를 학습합니다. 따라서 고해상도 이미지를 같은 배율의 저해상도 이미지로 만드는 작업과 각각의 이미지를 Crop 하여 맵핑하는 작업이 필요합니다.

"Zoom_img" 부분은 고해상도 이미지를 저해상도 이미지로 만드는 코드입니다. 고해상도 이미지 label을 upscale_factor만큼(여기서는 3) 줄이고, 다시 Bicubic 보간법으로 크기를 키워주면 한 장의 저해상도 이미지가 완성됩니다. (개념적으로 쉽게 이해하시려면 특정 이미지의 한 부분을 크게 확대하면 픽셀이 깨지는 것과 동일한 원리라고 생각하시면 됩니다.)

"Crop: img to sub_imgs" 부분은 저해상도 이미지와 고해상도 이미지를 각각 Crop 하여 Input & Label 관계로 맵핑하는 코드입니다. SRCNN은 Padding을 사용하지 않기 때문에 입력 이미지의 크기보다 출력 이미지의 크기가 작습니다. 이는 즉, 다시 말해 저해상도 이미지를 축소 복원하여 고해상도 이미지를 만드는 개념으로 동작한다는 의미입니다.

따라서 저해상도 이미지와 고해상도 이미지의 Crop 크기가 달라야 하는데요. 아래 코드에서 Input_size는 저해상도 이미지를 Crop 하는 크기를 의미합니다. 33으로 논문과 동일하게 설정했습니다. Output_size는 입력되는 저해상도 Sub Image가 SRCNN을 통과할 경우 최종 출력되는 Feature Map의 크기로 논문의 계산식 "33-9-5+3"에 의해 19로 지정합니다.

이미지를 Crop 할 때 Stride는 논문과 동일하게 14로 설정했습니다. 결과적으로 1장의 저해상도 이미지는 14씩 건너뛰면서 33 by 33의 정사각형 저해상도 Sub Images로 Crop 됩니다. 1장의 고해상도 이미지는 14씩 건너뛰면서 19 by 19의 정사각형 저해상도 Sub Images로 Crop 됩니다. 

sub_lr_imgs = []
sub_hr_imgs = []

for img_path in img_paths:
  img = cv2.imread(img_path, cv2.COLOR_BGR2RGB)
  
  # mod_crop
  h = img.shape[0] - np.mod(img.shape[0], upscale_factor)
  w = img.shape[1] - np.mod(img.shape[1], upscale_factor)
  img = img[:h, :w, :]

  # zoom_img
  label = img.astype('float') / 255
  temp_input = cv2.resize(label, dsize=(0,0), fx = 1/upscale_factor, fy = 1/upscale_factor,
                          interpolation = cv2.INTER_AREA)
  input = cv2.resize(temp_input, dsize=(0,0), fx = upscale_factor, fy = upscale_factor,
                     interpolation = cv2.INTER_CUBIC)
  
  # Crop: img to sub_imgs
  for h in range(0, input.shape[0] - input_size + 1, stride):
    for w in range(0, input.shape[1] - input_size + 1, stride):
      sub_lr_img = input[h:h+input_size, w:w+input_size, :]
      sub_hr_img = label[h+pad:h+pad+output_size, w+pad:w+pad+output_size, :]

      sub_lr_imgs.append(sub_lr_img)
      sub_hr_imgs.append(sub_hr_img)

sub_lr_imgs = np.asarray(sub_lr_imgs)
sub_hr_imgs = np.asarray(sub_hr_imgs)

 

Crop된 (저해상도, 고해상도) Paired 이미지 확인

fig, axes = plt.subplots(1,2, figsize = (5,5))
idx = random.randint(0, sub_lr_imgs.shape[0])

axes[0].imshow(sub_lr_imgs[idx])
axes[1].imshow(sub_hr_imgs[idx])

print(idx)
axes[0].set_title('lr_img')
axes[1].set_title('hr_img')

 

Crop된 저해상도, 고해상도 Sub Images

 

 

Model 선언 및 Parameters 초기화

GloroNormal은 Xaiver 정규 분포를 의미합니다.

initializer = initializers.GlorotNormal()

SRCNN = Sequential()
#SRCNN.add(Input(shape = (33, 33, 3)))
SRCNN.add(Conv2D(filters = n1, kernel_size = f1, activation = 'ReLU', input_shape = (33,33,3),
                 kernel_initializer = initializer, bias_initializer = 'zeros', name = 'Conv1'))
SRCNN.add(Conv2D(filters = n2, kernel_size = f2, activation = 'ReLU', 
                 kernel_initializer = initializer, bias_initializer = 'zeros', name = 'Conv2'))
SRCNN.add(Conv2D(filters = n3, kernel_size = f3, activation = 'linear', 
                 kernel_initializer = initializer, bias_initializer = 'zeros', name = 'Conv3'))

print(SRCNN.summary())

 

SRCNN Model의 구조

 

 

Model Compile 및 Fit

저는 200 Epochs만 학습해 보도록 하겠습니다. 

optimizer = Adam(lr = 0.0003)

SRCNN.compile(optimizer = optimizer, loss ='MSE', metrics = ['MSE'])
SRCNN.fit(sub_lr_imgs, sub_hr_imgs, batch_size = batch_size, epochs = epochs, verbose=1, callbacks = [callbacks])

 

성능 확인을 위한 Test Model 선언

initializer = initializers.GlorotNormal()

def predict_model():
  SRCNN = Sequential()
  SRCNN.add(Conv2D(filters = n1, kernel_size = f1, activation = 'ReLU', input_shape = (None,None,3),
                 kernel_initializer = initializer, bias_initializer = 'zeros', name = 'Conv1'))
  SRCNN.add(Conv2D(filters = n2, kernel_size = f2, activation = 'ReLU', 
                 kernel_initializer = initializer, bias_initializer = 'zeros', name = 'Conv2'))
  SRCNN.add(Conv2D(filters = n3, kernel_size = f3, activation = 'linear', 
                 kernel_initializer = initializer, bias_initializer = 'zeros', name = 'Conv3'))

  return SRCNN
  
SRCNN_Test = predict_model()
SRCNN_Test.load_weights(save_path)

 

이미지 비교

고해상도 이미지는 Set5의 Butterfly 이미지를 사용했습니다. Bicubic으로 Upscaling 된 이미지는 모델 학습 과정과 동일하게 처리해서 만들어 줍니다.

hr_img_path = '/content/drive/MyDrive/Blog/초해상도/SRCNN/set5_set14/Set5/butterfly.png'

hr_img = cv2.imread(hr_img_path)
hr_img = cv2.cvtColor(hr_img, cv2.COLOR_BGR2RGB)
print("img shape: {}".format(hr_img.shape))

plt.imshow(hr_img)

hr_img = hr_img.astype('float') / 255
temp_img = cv2.resize(hr_img, dsize=(0,0), fx = 1/upscale_factor, fy = 1/upscale_factor,
                interpolation = cv2.INTER_AREA)
bicubic_img = cv2.resize(temp_img, dsize=(0,0), fx = upscale_factor, fy = upscale_factor,
                interpolation = cv2.INTER_CUBIC)

input_img = bicubic_img[np.newaxis, :]
srcnn_img = SRCNN_Test.predict(input_img)

 

이제 3개의 이미지를 한눈에 비교해 보도록 하겠습니다.

fig, axes = plt.subplots(1,3, figsize = (10,5))

axes[0].imshow(hr_img)
axes[1].imshow(bicubic_img)
axes[2].imshow(np.squeeze(srcnn_img))

axes[0].set_title('hr_img')
axes[1].set_title('bicubic_img')
axes[2].set_title('srcnn_img')

 

SRCNN Model 성능 확인을 위한 Image 비교

 

 

여러분들의 눈에는 어떠신가요? 제 눈에는 논문에서 제시한 대로 YCrCb 채널로 굳이 바꿀 필요 없이 RGB 이미지를 그대로 활용해도 Bicubic에 비하면 성능이 훨씬 좋아 보이네요. 다만 아직 원본 고해상도 이미지에는 미치지 못하는 것 같습니다. 다음 시간에는 동일한 내용을 Pytorch로 구현해 보도록 하겠습니다.

 

 

참고로 저는 Study용으로 작성해서 모델 성능 평가 과정에서 PSNR 계산을 따로 하지는 않았는데요. 필요하신 분들은 아래 코드 참고하시면 도움이 될 것 같습니다.

from keras import backend as K

def PSNR(y_true, y_pred):
  rmse = K.mean(K.square((y_pred-y_true)**2))
  psnr = 20*tf.experimental.numpy.log10(255./rmse)
  return psnr

 

참고 문헌

https://github.com/yjn870/SRCNN-pytorch

GitHub - yjn870/SRCNN-pytorch: PyTorch implementation of Image Super-Resolution Using Deep Convolutional Networks (ECCV 2014)

https://github.com/MarkPrecursor/SRCNN-keras

GitHub - MarkPrecursor/SRCNN-keras

https://hwangtoemat.github.io/paper-review/2019-07-11-SRCNN-%EB%82%B4%EC%9A%A9/

[Super Resolution] 1-1. SRCNN 논문 리뷰