tensorflow图像处理：图像语义分割

图像分割参考：语义分割（sematic segmentation）

这里使用Oxford-IIIT Pet 数据集，该数据集由图像、图像标签、以及对像素逐一标记的掩码（分类）组成。像素的标签可以是：宠物、宠物轮廓、背景

 

下载安装必要的数据集和库，导入数据集 Oxford-IIIT Pet 

!pip install -q git+https://github.com/tensorflow/examples.git
import tensorflow as tf
from __future__ import absolute_import, division, print_function, unicode_literals

from tensorflow_examples.models.pix2pix import pix2pix

import tensorflow_datasets as tfds
tfds.disable_progress_bar()

from IPython.display import clear_output
import matplotlib.pyplot as plt
!pip install tensorflow_datasets
dataset, info = tfds.load('oxford_iiit_pet:3.1.0', with_info=True)

 

定义图像的归一化、增广函数

def normalize(input_image, input_mask): #归一化图像和掩码（0,1,2变为-1,0,1）
    input_image = tf.cast(input_image, tf.float32)/255.0
    input_mask -= 1
    return input_image, input_mask

@tf.function
def load_image_train(datapoint): #训练集
    input_image = tf.image.resize(datapoint['image'], (128, 128))
    input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128))

    if tf.random.uniform(()) > 0.5: #随机将图片和掩码同时水平翻转
        input_image = tf.image.flip_left_right(input_image)
        input_mask = tf.image.flip_left_right(input_mask)

    input_image, input_mask = normalize(input_image, input_mask) # 归一化

    return input_image, input_mask

def load_image_test(datapoint): #测试机
    input_image = tf.image.resize(datapoint['image'], (128, 128))
    input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128))

    input_image, input_mask = normalize(input_image, input_mask) #归一化

    return input_image, input_mask

 

训练数据的参数设置、测试数据调整，图像一般是三维，这里在前面加上一个维度，表示批量大小

TRAIN_LENGTH = info.splits['train'].num_examples
BATCH_SIZE = 64
BUFFER_SIZE = 1000
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE

train = dataset['train'].map(load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE) #训练图像归一化，可并行，这里由tf.data自动决定并行方式
test = dataset['test'].map(load_image_test) #测试图像归一化

train_dataset = train.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat() # cache 缓存数据；shuffle：打乱数据；batch：图像是三维，这里在前面加一维批量数；repeat：重复一次数据
train_dataset = train_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) #prefetch 在训练第s步时，预先提取s+1步需要的数据
test_dataset = test.batch(BATCH_SIZE) #测试数据加一维

def display(display_list):
    plt.figure(figsize=(15, 15))

    title = ['Input Image', 'True Mask', 'Predicted Mask']

    for i in range(len(display_list)):
        plt.subplot(1, len(display_list), i+1)
        plt.title(title[i])
        plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
        plt.axis('off')
    plt.show()

for image, mask in train.take(1):
    sample_image, sample_mask = image, mask
display([sample_image, sample_mask])

 

定义模型，借助预训练模型，将前向网络每一层输出提取出来，设置前向网络参数在训练时不可调整

OUTPUT_CHANNELS = 3

base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)

# 使用这些层的激活设置
layer_names = [
    'block_1_expand_relu',   # 64x64
    'block_3_expand_relu',   # 32x32
    'block_6_expand_relu',   # 16x16
    'block_13_expand_relu',  # 8x8
    'block_16_project',      # 4x4
]
layers = [base_model.get_layer(name).output for name in layer_names] #将每一层输出提取出来

# 创建特征提取模型
down_stack = tf.keras.Model(inputs=base_model.input, outputs=layers) # 前向网络

down_stack.trainable = False # 前向网络参数在训练时不可调整

定义模型的转置卷积流程即上采样，需要上采样的次数和之前传播的层数有关，下采样卷积中越靠后的层，经过的转置卷积层（上采样）越多。将每个上采样层的输出叠加作为模型总的输出。

up_stack = [
    pix2pix.upsample(512, 3),  # 4x4 -> 8x8
    pix2pix.upsample(256, 3),  # 8x8 -> 16x16
    pix2pix.upsample(128, 3),  # 16x16 -> 32x32
    pix2pix.upsample(64, 3),   # 32x32 -> 64x64
]

def unet_model(output_channels):
# 这是模型的最后一层
    last = tf.keras.layers.Conv2DTranspose( #定义转置卷积
        output_channels, 3, strides=2,
        padding='same', activation='softmax')  #64x64 -> 128x128

    inputs = tf.keras.layers.Input(shape=[128, 128, 3])
    x = inputs

# 在模型中降频取样
    skips = down_stack(x) #前向网络传播，这里获得了每一层的输出
    x = skips[-1]  #x作为最后一层的输出，将作为转置卷积层的输入
    skips = reversed(skips[:-1]) # 将之前每一层的顺序反转

# 升频取样然后建立跳跃连接
    for up, skip in zip(up_stack, skips):
        x = up(x) # 前向输出先做上采样
        concat = tf.keras.layers.Concatenate()
        x = concat([x, skip]) #上采样后拼接每一层输出

    x = last(x) # 做转置卷积

    return tf.keras.Model(inputs=inputs, outputs=x)

model = unet_model(OUTPUT_CHANNELS)

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy',
              metrics=['accuracy']) #损失函数为交叉熵损失

def create_mask(pred_mask):
    pred_mask = tf.argmax(pred_mask, axis=-1)
    pred_mask = pred_mask[..., tf.newaxis]
    return pred_mask[0]

def show_predictions(dataset=None, num=1):
    if dataset:
        for image, mask in dataset.take(num):
            pred_mask = model.predict(image)
            display([image[0], mask[0], create_mask(pred_mask)])
    else:
        display([sample_image, sample_mask,
            create_mask(model.predict(sample_image[tf.newaxis, ...]))])

show_predictions()

 

模型训练

class DisplayCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        clear_output(wait=True)
        show_predictions()
        print ('\nSample Prediction after epoch {}\n'.format(epoch+1))


EPOCHS = 20
VAL_SUBSPLITS = 5
VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS

model_history = model.fit(train_dataset, epochs=EPOCHS,
                          steps_per_epoch=STEPS_PER_EPOCH,
                          validation_steps=VALIDATION_STEPS,
                          validation_data=test_dataset,
                          callbacks=[DisplayCallback()])

 

画出损失函数

loss = model_history.history['loss']
val_loss = model_history.history['val_loss']

epochs = range(EPOCHS)

plt.figure()
plt.plot(epochs, loss, 'r', label='Training loss')
plt.plot(epochs, val_loss, 'bo', label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.ylim([0, 1])
plt.legend()
plt.show()

 

预测结果

show_predictions(test_dataset, 3)