[Tensorflow] 이미지/AutoEncoder/UNet - 오토인코더로 그림 인코딩, 디코딩 하기

이번 프로젝트는 오토인코더(Auto Encoder)를 활용하여 이미지를 출력하도록 모델을 학습시키는 기초적인 연습을 해보려 합니다.

08. 오토인코더 (AutoEncoder)

제가 이론을 공부할 때 참고한 블로그입니다. 잘 정리되어 있어요.

이번 프로젝트에서는 조금 새롭게 렌더링된 이미지를 사용합니다. 128x128픽셀 이미지가 아닌 조금 더 큰 사이즈의 이미지에, 128x128 이미지 때와 크기는 같은 구가 랜덤한 위치에 렌더링되어 있습니다. 해당 이미지를 128x128크기로 랜덤하게 잘라, 완전한 구 형태가 아닐 때에도 적용가능하도록 해보겠습니다.

1. openGL을 활용하여 렌더링 된 구(sphere)의 환경이미지를 100개 불러온다. (30,000개 중에서 같은 환경 당 100개만 선택) 환경이미지 예시 :

2. 이미지를 CNN을 활용하여 작고 많은 피쳐(코드)로 만들어주는 encoder, 인코딩된 코드를 다시 이미지로 deconvolution 해주는 decoder를 모델로 구성한다.
3. 1에서 불러온 이미지의 리스트를 모델의 input이자 label로 지정하여 학습시킨다.
4. 새로운 렌더링 된 구 환경이미지 32개(1,600개 중에서 선택)를 testset으로 불러와 학습된 모델로 인코딩-디코딩 한 후 원래의 이미지와 비교한다.

---

dataset hierarchy

📁 pbr_30000_env (training용 구 환경이미지)
ㄴ 구 이미지들(pbr_00000.jpg ~ pbr_29999.jpg)
📁 pbr_1600_env (test용 구 환경이미지)
ㄴ 구 이미지들(pbr_00000.jpg ~ pbr_01599.jpg)

* 해당 데이터셋은 전달/배포가 불가한 점 양해바랍니다.



이번 게시물엔 제가 프로젝트를 진행하며 만난 오류들을 해결하는 과정이 포함되어 있습니다.

전처리

train_env = []

def random_crop(image):
    cropped_image = tf.image.random_crop(image, size=[128, 128, 3])
    return cropped_image

for i in range (100):
    filename = 'nspheres\pbr_30000_env\pbr_' + str(i*300).zfill(5) + '.jpg'
    image = pilimg.open(filename)
    image = np.array(image)
    cropped_image = tf.image.random_crop(image, size=[128, 128, 3])
    train_env.insert(i, cropped_image) #concatnate (이미지 한장 먼저 넣고..)
    
train_env = np.array(train_env, dtype="float32")
train_env /= 255.0

train_env.shape

tf.image.random_crop  |  TensorFlow Core v2.4.1

위에서 언급한 바와 같이, random crop을 사용하여 이미지를 128x128 크기로 잘라 주었다. 참고할 점은, random crop을 하는 데에 시간이 꽤 오래 걸린다는 것이다.

모델 생성 첫 번째 시도 : 기본 오토인코더에 Encoder 레이어 추가

latent_dim = 64 

class Autoencoder(Model):
  def __init__(self, encoding_dim):
    super(Autoencoder, self).__init__()
    self.latent_dim = latent_dim   
    self.encoder = tf.keras.Sequential([ #여기에 convolution
      layers.Flatten(),
      layers.Dense(latent_dim, activation='relu'),
    ])
    self.decoder = tf.keras.Sequential([
      layers.Dense(784, activation='sigmoid'),
      layers.Reshape((28, 28))
    ])

  def call(self, x):
    encoded = self.encoder(x)
    decoded = self.decoder(encoded)
    return decoded

autoencoder = Autoencoder(latent_dim)

Autoencoder 소개  |  TensorFlow Core

위 코드는 TensorFlow Core의 오토인코더 소개에 제시된 가장 기본적인 형태의 오토인코더 이다. 이 오토인코더의 encoder 부분에 지난 프로젝트에서 했던 CNN 연산을 하여 이미지를 몇개의 파라미터 형태로 압축해줄 것이다. 그 다음 다시 decoder에 연산을 거꾸로 해주어 기존에 넣어주었던 이미지를 다시 복원해내도록 시도했다.

latent_dim = 64

