[Tensorflow] 이미지/AutoEncoder/UNet - 초록색 잎을 보라색으로 바꾸기

이번 프로젝트는 저번 프로젝트에 이어 U-Net 구조의 Autoencoder를 활용한 개인 프로젝트를 진행합니다. 개요는 다음과 같습니다.

 

 

 

1. 웹에서 나무(또는 풀)이 포함된 사진 200개 크롤링

2. Pinkl 앱으로 1번의 사진에 Plum.V라는 필터를 씌워 초록색 잎들을 다 보라색으로 변형한 데이터 셋을 만든다.

‎Pinkl - Infrared filter camera

데이터셋 예시

3. 지난 프로젝트에서 구현한 Autoencoder를 활용하여 200개의 초록색 잎 사진으로 200개의 보라색 잎 사진을 만들도록 모델을 학습시킨다.

4. 새로운 test 데이터셋을 생성하여 모델에 넣고, 출력된 결과를 확인한다.

 

---

dataset hierarchy

 

📁 green leaves (train용 초록색 잎 이미지 - 웹에서 크롤링)

ㄴ 초록 잎 이미지 200개 (green (1).jpeg ~ green (200).jpeg)

📁 purple leaves (train용 보라색 잎 이미지 - 앱으로 생성)

ㄴ 보라 잎 이미지 200개 (purple (1).jpg ~ purple (200).jpg)

📁 test leaves (test용 초록색 잎 이미지 - 웹에서 다운로드)

ㄴ 초록 잎 이미지 10개 (1.jpg ~ 10.jpg)

 

 

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

 

 

 

전처리

import cv2

train_green = []

for i in range (200):
    filename = 'leaves200\green leaves\green (' + str(i+1) + ').jpeg'
    try:
        image = cv2.imread(filename, cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # BGR->RGB
        image = cv2.resize(image,(256,256), interpolation=cv2.INTER_AREA)
    except Exception as e:
        print(str(e))
    image = np.array(image)
    train_green.append(image)
    
train_green = np.array(train_green, dtype="float32")
train_green /= 255.0

이미지를 웹에서 크롤링한 것이기 때문에, 이미지 크기가 다 제각각이다. 따라서 opencv를 활용해서 이미지를 읽고 (cv2.imread), 256x256 크기로 일괄 변경 (cv2.resize) 해주었다.

 

이 과정에서 !ssize.empty() in function 'cv::resize' 라는 오류가 생성되었는데, 알고보니 이 메세지는 opencv로 읽어온 이미지가 비어있을 때 나는 오류였다. 이런 오류가 생긴 이유는 filename 설정 시 주소를 표현하다 보니 \ 기호를 사용하는데, \v \n \t 처럼 특수한 문자가 오게 되는 경우 의도한 파일 이름을 설정할 수 없었기 때문이다. 이럴 때는 \를 두 번, \\ 이렇게 입력하여 정말 '\' 문자를 읽어달라고 표시해야 한다.

 

또 유의해야 할 점은, opencv는 이미지를 RGB 순으로 불러오지 않는다. BGR 순으로 읽는다. 따라서 cv2.imread()로 읽은 이미지는 다음과 같이 의도하지 않은 색상을 띌 수 있다. 위 코드에 주석으로 표시한 것처럼, cv2.COLOR_BGR2RGB를 사용하여 RGB 순으로 이미지를 다시 변경해주었다.

(좌) BGR 순으로 읽혀진 이미지 (우) RGB 순으로 변경한 이미지

train_purple = []

for i in range (200):
    filename = 'leaves200\purple leaves\purple (' + str(i+1) + ').jpg'
    try:
        image = cv2.imread(filename, cv2.IMREAD_COLOR)
        image = cv2.resize(image,(256,256), interpolation=cv2.INTER_AREA)
    except Exception as e:
        print(str(e)+str(i))
    image = np.array(image)
    train_purple.append(image)
    
train_purple = np.array(train_purple, dtype="float32")
train_purple /= 255.0

같은 방식으로 training을 위한 보라 잎 이미지들도 불러왔다.

 

 

 

모델 생성

input_g = tf.keras.Input(shape=(256,256,3))
        
conv0 = layers.Conv2D(16, (3,3), activation='relu', padding = 'same')(input_g)
mp0 = layers.MaxPooling2D((2,2))(conv0)
        
conv1 = layers.Conv2D(32, (3,3), activation='relu', padding = 'same')(mp0)
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)
        
convt5 = layers.Conv2DTranspose(16, (3,3), activation='relu', padding='same')(conv9)
upsamp5 = layers.UpSampling2D((2,2))(convt5)
skipcon5 = layers.Concatenate(axis=3)([conv0, upsamp5])
conv10 = layers.Conv2D(16, (3,3), activation='relu', padding='same')(skipcon5)

output_p = layers.Conv2DTranspose(3, (3,3), activation='relu', padding='same')(conv10)

지난 프로젝트와 비슷하지만, 이미지 크기가 128x128이 아닌 256x256이므로 레이어를 한 단계씩 더 추가해 주었다. conv0, mp0, conv5, upsamp5, skipcon5, conv10 이 새로 추가된 레이어들이다.

 

 

 

학습

model = Model(inputs=input_g, outputs=output_p)
model.compile(optimizer='adam', loss=losses.MeanSquaredError())
model.fit(train_green, train_purple, validation_split=0.2, epochs=1000, verbose=2)

이번 모델을 학습시킬 때에는, train_green을 넣고 train_purple이 타겟이 되도록 설정하는 것을 잊지 말자.

 

 

테스트

test_green = []

for i in range (10):
    filename = 'leaves200\\test leaves\\' + str(i+1) + '.jpg'
    try:
        image = cv2.imread(filename, cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, (256,256), interpolation=cv2.INTER_AREA)
    except Exception as e:
        print(str(e)+str(i))
    image = np.array(image)
    test_green.append(image)

test_green = np.array(test_green, dtype="float32")
test_green /= 255.0

테스트를 할 때도 문자에 유의하고, RGB 순서를 변경해준다.

 

def generate_images(test_input, test_output):
    plt.figure(figsize=(10, 10))
    display_list = [test_input, test_output]
    title = ['Input Image', 'Predicted Image']
    
    for i in range(2):
        plt.subplot(1, 2, i+1)
        plt.title(title[i])
        plt.imshow(display_list[i])
        plt.axis('off')
    plt.show()

이미지를 나란히 출력하기 위해 이미지를 출력하는 함수를 하나 만들었다.

 

test_prediction = model(test_green)

i = 0
for image in test_prediction:
    generate_images(test_green[i], image)
    i += 1

test_green을 모델에 넣어 생성한 함수로 결과를 출력하면 다음과 같이 비교를 할 수 있다.

 

원한다면, 데이터셋을 만들었던 앱으로 같은 이미지를 보라색으로 변경해보고, 모델의 출력결과와 비교해보아도 재밌을 것이다.

순서대로 원본 이미지, 모델이 예측한 이미지, 앱으로 변형한 이미지

결과를 보면 모델이 생성한 이미지와 사용한 앱이 생성한 이미지가 거의 똑같은 결과를 출력함을 알 수 있다. 여기에서 또 다른 도전을 하자면, 숫자와 함께 필터의 정도를 조정하도록 구현할 수도 있을 것이고, 단순히 색만 변형하는 데이터가 아닌 형태도 변형하는 데이터셋을 만들어 같은 결과를 기대해 볼 수도 있을 것이다.

 

 

 

---

다음 프로젝트는 GAN을 활용한 (드디어) 프로젝트를 해보려고 합니다!