tensorflow 实现像素级图像分割算法

#!pip install -q git+https://github.com/tensorflow/examples.git#!pip install -q -U tfds-nightly
import tensorflow as tffrom tensorflow_examples.models.pix2pix import pix2piximport urllib3
import tensorflow_datasets as tfds
from IPython.display import clear_output
import matplotlib.pyplot as plt
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


for i in [tf,tfds]:    print(i.__name__,": ",i.__version__,sep="")

#这个数据集已经集成在 Tensorflow datasets 中，只需下载即可。图像分割掩码在版本 3.0.0 中才被加入，因此我们特别选用这个版本。dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)
#下图像标准化到 [0,1]。为了方便起见，我们将分割掩码都减 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.functiondef 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_examplesBATCH_SIZE = 64BUFFER_SIZE = 1000STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE
train = dataset['train'].map(load_image_train, num_parallel_calls=tf.data.experimental.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.experimental.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, maskdisplay([sample_image, sample_mask])
# 这里用到的模型是一个改版的 U-Net。U-Net 由一个编码器（下采样器（downsampler））和一个解码器（上采样器（upsampler））组成。#为了学习到鲁棒的特征，同时减少可训练参数的数量，这里可以使用一个预训练模型作为编码器。#因此，这项任务中的编码器将使用一个预训练的 MobileNetV2 模型，它的中间输出值将被使用。#解码器将使用在 TensorFlow Examples 中的 Pix2pix tutorial 里实施过的升频取样模块。# 输出信道数量为 3 是因为每个像素有三种可能的标签。把这想象成一个多类别分类，每个像素都将被分到三个类别当中。OUTPUT_CHANNELS = 3
#编码器是一个预训练的 MobileNetV2 模型，它在 tf.keras.applications 中已被准备好并可以直接使用。#编码器中包含模型中间层的一些特定输出。注意编码器在模型的训练过程中是不会被训练的。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):  inputs = tf.keras.layers.Input(shape=[128, 128, 3])  x = inputs
  # 在模型中降频取样  skips = down_stack(x)  x = skips[-1]  skips = reversed(skips[:-1])
  # 升频取样然后建立跳跃连接  for up, skip in zip(up_stack, skips):    x = up(x)    concat = tf.keras.layers.Concatenate()    x = concat([x, skip])
  # 这是模型的最后一层  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)
#训练模型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()
#我们来观察模型是怎样随着训练而改善的。为达成这一目的，下面将定义一个 callback 函数。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 = 20VAL_SUBSPLITS = 5VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS

# Fitmodel_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()

#我们来做几个预测。为了节省时间，这里只使用很少的周期（epoch）数，但是你可以设置更多的数量以获得更准确的结果。show_predictions(test_dataset, 3)

Sample Prediction after epoch 20
57/57 [==============================] - 224s 4s/step - loss: 0.1375 - accuracy: 0.9374 - val_loss: 0.3113 - val_accuracy: 0.8930