class Autoencoder(Model):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.latent_dim = latent_dim
        self.encoder = tf.keras.Sequential([ #여기에 convolution 추가
            layers.Conv2D(32, (3,3), activation='relu', input_shape=(128,128,3)),
            layers.MaxPooling2D((2,2)),
            layers.Conv2D(64, (3,3), activation='relu'),
            layers.MaxPooling2D((2,2)),
            layers.Conv2D(128, (3,3), activation='relu'),
            layers.MaxPooling2D((2,2)),
            layers.Conv2D(256, (3,3), activation='relu'),
            layers.MaxPooling2D((2,2)),
            layers.Conv2D(512, (3,3), activation='relu'),
            layers.Flatten(),
            layers.Dense(512, activation='relu'),
            layers.Dense(latent_dim, activation='relu')
        ])
        self.decoder = tf.keras.Sequential([
            layers.Dense(512, activation='sigmoid'),
            layers.Reshape((128,128))
        ])
        
    def call(self, x):
        encoded = self.encoder(x) #x->encoded
        decoded = self.decoder(encoded) #encoded->decoded
        return decoded

일단 위 코드처럼 convolution 레이어들을 추가해주었다. 이 정도에서 끝날줄 알았던 것은 크나큰 착각이었다. 당연히 결과로는 비슷하지도 않은 그림들만 잔뜩 출력되었다.

모델 생성 두 번째 시도 : Decoder 레이어 생성

Encoder에서 convolution 연산을 해주었다면 Decoder에서도 똑같은 횟수만큼 deconvolution 연산을 해주어야 하는 것을 알게 되었다. 그리고 마찬가지로, MaxPooling 연산에 대해서도 반대 연산인 UpSampling 연산을 Decoder에 추가해주어야 한다.

tf.keras.layers.Conv2DTranspose  |  TensorFlow Core v2.4.1

[신경망] 18. Deconvolution

tf.keras.layers.UpSampling2D  |  TensorFlow Core v2.4.1

메소드 이름이 몇번 변경되었다. 오래된 자료에서는 다른 이름으로 소개되기도 하지만, Conv2DTranspose와 UpSampling2D가 이번 프로젝트에서 사용할 Conv2D와 MaxPooling2D의 반대 연산 메소드라고 생각하면 된다. (2021년 3월 기준)

Encoder 마지막 부분에 추가한 Dense 레이어들을 삭제했다. Encoder의 레이어들:

layers.Conv2D(32, (3,3), activation='relu', input_shape=(128,128,3)),
layers.MaxPooling2D((2,2))(conv1),
layers.Conv2D(64, (3,3), activation='relu'),
layers.MaxPooling2D((2,2))(conv2),
layers.Conv2D(128, (3,3), activation='relu'),
layers.MaxPooling2D((2,2))(conv3),
layers.Conv2D(256, (3,3), activation='relu'),
layers.MaxPooling2D((2,2))(conv4),
layers.Conv2D(512, (3,3), activation='relu')

위에서 첫 번째 Conv2D 레이어는 128x128x3(채널, RGB)의 이미지를 128x128x32(Conv2D 레이어의 첫번째 인자)로 바꾸고, 파라미터 개수를 늘린다. 두 번째 레이어인 MaxPooling2D는 전 레이어에서 생성한 128x128 크기의 이미지들을 모두 반으로, 64x64x32(채널)로 만들어 준다.

파라미터, 이미지크기, 채널 수에 유의하며 Decoder에 레이어를 추가해주었다. Decoder의 레이어들:

layers.Conv2DTranspose(256, (3,3), activation='relu'), #512아니고 256
layers.UpSampling2D((2,2)),
layers.Conv2DTranspose(128, (3,3), activation='relu'),
layers.UpSampling2D((2,2)),
layers.Conv2DTranspose(64, (3,3), activation='relu'),
layers.UpSampling2D((2,2)),
layers.Conv2DTranspose(32, (3,3), activation='relu'),
layers.UpSampling2D((2,2)),
layers.Conv2DTranspose(3, (3,3), activation='relu')

파라미터 수가 하나라도 맞지 않으면 오류를 낸다. 그래서 아예 새로운 ipynb 파일을 하나 만들어서, 레이어들만 놓고 모델을 컴파일 해봤다. summary를 출력하면서 계산이 맞도록 손을 봐주면 된다.

나 같은 경우에는, 이미지 크기가 홀수일 때 MaxPooling을 하고 다시 UpSampling을 하면 크기가 달라지는 문제를 해결하기 위해, 이미지의 크기를 2의 거듭제곱으로 설정하고 Conv2D 레이어와 Conv2DTranspose 레이어에 padding을 설정해 주었다.

시행착오를 겪는 중

MaxPooling2d() turns odd input into even numbers, is there anyway to go around this? · Issue #4818 · keras-team/keras

