[TensorFlow] 오토 인코더를 이용한 이미지 노이즈 제거

오토 인코더를 설계하여 MNIST 데이터의 노이즈 제거를 해보겠습니다.

 

import numpy as np
import tensorflow as tf
from tensorflow.keras import datasets
import matplotlib.pyplot as plt

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import MaxPool2D
from tensorflow.keras.layers import UpSampling2D

# 랜덤 시드 설정
np.random.seed(0)
tf.random.set_seed(0)

# 데이터 불러오기
(X_tn0, y_tn0), (X_te0, y_te0) = datasets.mnist.load_data()

# 원본 데이터 차원 확인
print(X_tn0.shape)
print(y_tn0.shape)
print(X_te0.shape)
print(y_te0.shape)

# 원본 데이터 시각화
plt.figure(figsize=(10, 5))
for i in range(2*5):
    plt.subplot(2, 5, i+1)
    plt.imshow(X_tn0[i].reshape((28, 28)),
               cmap='Greys')
plt.show()

# 피쳐 데이터 스케일 조정
X_tn_re = X_tn0.reshape(60000, 28, 28, 1)
X_tn = X_tn_re/255
print(X_tn.shape)

X_te_re = X_te0.reshape(10000, 28, 28, 1)
X_te = X_te_re/255
print(X_te.shape)

# 노이즈 피쳐 데이터
X_tn_noise = X_tn + np.random.uniform(-1, 1, size=X_tn.shape)
X_te_noise = X_te + np.random.uniform(-1, 1, size=X_te.shape)

# 노이즈 데이터 스케일링
X_tn_ns = np.clip(X_tn_noise, a_min=0, a_max=1)
X_te_ns = np.clip(X_te_noise, a_min=0, a_max=1)

# 노이즈 데이터 시각화
plt.figure(figsize=(10, 5))
for i in range(2*5):
    plt.subplot(2, 5, i+1)
    plt.imshow(X_tn_ns[i].reshape((28, 28)), cmap='Greys')
plt.show()

# 오토인코더
model = Sequential()
# 인코딩
model.add(Conv2D(20, kernel_size=(5, 5),
                 input_shape=(28, 28, 1),
                 padding='same',
                 activation='relu'))
model.add(MaxPool2D(pool_size=2,
                    padding='same'))
# 디코딩
model.add(Conv2D(10, kernel_size=(5, 5),
                 padding='same',
                 activation='relu'))
model.add(UpSampling2D())
model.add(Conv2D(1, kernel_size=(5, 5),
                 padding='same',
                 activation='relu'))

# 모형 컴파일
model.compile(loss='mean_squared_error',
              optimizer='adam',
              metrics=['mean_squared_error'])

# 학습
hist = model.fit(X_tn_ns, X_tn,
                 epochs=1,
                 batch_size=100)

# 예측값
X_pred = model.predict(X_tn_ns)

# 오토인코딩 데이터 시각화
plt.figure(figsize=(10, 5))
for i in range(2*5):
    plt.subplot(2, 5, i+1)
    plt.imshow(X_pred[i].reshape((28, 28)),
               cmap='Greys')
plt.show()

 

먼저 MNIST 원본 데이터는 다음과 같습니다.

 

Conv2d 기반 오토 인코더를 사용하기 위해서는 입력 이미지 데이터의 차원(shape)이 4개로 표현되어야 합니다.

shape = (이미지 개수, 가로, 세로, 채널 수)

 

또한 0~255로 표현되어 있는 이미지 데이터를 0~1 값을 가지는 데이터로 스케일링 합니다.

스케일링을 하는 이유는 데이터의 편향성을 제거할 수 있고, 학습 중 0에 수렴을 예방하거나 발산을 억제할 수 있습니다.

 

다음으로, 노이즈 데이터를 랜덤으로 생성하고 앞과 같이 스케일링을 해줍니다.

노이즈 데이터를 시각화하면 다음과 같습니다.

 

모델 설계에서는 MaxPool2D()를 이용해 인코딩하고, UpSampling2D()를 이용해 디코딩합니다.

위에 노이즈 데이터를 입력으로 모델이 예측한 노이즈가 제거된 데이터는 아래와 같습니다.

 

원본 데이터보다는 잘 보이지 않지만 좋은 결과를 얻을 수 있습니다.