Tensorflow

케라스 번역: Image segmentation with a U-Net-like architecture

카카오그래놀라 2020. 11. 7. 02:04

Image segmentation with a U-Net-like architecture

Image segmentation

Author: fchollet
Date created: 2019/03/20
Last modified: 2020/04/20
Description: Image segmentation model trained from scratch on the Oxford Pets dataset.

 

데이터 다운로드

Download the data
# 윈도우 사용자는 해당 URL을 웹브라우져에 복사 붙여넣기 하면 다운로드 가능합니다.
#
# !curl -O http://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz
# !curl -O http://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz
# !tar -xf images.tar.gz
# !tar -xf annotations.tar.gz

 

Input 이미지와 Target인 segmentation masks의 경로를 지정합니다.

Prepare paths of input images and target segmentation masks
import os

input_dir = "images/"
target_dir = "annotations/trimaps/"
img_size = (160, 160)
num_classes = 4
batch_size = 32

input_img_paths = sorted(
    [
        os.path.join(input_dir, fname)
        for fname in os.listdir(input_dir)
        if fname.endswith(".jpg")
    ] # Input이미지의 주소값을 담은 리스트: jpg로 끝나는 이미지만 담았으며, 정렬 적용
)
target_img_paths = sorted(
    [
        os.path.join(target_dir, fname)
        for fname in os.listdir(target_dir)
        if fname.endswith(".png") and not fname.startswith(".")
    ] # Target이미지의 주소값을 담은 리스트:
      # png로 끝나고, `.`으로 시작하지 않는 이미지만 담은 리스트. 정렬 적용
)

# Input 데이터 개수
print("Number of samples:", len(input_img_paths))

# 주소가 잘 들어갔는지, 10개 정도 확인해봅시다.
for input_path, target_path in zip(input_img_paths[:10], target_img_paths[:10]):
    print(input_path, "|", target_path)

 

Input image에 대응하는 segmentation mask를 함께 봐봅시다.

What does one input image and corresponding segmentation mask look like?
from IPython.display import Image, display
from tensorflow.keras.preprocessing.image import load_img
import PIL
from PIL import ImageOps

# 9번째 이미지를 봅시다.
display(Image(filename=input_img_paths[9]))

# auto-contrast를 적용한 target: segmentation mask 이미지를 봅시다.
# mask의 값들을 실제로 보면, 0,1,2 이렇게 굉장히 비슷한 값을 가집니다.
# 그래서 그냥 display하면 육안으로 확인하기가 힘들어 auto-contrast를 적용해서 봅니다.

# Display auto-contrast version of corresponding target (per-pixel categories)
img = PIL.ImageOps.autocontrast(load_img(target_img_paths[9]))
display(img)

 

배치 단위의 데이터를 vectorize하고 load하는 Sequence 클래스를 정의합니다.

Prepare Sequence class to load & vectorize batches of data
from tensorflow import keras
import numpy as np
from tensorflow.keras.preprocessing.image import load_img


class OxfordPets(keras.utils.Sequence):
    """Helper to iterate over the data (as Numpy arrays)."""

    def __init__(self, batch_size, img_size, input_img_paths, target_img_paths):
        self.batch_size = batch_size
        self.img_size = img_size
        self.input_img_paths = input_img_paths
        self.target_img_paths = target_img_paths

    def __len__(self):
        return len(self.target_img_paths) // self.batch_size

    def __getitem__(self, idx):
        """Returns tuple (input, target) correspond to batch #idx."""
        """배치 순서에 대응하는 튜플(input, target)을 반환합니다."""
        
        # 배치가 4라면, 0,4,8,12... 이렇게 시작 인덱스 지정
        i = idx * self.batch_size 
        
        # 배치 단위로 이미지 담음. 배치가 4라면 데이터가 4개씩 담김
        batch_input_img_paths = self.input_img_paths[i : i + self.batch_size]
        batch_target_img_paths = self.target_img_paths[i : i + self.batch_size]
        
        # 배치가 4라면 (4, IMG_SIZE, IMG_SIZE, 3)의 shape를 가지고,
        # 모든 값은 0으로 채워진 x를 정의합니다.
        x = np.zeros((batch_size,) + self.img_size + (3,), dtype="float32")
        
        # 배치에 담긴 데이터를 하나씩 꺼내서
        # 위에서 만든 np.array에 대입합니다.
        # j는 배치속에서 몇번째 인지 나타냅니다.
        # enumerate => 루프문에서 리스트 순서와 값을 동시에 반환합니다.
        # 배치가 4라면, (0, path0), (1, path1), (2, path2), (3, path3)을 반환합니다.
        for j, path in enumerate(batch_input_img_paths):
            img = load_img(path, target_size=self.img_size)
            x[j] = img
           
        # 배치가 4라면 (4, IMG_SIZE, IMG_SIZE, 1)의 shape를 가지고,
        # 모든 값은 0으로 채워진 y를 정의합니다.
        # 배치에 담긴 데이터(여기서는 target값, segmentation mask)를 하나씩 꺼내서
        # 위에서 만든 np.array에 대입합니다.
        y = np.zeros((batch_size,) + self.img_size + (1,), dtype="uint8")
        for j, path in enumerate(batch_target_img_paths):
            img = load_img(path, target_size=self.img_size, color_mode="grayscale")
            # 이미지가 1차원으로 되어있어서 불러올때 (IMG_SIZE, IMG_SIZE) 2차원으로 불러옵니다.
            # 따라서 차원을 증가시켜준 후에, y에 대입합니다. (IMG_SIZE, IMG_SIZE, 1)
            y[j] = np.expand_dims(img, 2)
        return x, y

 