TypeError: ('Keyword argument not understood:', 'border_mode') · Issue #6765 · Theano/Theano

#Encoder
input_e = tf.keras.Input(shape=(128,128,3))

conv1 = layers.Conv2D(32, (3,3), activation='relu', padding = 'same')(input_e)
mp1 = layers.MaxPooling2D((2,2))(conv1)
        
conv2 = layers.Conv2D(64, (3,3), activation='relu', padding = 'same')(mp1)
mp2 = layers.MaxPooling2D((2,2))(conv2)
        
conv3 = layers.Conv2D(128, (3,3), activation='relu', padding = 'same')(mp2)
mp3 = layers.MaxPooling2D((2,2))(conv3)

conv4 = layers.Conv2D(256, (3,3), activation='relu', padding = 'same')(mp3)
mp4 = layers.MaxPooling2D((2,2))(conv4)
        
output_e = layers.Conv2D(512, (3,3), activation='relu', padding = 'same')(mp4)

#Decoder
convt1 = layers.Conv2DTranspose(256, (3,3), activation='relu', padding='same')(output_e) #512아니고 256
upsamp1 = layers.UpSampling2D((2,2))(convt1)

convt2 = layers.Conv2DTranspose(128, (3,3), activation='relu', padding='same')(upsamp1)
upsamp2 = layers.UpSampling2D((2,2))(convt2)

convt3 = layers.Conv2DTranspose(64, (3,3), activation='relu', padding='same')(upsamp2)
upsamp3 = layers.UpSampling2D((2,2))(convt3)

convt4 = layers.Conv2DTranspose(32, (3,3), activation='relu', padding='same')(upsamp3)
upsamp4 = layers.UpSampling2D((2,2))(convt4)

output_d = layers.Conv2DTranspose(3, (3,3), activation='relu', padding='same')(upsamp4)

위의 코드가 functional API로 만들어준 Encoder와 Decoder 레이어들이다. (튜토리얼처럼 모델을 두 개로 나눠 클래스화 했더니 graph connection 오류가 생겨서 변경했다 : stackoverflow.com/questions/51522848/graph-disconnected-cannot-obtain-value-for-tensor-tensor-input-keras-python/51523029)

여기까지 구현하면 뚜렷하게 보이지는 않지만 뭔가 학습이 된 듯한 결과를 출력한다.

Input 이미지

Output 이미지

모델 생성 세 번째 시도 : U-Net 구조 만들기 (Skip Connection 추가)

U-Net 구조는 초반 부분의 레이어와 후반 부분의 레이어에 skip connection을 추가함으로서 높은 공간 frequency 정보를 유지하고자 하는 방법이다. 출처 :

딥러닝을 위한 Atrous Convolution과 U-Net 구조: 간략한 역사

Skip Connection을 이용한 U-net 구조를 그림으로 도식화 해보면 다음과 같다.

