[컴퓨터비전/실습] Super Resolution(VDSR) 실습

사용한 Super Resolution Deep Learning 모델

VDSR(Very Deep Super Resolution)

공식자료 : https://cv.snu.ac.kr/research/VDSR/

참고리뷰 : https://mole-starseeker.tistory.com/83

[논문 리뷰] Super Resolution - VDSR (CVPR 2016)

SRCNN(SOTA 모델) 한계점→ 해결

1. 조그만 이미지 영역의 문맥에 의존 → CNN 모델 구조를 깊게 만듦
2. 훈련이 너무 천천히 수렴됨(그래디언트 소실/폭주 문제) → residual learning & gradient clipping
3. 단일 scale(3)에 대해서만 네트워크 작동 →multi scale factor(2, 3, 4) 처리 가능

 

실습

실습 코드(git clone 필수) : https://github.com/Lornatang/VDSR-PyTorch

GitHub - Lornatang/VDSR-PyTorch: PyTorch implements "Accurate Image Super-Resolution Using Very Deep Convolutional Networks"

 

코드 수정 위해 참고한 자료 : https://kr.mathworks.com/help/images/single-image-super-resolution-using-deep-learning.html

딥러닝을 사용하여 영상 해상도 높이기 - MATLAB & Simulink Example - MathWorks 한국

 

 

Train

 

1. dataset 이미지 구하기

dataset으로 쓸 high resolution 이미지(본인은 KITTI data_object_image_2) 5~10장을 data 경로에 추가

2. scripts/run.py 파일을 이용하여 dataset 준비

## scripts/prepare_dataset.py ##
import os

# Prepare dataset
os.system("python ./prepare_dataset.py --images_dir ../data/KITTI/original --output_dir ../data/KITTI/VDSR/train --image_size 42 --step 42 --scale 2 --num_workers 10")
os.system("python ./prepare_dataset.py --images_dir ../data/KITTI/original --output_dir ../data/KITTI/VDSR/train --image_size 42 --step 42 --scale 3 --num_workers 10")
os.system("python ./prepare_dataset.py --images_dir ../data/KITTI/original --output_dir ../data/KITTI/VDSR/train --image_size 44 --step 44 --scale 4 --num_workers 10")

# Split train and valid
os.system("python ./split_train_valid_dataset.py --train_images_dir ../data/KITTI/VDSR/train --valid_images_dir ../data/KITTI/VDSR/valid --valid_samples_ratio 0.1")

VDSR-PyTorch$ python scripts/run.py
Prepare split image: 100%|███████████████████| 41/41 [00:06<00:00,  6.26image/s]
Prepare split image: 100%|███████████████████| 41/41 [00:06<00:00,  6.67image/s]
Prepare split image: 100%|███████████████████| 41/41 [00:05<00:00,  7.11image/s]
Split train/valid dataset: 100%|██████| 2820/2820 [00:00<00:00, 13806.76image/s]

 

1. Prepare dataset

prepare_dataset.py를 통해 원본 이미지를 정해진 해상도(42*42 or 44*44, 본인은 건들지 않음)만큼 조각내어 이미지 수를 늘리고, high resolution, low resolution 이미지 생성 및 폴더에 각각 정리해줌 (경로는 본인에 맞게 수정)

 

crop hr (42*42)

crop hr -> lr (42*42)

2. Split train and valid

split_train_valid_dataset.py를 통해 train, valid dataset으로 split해줌

3. 훈련 진행

## config.py ##
# line 30: upscale_factor 변경(2, 3, 4)
upscale_factor = 2

# line 32: mode 변경 (훈련 시 train, 테스트 시 valid)
mode = "train"

# line 38, 39, 40 / line 71: dataset 경로 변경
train_image_dir = "data/KITTI/train"
valid_image_dir = "data/KITTI/valid"
test_image_dir = "data/KITTI/test"

# 그 외 훈련 parameter 수정하면서 custom(본인은 하지 않음)
...
if mode == "valid":
    # Test data address
    sr_dir = f"results/test/{exp_name}"
    hr_dir = f"data/test"

 

