说明：本文章为作者在学习Tensorflow官方教程时的学习笔记，现整理出来供大家学习参考。您可以将本文章当作官方教程的中文翻译来阅读学习。本教程代码与官方代码一致。
Tensorflow官方教程1链接附在文章末。
图像分割
什么是图像分割？
图像分割是计算机视觉中的一个关键过程。它包括将视觉输入分割成片段以简化图像分析。片段表示目标或目标的一部分，并由像素集或“超像素”组成。图像分割将像素组织成更大的部分，消除了将单个像素作为观察单位的需要。
图像分割的任务是训练一个神经网络来输出图像的像素范围蒙版。这可以帮助我们在更低层次（比如像素层次）来理解图像。
图像分割广泛应用在例如医学成像、自动驾驶、卫星成像等方面。
本教程将使用 Oxford-IIIT Pet 数据集，该数据集包含图片的标签和像素蒙版。蒙版是每个像素的基本标签。
每个像素包含下列三个中的一个：
- Class 1：属于宠物的像素
- Class 2：宠物的像素边界
- Class 3：不包含上面/环绕像素
导入模块
pip install git+https://github.com/tensorflow/examples.git
import tensorflow as tf
from tensorflow_examples.models.pix2pix import pix2pix
import tensorflow_datasets as tfds
from IPython.display import clear_output
import matplotlib.pyplot as plt
开始
下载Oxford-IIIT Pets 数据集 & 预处理
该数据集已经包含了所需要的数据。分割蒙版包含在版本3及以上版本。
dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)
下面的代码通过翻转图片来增强我们的数据
- 分割蒙版中的像素已经标记为{1, 2, 3}，为了方便，我们将分割蒙版中的标记减1，得到新的标签结果为{0, 1, 2}；
def normalize(input_image, input_mask): #标准化图像
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.AUTOTUNE)
test = dataset['test'].map(load_image_test)
train_dataset = train.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat()
train_dataset = train_dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
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])
定义模型
我们使用的模型是一个已经改进的U-Net。U-Net由编码器（下采样器）和解码器组成。
一个预先训练好的模型用作编码器，使网络学习到健壮的特征，并且减少能够训练的参数的数量。
我们使用已经训练好的 MobileNetV2 模型作为编码器，我们将使用它的中间输出。
解码器将使用 Tensorflow Examples 中的 Pix2pix 教程中已经实现的上采样器。
OUTPUT_CHANNELS = 3
因为每个像素有三个标签，所以我们的输出通道设置为3
MobileNetV2 模型我们可以通过tf.keras.applications来调用。编码器由模型中间层的特殊输出组成。注意在模型训练的过程中将不会训练编码器。
base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)
# Use the activations of these layers
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
]
base_model_outputs = [base_model.get_layer(name).output for name in layer_names]
# Create the feature extraction model
down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)
down_stack.trainable = False
编码器是一系列已经在Tensorflow Examples中实现的上采样器。
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):
inputs = tf.keras.layers.Input(shape=[128, 128, 3])
# Downsampling through the model
skips = down_stack(inputs)
x = skips[-1]
skips = reversed(skips[:-1])
# Upsampling and establishing the skip connections
for up, skip in zip(up_stack, skips):
x = up(x)
concat = tf.keras.layers.Concatenate()
x = concat([x, skip])
# This is the last layer of the model
last = tf.keras.layers.Conv2DTranspose(
output_channels, 3, strides=2,
padding='same') #64x64 -> 128x128
x = last(x)
return tf.keras.Model(inputs=inputs, outputs=x)
训练模型
现在我们开始编译并训练模型。我们将使用losses.SparseCategoricalCrossentropy(from_logits=True)损失函数。因为网络会像多类别预测一样，为每一个像素都分配一个标签。
在实际分离蒙版中，每个像素都会有{0, 1, 2}三个标签。网络将会输出三个通道。本质上，每个通道都会学习去预测一个类别，并且该损失函数是这类方案的推荐函数。
使用网络的输出，分配到像素上的标签代表了最高值的通道。
model = unet_model(OUTPUT_CHANNELS)
model.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])
快速浏览一下结果模型的构造：
tf.keras.utils.plot_model(model, show_shapes=True)
测试一下模型看看在训练之前会预测出什么：
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']
plt.figure()
plt.plot(model_history.epoch, loss, 'r', label='Training loss')
plt.plot(model_history.epoch, 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()
开始预测
- 为了节省时间，我们继续使用较小的epochs，但是如果你想获得更加精确的结果可以把它调高。
show_predictions(test_dataset, 3)
结束
可选项：非平衡类与类权重
有兴趣可以参考官方教程
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
