SRGAN Custom Training (2) - Dataset Custom Setting

Dataset Custom

나의 경우 데이터의 갯수가 적고, 화질이 낮아 파일을 불러온 후 고해상도 이미지를 저해상도로 다운샘플링한 후, 저해상도 이미지를 다시 고해상도 크기로 리사이징하여 저장하는 작업을 수행했다.

import os
import glob
import tensorflow as tf

• os: 파일 시스템 경로 관련 작업을 수행하는 파이썬 내장 모듈이다.
• glob: 파일 경로 패턴 매칭을 위한 모듈이다. 특정 디렉토리에서 원하는 파일 확장자나 패턴을 통해 파일 목록을 가져올 수 있다.
• tensorflow: 텐서플로우 라이브러리로 딥러닝 및 다양한 데이터 처리를 수행한다.

 

---

이미지 로딩 함수 정의

def load_image(path):
    img = tf.io.read_file(path)
    img = tf.image.decode_png(img, channels=3)
    return img

• load_image: 지정된 경로에서 PNG 이미지를 읽어들이는 함수이다.

 

---

전처리 함수 정의 (이미지 패딩 및 리사이즈)

def preprocessing(hr, target_size=384):
    hr = tf.image.convert_image_dtype(hr, tf.float32)
    hr_padded = tf.image.resize_with_pad(hr, target_size, target_size)
    lr_padded = tf.image.resize(hr_padded, [target_size // 4, target_size // 4], method='bicubic')
    lr_resized = tf.image.resize(lr_padded, [target_size, target_size], method='bicubic')
    return lr_resized, hr_padded

• preprocessing: 고해상도 이미지를 전처리하고 저해상도 이미지를 생성하는 함수이다. 이미지 데이터를 float32 타입으로 변환하고 [0, 1] 범위로 정규화한다. 종횡비 유지 및 패딩 추가로 이미지 크기를 조정한다. 이후 고해상도 이미지를 다운샘플링하여 저해상도 이미지를 생성한 후, 저해상도 이미지를 다시 고해상도 크기로 리사이즈한다.

 

---

전처리된 이미지를 저장하는 함수

def preprocess_and_save(path, index, hr_save_dir, lr_save_dir, target_size=384):
    hr = load_image(path)
    lr_resized, hr_padded = preprocessing(hr, target_size)
    hr_padded = tf.image.convert_image_dtype(hr_padded, tf.uint8)
    lr_resized = tf.image.convert_image_dtype(lr_resized, tf.uint8)
    hr_path = os.path.join(hr_save_dir, f'hr_{index:05d}.png')
    lr_path = os.path.join(lr_save_dir, f'lr_{index:05d}.png')
    hr_encoded = tf.image.encode_png(hr_padded)
    lr_encoded = tf.image.encode_png(lr_resized)
    tf.io.write_file(hr_path, hr_encoded)
    tf.io.write_file(lr_path, lr_encoded)

• preprocess_and_save: 이미지를 전처리한 후 디스크에 저장하는 함수이다. 원본 이미지를 로드하여 전처리를 통한 저해상도 및 고해상도 이미지 쌍을 생성한다. 이미지를 uint8 타입으로 변환하고 [0, 255] 범위로 스케일링한다. 이미지 저장 경로를 설정하여 이미지를 PNG 형식으로 인코딩하여 디스크에 이미지를 저장한다.

data_list = glob.glob('_srgan_data/dataset/*.png')

• glob.glob: 주어진 경로에서 .png 파일 목록을 가져온다. 이 파일 경로 리스트는 전처리할 이미지들의 경로이다.

hr_save_dir = 'processed_data/hr_images'  
lr_save_dir = 'processed_data/lr_images' 

os.makedirs(hr_save_dir, exist_ok=True)
os.makedirs(lr_save_dir, exist_ok=True)

• os.makedirs: 지정된 디렉토리가 존재하지 않으면 새로 생성한다. 고해상도와 저해상도 이미지를 각각 저장할 디렉토리를 설정한다.

for idx, path in enumerate(data_list):
    preprocess_and_save(path, idx, hr_save_dir, lr_save_dir)
    if idx % 100 == 0:
        print(f'{idx}개의 이미지를 처리했습니다.')
print('데이터셋 생성이 완료되었습니다.')

• for idx, path in enumerate: 경로 리스트에서 이미지를 하나씩 불러와 전처리하고 저장한다. 100개 단위로 처리 상태를 출력한다. 이후 데이터 생성 완료시 출력한다.

 

---

Data Architecture

먼저, 학습하고자 하는 데이터셋을 구성한다. test가 있을 경우 8:1:1, 또는 8:2, 7:3의 비율로 진행한다.

 

데이터셋의 구조는 아래와 같다.

# 2,500EA, (train 1,000/1,000, validation 250/250)
# 1,358EA, (train 543/544, validation 136/135)
├─dataset root
│ ├─train
│ │ ├─hr_images
│ │ ├─lr_images
│ ├─valid
│ │ ├─hr_images
│ │ ├─lr_images

 

 

---

Config.py

SRGAN의 config.py를 통해 dataset path를 지정한다.

## train set location
config.TRAIN.hr_img_path = 'dataset/train/hr_images/'
config.TRAIN.lr_img_path = 'dataset/train/lr_images/'

config.VALID = edict()
## test set location
config.VALID.hr_img_path = 'dataset/valid/hr_images/'
config.VALID.lr_img_path = 'dataset/valid/lr_images/'

 

 

---

참고

https://arxiv.org/abs/1609.04802

Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network

https://github.com/tensorlayer/SRGAN

GitHub - tensorlayer/SRGAN: Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network

https://panggu15.github.io/gan/SRGAN/#srgan-%EC%84%A4%EB%AA%85

SRGAN 설명 및 코드 구현하기