- config.py에서 수정해야할 사항: upscale_factor = 2 or 3 or 4, mode = “train”

- tensorboard 설치는 처음 실행할 때만 설치

- 정확한 test 성능을 확인하기 위해 train에서 사용되는 test 이미지와 valid 시 사용되는 test 이미지 데이터를 다르게 함

- test 이미지에는 high resolution(원본) 이미지를 사용함

 

VDSR-PyTorch$ pip install tensorboard
VDSR-PyTorch$ python train.py
2023-11-29 16:55:17.088318: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.
...
Read lr dataset into memory: 100%|██| 28208/28208 [00:02<00:00, 12529.97image/s]
Read hr dataset into memory: 100%|██| 28208/28208 [00:02<00:00, 12443.49image/s]
Read lr dataset into memory: 100%|████| 2820/2820 [00:00<00:00, 12621.19image/s]
Read hr dataset into memory: 100%|████| 2820/2820 [00:00<00:00, 12382.51image/s]
Read test dataset into memory: 100%|███████████| 6/6 [00:00<00:00,  8.63image/s]
Load train dataset and valid dataset successfully.
Build VDSR model successfully.
Define all loss functions successfully.
Define all optimizer functions successfully.
Define all optimizer scheduler successfully.
Check whether the pretrained model is restored...
Epoch: [1][   0/1763]	Time  0.512 ( 0.512)	Data  0.007 ( 0.007)	Loss 13404.875000 (13404.875000)	PSNR 3.02 (3.02)
Epoch: [1][ 200/1763]	Time  0.007 ( 0.009)	Data  0.000 ( 0.000)	Loss 51.034664 (783.978808)	PSNR 27.22 (28.19)
Epoch: [1][ 400/1763]	Time  0.006 ( 0.008)	Data  0.000 ( 0.000)	Loss 20.585293 (404.789111)	PSNR 31.16 (29.53)
Epoch: [1][ 600/1763]	Time  0.007 ( 0.007)	Data  0.000 ( 0.000)	Loss 13.845083 (277.809261)	PSNR 32.88 (30.00)
Epoch: [1][ 800/1763]	Time  0.005 ( 0.007)	Data  0.000 ( 0.000)	Loss 13.586741 (214.185162)	PSNR 32.97 (30.26)
Epoch: [1][1000/1763]	Time  0.006 ( 0.007)	Data  0.000 ( 0.000)	Loss 25.427002 (175.798263)	PSNR 30.24 (30.45)
Epoch: [1][1200/1763]	Time  0.006 ( 0.007)	Data  0.000 ( 0.000)	Loss 20.634468 (150.000494)	PSNR 31.15 (30.62)
Epoch: [1][1400/1763]	Time  0.006 ( 0.007)	Data  0.000 ( 0.000)	Loss 23.769402 (131.790329)	PSNR 30.54 (30.69)
Epoch: [1][1600/1763]	Time  0.006 ( 0.007)	Data  0.000 ( 0.000)	Loss 8.729490 (117.887938)	PSNR 34.89 (30.80)
Valid: [  0/177]	Time  0.007 ( 0.007)	PSNR 27.36 (27.36)
 *  PSNR 31.30
Test: [0/6]	Time  0.096 ( 0.096)	PSNR 38.35 (38.35)
 *  PSNR 35.65