(https://www.tensorflow.org/tutorials/generative/pix2pix)

Convolution 연산을 진행하면서, 인코딩 되는 과정의 이미지들을 기억해놨다가, 디코딩을 하는 과정에 같이 소스로 추가해주는 것이다. 물론 이 프로젝트에서처럼 하나의 이미지를 똑같이 복원해내는 오토인코더에서는 Skip Connection이 있는 모델의 의미가 무엇인지 의문이 들 수 있다. 하지만 해당 U-net 구조를 통해 이미지에 변형을 주는 경우에도, 조금 더 선명하고 정확한 결과를 도출할 수 있다.

따라서 Decoder 부분에 Skip Connection 레이어를 추가하여, Encoder의 Convolution 레이어와 Decoder의 Upsampling 레이어를 연결해준다. Concatenate를 사용할 경우 8x8x256의 레이어 두개를 연산하면 8x8x512의 결과를 출력하므로, Conv2D 레이어도 추가로 넣어주어 다시 8x8x256의 결과를 출력하도록 연산한다. Decoder 레이어:

convt1 = layers.Conv2DTranspose(256, (3,3), activation='relu', padding='same')(output_e)
upsamp1 = layers.UpSampling2D((2,2))(convt1)
skipcon1 = layers.Concatenate(axis=3)([conv4, upsamp1])
conv6 = layers.Conv2D(256, (3,3), activation = 'relu', padding='same')(skipcon1)

convt2 = layers.Conv2DTranspose(128, (3,3), activation='relu', padding='same')(conv6)
upsamp2 = layers.UpSampling2D((2,2))(convt2)
skipcon2 = layers.Concatenate(axis=3)([conv3, upsamp2])
conv7 = layers.Conv2D(128, (3,3), activation = 'relu', padding='same')(skipcon2)

convt3 = layers.Conv2DTranspose(64, (3,3), activation='relu', padding='same')(conv7)
upsamp3 = layers.UpSampling2D((2,2))(convt3)
skipcon3 = layers.Concatenate(axis=3)([conv2, upsamp3])
conv8 = layers.Conv2D(64, (3,3), activation='relu', padding='same')(skipcon3)

convt4 = layers.Conv2DTranspose(32, (3,3), activation='relu', padding='same')(conv8)
upsamp4 = layers.UpSampling2D((2,2))(convt4)
skipcon4 = layers.Concatenate(axis=3)([conv1, upsamp4])
conv9 = layers.Conv2D(32, (3,3), activation='relu', padding='same')(skipcon4)

output_d = layers.Conv2DTranspose(3, (3,3), activation='relu', padding='same')(conv9)

변경된 모델의 summary를 출력해보면 다음과 같다.

input_e = tf.keras.Input(shape=(128,128,3))
        
conv1 = layers.Conv2D(32, (3,3), activation='relu', padding = 'same')(input_e)
mp1 = layers.MaxPooling2D((2,2))(conv1)

conv2 = layers.Conv2D(64, (3,3), activation='relu', padding = 'same')(mp1)
mp2 = layers.MaxPooling2D((2,2))(conv2)
        
conv3 = layers.Conv2D(128, (3,3), activation='relu', padding = 'same')(mp2)
mp3 = layers.MaxPooling2D((2,2))(conv3)
        
conv4 = layers.Conv2D(256, (3,3), activation='relu', padding = 'same')(mp3)
mp4 = layers.MaxPooling2D((2,2))(conv4)
        
output_e = layers.Conv2D(512, (3,3), activation='relu', padding = 'same')(mp4)

convt1 = layers.Conv2DTranspose(256, (3,3), activation='relu', padding='same')(output_e) #512아니고 256
upsamp1 = layers.UpSampling2D((2,2))(convt1)
skipcon1 = layers.Concatenate(axis=3)([conv4, upsamp1])
conv6 = layers.Conv2D(256, (3,3), activation = 'relu', padding='same')(skipcon1)

convt2 = layers.Conv2DTranspose(128, (3,3), activation='relu', padding='same')(conv6)
upsamp2 = layers.UpSampling2D((2,2))(convt2)
skipcon2 = layers.Concatenate(axis=3)([conv3, upsamp2])
conv7 = layers.Conv2D(128, (3,3), activation = 'relu', padding='same')(skipcon2)

convt3 = layers.Conv2DTranspose(64, (3,3), activation='relu', padding='same')(conv7)
upsamp3 = layers.UpSampling2D((2,2))(convt3)
skipcon3 = layers.Concatenate(axis=3)([conv2, upsamp3])
conv8 = layers.Conv2D(64, (3,3), activation='relu', padding='same')(skipcon3)

convt4 = layers.Conv2DTranspose(32, (3,3), activation='relu', padding='same')(conv8)
upsamp4 = layers.UpSampling2D((2,2))(convt4)
skipcon4 = layers.Concatenate(axis=3)([conv1, upsamp4])
conv9 = layers.Conv2D(32, (3,3), activation='relu', padding='same')(skipcon4)

output_d = layers.Conv2DTranspose(3, (3,3), activation='relu', padding='same')(conv9)
        
model = Model(inputs=input_e, outputs=output_d)
model.compile(optimizer='adam', loss=losses.MeanSquaredError())

학습

model.fit(train_env, train_env, validation_split=0.2, epochs=300, verbose=2)

300번만 training을 시켜주었다.

테스트

test_env = []

def random_crop(image):
    cropped_image = tf.image.random_crop(image, size=[128, 128, 3])
    return cropped_image

for i in range (32):
    filename = 'nspheres\pbr_1600_env\pbr_' + str(i*50).zfill(5) + '.jpg'
    image = pilimg.open(filename)
    image = np.array(image)
    cropped_image = tf.image.random_crop(image, size=[128, 128, 3])
    test_env.append(cropped_image)
    
test_env = np.array(test_env, dtype="float32")
test_env /= 255.0

test_env.shape

result = model.predict(test_env)

위 처럼 테스트를 진행하고, 이미지를 출력해보면 결과는 다음과 같다.

Input 이미지

Output 이미지

거의 완벽하게 복원된 것을 확인할 수 있다.

---

다음 프로젝트는 이 구축된 U-Net 구조의 Autoencoder를 활용해서, 완전복원이 아닌 약간의 변형을 곁들인 개인 프로젝트를 진행해보겠습니다!