U-Net 과 X-ception 스타일의 모델을 정의합니다.

Perpare U-Net Xception-style model
from tensorflow.keras import layers


def get_model(img_size, num_classes):
    inputs = keras.Input(shape=img_size + (3,))

    ### [First half of the network: downsampling inputs] ###
    ### 네트워크의 앞부분은 inputs을 downsampling 합니다. (이미지 크기를 줄이고, 특징을 획득)

    # Entry block
    x = layers.Conv2D(32, 3, strides=2, padding="same")(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)

    previous_block_activation = x  # 잔차를 위해 놔둡시다.

    # Blocks 1, 2, 3 are identical apart from the feature depth.
    for filters in [64, 128, 256]:
        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.MaxPooling2D(3, strides=2, padding="same")(x)

        # 잔차를 레이어 크기에 맞춥시다.
        residual = layers.Conv2D(filters, 1, strides=2, padding="same")(
            previous_block_activation
        )
        x = layers.add([x, residual])  # 잔차를 더해줍시다.
        previous_block_activation = x  # 다음 잔차 계산을 위해 놔둡시다.

    ### [Second half of the network: upsampling inputs] ###
    ### 네트워크의 앞부분은 inputs을 upsampling 합니다. (이미지 크기를 원 상태로 복구)

    for filters in [256, 128, 64, 32]:
        x = layers.Activation("relu")(x)
        x = layers.Conv2DTranspose(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.Activation("relu")(x)
        x = layers.Conv2DTranspose(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.UpSampling2D(2)(x)

        # 잔차를 레이어 크기에 맞춥시다.
        residual = layers.UpSampling2D(2)(previous_block_activation)
        residual = layers.Conv2D(filters, 1, padding="same")(residual)
        x = layers.add([x, residual])  # 잔차를 더해줍시다.
        previous_block_activation = x  # 다음 잔차 계산을 위해 놔둡시다.

    # 픽셀별 분류 레이어를 추가합니다.
    outputs = layers.Conv2D(num_classes, 3, activation="softmax", padding="same")(x)

    # inputs, outputs을 토대로 모델을 정의합시다.
    model = keras.Model(inputs, outputs)
    return model


# 모델 정의 셀(jupyter notebook 상에서 하나의 셀)이 여러번 실행된 경우, 램 확보를 합시다.
keras.backend.clear_session()

# 모델을 만듭시다.
model = get_model(img_size, num_classes)
model.summary()

 

데이터셋을 Train과 Validation 셋으로 나눕니다.

Set aside a validation split
import random

# Split our img paths into a training and a validation set
# training과 validation 셋으로 나눕니다.

# 1000개의 validation 셋을 만들겠습니다.
val_samples = 1000

# seed 값을 동일하게 주고, 섞었기 때문에, input과 target이 서로 일치하게 섞였습니다.
random.Random(1337).shuffle(input_img_paths)
random.Random(1337).shuffle(target_img_paths)

# 원래 데이터셋 개수 - 1000개의 데이터를 가지게 된 training set
train_input_img_paths = input_img_paths[:-val_samples]
train_target_img_paths = target_img_paths[:-val_samples]

# 1000개의 validation 셋
val_input_img_paths = input_img_paths[-val_samples:]
val_target_img_paths = target_img_paths[-val_samples:]

# OxfordPets의 인스턴스인 train_gen, val_gen을 만듭니다.
# Instantiate data Sequences for each split
train_gen = OxfordPets(
    batch_size, img_size, train_input_img_paths, train_target_img_paths
)
val_gen = OxfordPets(batch_size, img_size, val_input_img_paths, val_target_img_paths)

 

모델을 학습합니다.

Train the model
# Configure the model for training.
# traing을 하기위해 모델을 compile 합시다.
# 우리는 sparse categorical_crossentropy를 사용할 것입니다.
# target값이 one hot encoding 아닌 정수로 되어있기 때문입니다.

model.compile(optimizer="rmsprop", loss="sparse_categorical_crossentropy")

callbacks = [
    keras.callbacks.ModelCheckpoint("oxford_segmentation.h5", save_best_only=True)
]

# Train 하면서, 매 에폭이 끝날 때마다 validation 결과값을 봅시다.
epochs = 15
model.fit(train_gen, epochs=epochs, validation_data=val_gen, callbacks=callbacks)

 

예측값을 구하고, 시각화해봅시다.

Visualize predictions
# Generate predictions for all images in the validation set

# valiation 데이터셋 generator를 만듭니다.
val_gen = OxfordPets(batch_size, img_size, val_input_img_paths, val_target_img_paths)
val_preds = model.predict(val_gen)


# 시각화해봅시다.
def display_mask(i):
    """Quick utility to display a model's prediction."""
    mask = np.argmax(val_preds[i], axis=-1)
    mask = np.expand_dims(mask, axis=-1)
    img = PIL.ImageOps.autocontrast(keras.preprocessing.image.array_to_img(mask))
    display(img)


# 10번째 이미지를 시각화해봅시다.
# Display results for validation image #10
i = 10

# 이미지를 봅시다.
display(Image(filename=val_input_img_paths[i]))

# 실제 segmentation 값을 봅시다.
img = PIL.ImageOps.autocontrast(load_img(val_target_img_paths[i]))
display(img)

# 예측한 segmentation 값을 봅시다.
# Display mask predicted by our model
display_mask(i)  # Note that the model only sees inputs at 150x150.