[파이썬] Keras 이미지 세그멘테이션

세그멘테이션은 컴퓨터 비전 분야에서 매우 중요한 작업 중 하나입니다. 이를 통해 이미지 내에서 각 픽셀에 레이블을 할당하여 객체의 경계를 정확하게 식별하는 것이 가능해집니다. 세그멘테이션 모델을 훈련하는 것은 대부분의 관심을 끌며, 최근에는 딥러닝 모델을 사용한 세그멘테이션의 인기가 높아지고 있습니다.

이번 포스트에서는 Keras를 사용하여 이미지 세그멘테이션을 수행하는 방법에 대해 알아보겠습니다. Keras는 간편한 API와 딥러닝 모델을 훈련하기 위한 다양한 도구를 제공하므로 이를 사용하는 것이 유용합니다.

데이터셋 준비

먼저 훈련에 사용할 이미지와 해당 이미지의 세그멘테이션 맵으로 구성된 데이터셋을 준비해야 합니다. 일반적으로 이러한 데이터셋은 픽셀별 레이블로 이루어진 이미지 세그멘테이션 맵과 해당 이미지를 포함하는 파일로 구성됩니다. 이 예제에서는 PASCAL VOC라고 불리는 널리 사용되는 데이터셋을 사용하겠습니다.

모델 아키텍처 정의

이미지 세그멘테이션을 위한 모델을 정의해야 합니다. Keras에서는 함수형 API를 사용하여 모델을 정의할 수 있습니다. 이 예제에서는 U-Net이라는 유명한 세그멘테이션 네트워크를 사용하겠습니다. 먼저 모델의 입력층을 정의합니다.

from keras.models import Model
from keras.layers import Input

input_shape = (256, 256, 3)
input_layer = Input(shape=input_shape)

U-Net은 인코더와 디코더로 구성되어 있으며, 이들은 컨볼루션, 풀링 및 업샘플링 레이어로 이루어져 있습니다. 아래 코드는 모델의 아키텍처를 정의하는 함수를 보여줍니다.

from keras.layers import Conv2D, MaxPooling2D, UpSampling2D

def build_model(input_layer):
    # 인코더
    conv1 = Conv2D(64, 3, activation='relu', padding='same')(input_layer)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
    conv2 = Conv2D(128, 3, activation='relu', padding='same')(pool1)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
    # ...
    
    # 디코더
    up1 = UpSampling2D(size=(2, 2))(conv7)
    conv8 = Conv2D(64, 3, activation='relu', padding='same')(up1)
    up2 = UpSampling2D(size=(2, 2))(conv8)
    conv9 = Conv2D(num_classes, 1, activation='softmax')(up2)

    return conv9

model_output = build_model(input_layer)

모델 훈련하기

모델을 훈련하기 전에 데이터셋을 준비해야 합니다. 이를 위해 훈련 이미지와 해당 세그멘테이션 맵을 로드하고, 필요한 전처리 작업을 수행합니다. 예를 들어, 이미지를 정규화하거나 레이블을 원-핫 인코딩으로 변환할 수 있습니다.

이제 모델을 컴파일하고 훈련을 시작할 수 있습니다. Keras에서는 model.compile()을 사용하여 모델을 컴파일하고, model.fit()을 사용하여 훈련을 수행할 수 있습니다. 아래는 간단한 예제입니다.

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(X_train, y_train, batch_size=16, epochs=10, validation_data=(X_val, y_val))

모델 평가하기

훈련이 완료되면 모델을 평가할 수 있습니다. model.evaluate()를 사용하여 테스트 데이터에서 모델의 성능을 평가할 수 있습니다.

score = model.evaluate(X_test, y_test)
print('Test loss:', score[0])
print('Test accuracy:', score[1])

세그멘테이션 결과 시각화하기

마지막으로, 훈련된 모델로 이미지 세그멘테이션을 수행하고 결과를 시각화하는 방법을 알아보겠습니다. 테스트 이미지에 대해 모델의 출력을 얻은 후, 이를 세그멘테이션 맵으로 변환하고 시각화할 수 있습니다.

import numpy as np
import matplotlib.pyplot as plt

# 테스트 이미지 선택
test_image = X_test[0]

# 세그멘테이션 수행
segmentation_map = model.predict(np.expand_dims(test_image, axis=0))

# 시각화
plt.imshow(test_image)
plt.imshow(segmentation_map[0], alpha=0.5)
plt.show()

이제 Keras를 사용하여 이미지 세그멘테이션을 수행하는 방법을 알고 있습니다. Keras의 간편한 API를 활용하여 다양한 세그멘테이션 모델을 실험하고, 컴퓨터 비전 문제를 해결하는 데 활용해보세요.