...
Epoch: [79][   0/1763]	Time  0.011 ( 0.011)	Data  0.005 ( 0.005)	Loss 22.031227 (22.031227)	PSNR 30.87 (30.87)
Epoch: [79][ 200/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 15.408798 (14.234085)	PSNR 32.42 (33.27)
Epoch: [79][ 400/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 7.400876 (15.074915)	PSNR 35.60 (33.04)
Epoch: [79][ 600/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 10.699713 (15.141599)	PSNR 34.00 (33.02)
Epoch: [79][ 800/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 21.873383 (14.955173)	PSNR 30.90 (33.04)
Epoch: [79][1000/1763]	Time  0.007 ( 0.006)	Data  0.000 ( 0.000)	Loss 19.263966 (15.016066)	PSNR 31.45 (33.00)
Epoch: [79][1200/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 9.317989 (14.819112)	PSNR 34.60 (33.07)
Epoch: [79][1400/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 17.098080 (14.958444)	PSNR 31.97 (33.02)
Epoch: [79][1600/1763]	Time  0.007 ( 0.006)	Data  0.000 ( 0.000)	Loss 7.460383 (14.863348)	PSNR 35.57 (33.04)
Valid: [  0/177]	Time  0.010 ( 0.010)	PSNR 28.46 (28.46)
 *  PSNR 33.01
Test: [0/6]	Time  0.041 ( 0.041)	PSNR 41.08 (41.08)
 *  PSNR 37.83


Epoch: [80][   0/1763]	Time  0.015 ( 0.015)	Data  0.006 ( 0.006)	Loss 11.411203 (11.411203)	PSNR 33.72 (33.72)
Epoch: [80][ 200/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 5.546063 (14.348182)	PSNR 36.86 (33.13)
Epoch: [80][ 400/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 21.118137 (14.629654)	PSNR 31.05 (33.07)
Epoch: [80][ 600/1763]	Time  0.007 ( 0.006)	Data  0.000 ( 0.000)	Loss 29.585289 (14.733129)	PSNR 29.59 (33.12)
Epoch: [80][ 800/1763]	Time  0.005 ( 0.006)	Data  0.000 ( 0.000)	Loss 18.407265 (14.824776)	PSNR 31.65 (33.08)
Epoch: [80][1000/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 10.541956 (14.634793)	PSNR 34.07 (33.14)
Epoch: [80][1200/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 10.815474 (14.571670)	PSNR 33.96 (33.16)
Epoch: [80][1400/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 9.329271 (14.746994)	PSNR 34.60 (33.11)
Epoch: [80][1600/1763]	Time  0.006 ( 0.006)	Data  0.000 ( 0.000)	Loss 35.551880 (14.777465)	PSNR 28.79 (33.10)
Valid: [  0/177]	Time  0.006 ( 0.006)	PSNR 28.52 (28.52)
 *  PSNR 33.01
Test: [0/6]	Time  0.042 ( 0.042)	PSNR 41.06 (41.06)
 *  PSNR 37.85

 

- results 폴더에 VDSR 훈련 모델 생성됨 확인(vdsr_baseline)
- 훈련 약 10~20분 정도 소요

 

훈련시 발생한 에러: RuntimeError tensor size

Epoch: [1][  0/561]	Time  0.436 ( 0.436)	Data  0.006 ( 0.006)	Loss 6990.436523 (6990.436523)	PSNR 5.85 (5.85)
Epoch: [1][200/561]	Time  0.007 ( 0.009)	Data  0.000 ( 0.000)	Loss 3934.592285 (4318.330361)	PSNR 8.35 (8.27)
Epoch: [1][400/561]	Time  0.007 ( 0.008)	Data  0.000 ( 0.000)	Loss 3170.051270 (3971.348025)	PSNR 9.29 (8.56)
Valid: [  0/188]	Time  0.009 ( 0.009)	PSNR 8.78 (8.78)
 *  PSNR 9.22
...
Traceback (most recent call last):
  File "train.py", line 369, in <module>
    main()
  File "train.py", line 89, in main
    psnr = validate(model, test_prefetcher, psnr_criterion, epoch, writer, "Test")
...
RuntimeError: The size of tensor a (1244) must match the size of tensor b (1242) at non-singleton dimension 3

원인

lr에 적용되는 upscale_factor를 hr의 image 크기를 벗어나면 나는 에러였다.

original image 사이즈가 1242*375라는 나누기 애매한 사이즈를 갖고있기 때문이었다.

• /4로 해상도를 줄였다가 *4로 키울 때, 1242, 375 모두 나누어떨어지지 않기 때문에 두 값이 달라 줄였다 키운 target 이미지와 원본 input 이미지의 사이즈가 달라지는 문제였다.(비교 불가)

해결

original image를 나누기 좋은 1232*372 정도로 resize(2, 4 모두 나누어 떨어짐)

• 이미지 사이즈가 12로 나누어떨어지도록 설정하는 것도 좋을 듯(upscale factor 2, 3, 4 모두 가능)

 

Test

 

!config.py 에서 mode="valid'으로 변경!

VDSR-PyTorch$ python validate.py

high resolution(input)

super resolution(output)

low resolution(비교용)

super resolution 일부

low resolution 일부

 

 

위 코드는 high resolution 이미지를 test input으로 사용하고, 코드 내에서 이를 저해상도 이미지로 만든 후
super resolution 진행한다.
하지만 실생활에서는 low resolution 이미지를 input으로 하고, 이를 super resolution할 경우가 많다고 생각하므로, 이를 실현하기 위해 코드를 수정하고자 한다.

 

코드 수정(validate.py)

    for index in range(total_files):
        sr_image_path = os.path.join(config.sr_dir, file_names[index])
        # hr_image_path : low resolution test input image 들어있음
        hr_image_path = os.path.join(config.hr_dir, file_names[index])
        print(f"Processing `{os.path.abspath(hr_image_path)}`...")

        # Read low-resolution image
        lr_image = cv2.imread(hr_image_path).astype(np.float32) / 255.0
        lr_image_height, lr_image_width = lr_image.shape[:2]
        lr_image_height_remainder = lr_image_height % 12
        lr_image_width_remainder = lr_image_width % 12
        lr_image = lr_image[:lr_image_height - lr_image_height_remainder, :lr_image_width - lr_image_width_remainder, ...]

        # Convert BGR image to YCbCr image
        lr_ycbrc = imgproc.bgr2ycbcr(lr_image, use_y_channel=False)

        # Split YCbCr image data
        lr_y, lr_cb, lr_cr = cv2.split(lr_ycbrc)

        # upsampling
        y_bicubic = cv2.resize(lr_y, (lr_image_width*config.upscale_factor, lr_image_height*config.upscale_factor), interpolation=cv2.INTER_CUBIC)
        cb_bicubic = cv2.resize(lr_cb, (lr_image_width*config.upscale_factor, lr_image_height*config.upscale_factor), interpolation=cv2.INTER_CUBIC)
        cr_bicubic = cv2.resize(lr_cr, (lr_image_width*config.upscale_factor, lr_image_height*config.upscale_factor), interpolation=cv2.INTER_CUBIC)

        # Convert Y image data convert to Y tensor data
        lr_y_tensor = imgproc.image2tensor(y_bicubic, range_norm=False, half=True).to(config.device).unsqueeze_(0)

        # Only reconstruct the Y channel image data.
        with torch.no_grad():
            sr_y_tensor = model(lr_y_tensor).clamp_(0, 1.0)
            
        # Save image
        sr_y_image = imgproc.tensor2image(sr_y_tensor, range_norm=False, half=True)
        sr_y_image = sr_y_image.astype(np.float32) / 255.0
        sr_ycbcr_image = cv2.merge([sr_y_image, cb_bicubic, cr_bicubic])
        sr_image = imgproc.ycbcr2bgr(sr_ycbcr_image)
        cv2.imwrite(sr_image_path, sr_image * 255.0)

 

원본 코드의 경우, high resolution의 y_tensor가 Save image 부분에서 cv2.merge 시 YCbCr 중 cb와 cr 값이 필요하다.

하지만, 이를 low resolution 이미지를 bicubic upsampling 방법을 이용하여 resolution을 high resolution 값에 맞게 높여주고, 이에 대한 cb, cr 값을 sr_y_image에 merge시켜주었다.

 

Result

high resolution(비교용)

super resolution(output)

low resolution(input)