영역 분할

import os
from glob import glob

data_dir = '../data/instance-level-human-parsing/instance-level_human_parsing/instance-level_human_parsing/'
# 훈련 데이터셋 파일 경로
train_image_list = glob(os.path.join(data_dir, 'Training', 'Images/*.jpg'))
train_mask_list = glob(os.path.join(data_dir, 'Training', 'Category_ids/*.png'))
# 검증 데이터셋 파일 경로
val_image_list = glob(os.path.join(data_dir, 'Validation', 'Images/*.jpg'))
val_mask_list = glob(os.path.join(data_dir, 'Validation', 'Category_ids/*.png'))
# 테스트 데이터셋 파일 경로
test_image_list = glob(os.path.join(data_dir, 'Testing', 'Images/*.jpg'))

print(len(train_image_list), len(train_mask_list))
print(len(val_image_list), len(val_mask_list))
print(len(test_image_list))

from collections import namedtuple
import tensorflow as tf

configs = namedtuple('configs', ['image_size', 'batch_size', 'num_classes'])
configs.image_size = (512, 512)
configs.batch_size = 4
configs.num_classes = 20

def read_image(image_path, image_size, mask=False):
    image = tf.io.read_file(image_path)
    if mask:
        image = tf.image.decode_png(image, channels=1)
        image.set_shape([None, None, 1])
        image = tf.image.resize(images=image, size=image_size)
    else:
        image = tf.image.decode_jpeg(image, channels=3)
        image.set_shape([None, None, 3])
        image = tf.image.resize(images=image, size=image_size)
        image = tf.keras.applications.resnet50.preprocess_input(image)
    return image

def load_data(image_list, mask_list):
    image = read_image(image_list, image_size=configs.image_size)
    mask = read_image(mask_list, image_size=configs.image_size, mask=True)
    return image, mask

def data_generator(image_list, mask_list):
    dataset = tf.data.Dataset.from_tensor_slices((image_list, mask_list))
    dataset = dataset.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(configs.batch_size, drop_remainder=True)
    return dataset

train_dataset = data_generator(train_image_list[:1000], train_mask_list[:1000])
valid_dataset = data_generator(val_image_list[:50], val_mask_list[:50])

for image, mask in train_dataset.take(1):
    print(image.shape, mask.shape)

import tensorflow.keras.layers as layers

def conv_block(inputs, filters=256, kernel_size=3, dilation_rate=1, padding='same', use_bias=False):
    x = layers.Conv2D(filters=filters, kernel_size=kernel_size, dilation_rate=dilation_rate, padding=padding, use_bias=use_bias)(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    return x

def DilatedSpatialPyramidPooling(inputs):
    dims = inputs.shape
    x = layers.AveragePooling2D(pool_size=(dims[-3], dims[-2]))(inputs)
    x = conv_block(x, kernel_size=1, use_bias=True)
    out_pool = layers.UpSampling2D(size=(dims[-3]//x.shape[1], dims[-2]//x.shape[2]), interpolation='bilinear')(x)

    out_1 = conv_block(inputs, kernel_size=1, dilation_rate=1)
    out_6 = conv_block(inputs, kernel_size=3, dilation_rate=6)
    out_12 = conv_block(inputs, kernel_size=3, dilation_rate=12)
    out_18 = conv_block(inputs, kernel_size=3, dilation_rate=18)

    x = layers.Concatenate(axis=-1)([out_pool, out_1, out_6, out_12, out_18])
    output = conv_block(x, kernel_size=1)
    return output

def build_model(input_shape=(512, 512, 3), num_classes=20):
    inputs = tf.keras.Input(shape=input_shape)
    # Encoder
    resnet50 = tf.keras.applications.ResNet50(include_top=False, input_tensor=inputs, weights='imagenet')
    x = resnet50.get_layer('conv4_block6_2_relu').output
    x = DilatedSpatialPyramidPooling(x)
    # Decoder
    IMAGE_SIZE = input_shape[0]
    input_a = layers.UpSampling2D(size=(IMAGE_SIZE // 4 // x.shape[1], IMAGE_SIZE // 4 // x.shape[2]), interpolation='bilinear')(x)
    input_b = resnet50.get_layer('conv2_block3_2_relu').output
    input_b = conv_block(input_b, filters=48, kernel_size=1)
    
    x = layers.Concatenate(axis=-1)([input_a, input_b])
    x = conv_block(x)
    x = conv_block(x)
    x = layers.UpSampling2D(size=(IMAGE_SIZE // x.shape[1], IMAGE_SIZE // x.shape[2]), interpolation='bilinear')(x)

    outputs = layers.Conv2D(num_classes, kernel_size=1, padding='same')(x)
    return tf.keras.Model(inputs=inputs, outputs=outputs, name='DeepLabV3_Plus')

model = build_model(input_shape=configs.image_size + (3,), num_classes=configs.num_classes)
model.summary()

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3), 
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), 
    metrics=['acc'])

callbacks = [
    tf.keras.callbacks.ModelCheckpoint('./checkpoints/DeepLabV3_Plus.h5', save_best_only=True)
]

history = model.fit(
    train_dataset,
    validation_data=valid_dataset,
    epochs=2,
    callbacks=callbacks
)