ICode9

精准搜索请尝试: 精确搜索
首页 > 其他分享> 文章详细

图像分割之Unet网络详解

2021-10-20 22:00:17  阅读:438  来源: 互联网

标签:image mask sample Unet 详解 图像 tf input model


本笔记基于tensorflow-2版本,先贡献代码或者下载代码(可能需要科学上网)。

什么是图像分割

在图像分类任务中,网络为每个输入图像分配一个标签(或类别)。然而,假设你想知道该物体的形状,哪个像素属于哪个物体,等等。在这种情况下,你会想给图像的每个像素分配一个类别。这项任务被称为分割。一个分割模型会返回关于图像的更详细的信息。图像分割在医学成像、自动驾驶汽车和卫星成像方面有许多应用,举例说明一下。

这里有一个数据集,该数据集由37个宠物品种的图像组成,每个品种有200张图像(在训练和测试部分各约100张)。每张图片都包括相应的标签和像素级的掩码。掩码是每个像素的类别标签。每个像素被赋予三个类别中的一个:

  • 第一类:表现宠物的像素
  • 第二类:与宠物边缘接壤的像素
  • 第三类:以上都不是的像素
    该例子可以通过下面的程序安装:
pip install git+https://github.com/tensorflow/examples.git

如果无法链接GitHub或者链接网速很慢,可以通过这个镜像进行安装,同样其它任何GitHub的包都可以使用这个镜像进行下载或者安装:

pip install git+https://github.com.cnpmjs.org/tensorflow/examples.git

导入TensorFlow相关包:

import tensorflow as tf
from tensorflow.keras.layers.experimental import preprocessing 

import tensorflow_datasets as tfds

from tensorflow_examples.models.pix2pix import pix2pix

from IPython.display import clear_output
import matplotlib.pyplot as plt

下载Oxford-IIIT宠物数据(下载大小773.52M,数据集大小774.69M):

dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)

也可以手动下载
图片展示:
在这里插入图片描述
此外,图像颜色值被归一化为[0,1]范围。最后,如上所述,分割中的像素被标记为{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

整理:

def load_image(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

下面的类执行了一个简单的扩增,即对一个图像进行随机翻转操作:

class Augment(tf.keras.layers.Layer):
  def __init__(self, seed=42):
    super().__init__()
    # both use the same seed, so they'll make the same randomn changes.
    self.augment_inputs = preprocessing.RandomFlip(mode="horizontal", seed=seed)
    self.augment_labels = preprocessing.RandomFlip(mode="horizontal", seed=seed)

  def call(self, inputs, labels):
    inputs = self.augment_inputs(inputs)
    labels = self.augment_labels(labels)
    return inputs, labels

构建输入的pipeline,在对输入进行批处理后然后使用Augmentation:

train_batches = (
    train_images
    .cache()
    .shuffle(BUFFER_SIZE)
    .batch(BATCH_SIZE)
    .repeat()
    .map(Augment())
    .prefetch(buffer_size=tf.data.AUTOTUNE))

test_batches = test_images.batch(BATCH_SIZE)

图像例子和它的MASK图:

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 images, masks in train_batches.take(2):
  sample_image, sample_mask = images[0], masks[0]
  display([sample_image, sample_mask])

下面定义Unet模型

我们这里使用的是改进的unet网络模型。一个unet由一个编码器(下采样器)和解码器(上采样器)组成。为了更稳健的学习到图像特征并减少可训练参数的数量,这里将使用一个预训练的模型 MobileNetV2作为编码器。对于解码器,可以使用上采样模块,该模块已经在TensorFlow实例库中的Pix2pix中实现了。

编码器是训练好的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中实现的一系列上采样模块。

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:int):
  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(
      filters=output_channels, kernel_size=3, strides=2,
      padding='same')  #64x64 -> 128x128

  x = last(x)

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

注意:上一层过滤器数量设置为输出通道数量,这也是该层的输出通道数。

训练模型

这个是一个多分类问题,采用CetegforicalCrossentropy(from_logits=True)作为标准的损失函数,即使用 losses.SparseCategoricalCrossentropy(from_logits=True),原因是因为标签是标量的,而不是每个类的分数向量。

OUTPUT_CLASSES = 3

model = unet_model(output_channels=OUTPUT_CLASSES)
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_batches, epochs=EPOCHS,
                          steps_per_epoch=STEPS_PER_EPOCH,
                          validation_steps=VALIDATION_STEPS,
                          validation_data=test_batches,
                          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()

预测:

show_predictions(test_batches, 3)

语义分割数据集可能是高度不平衡的,这就意味着特定类别的像素可以比其他像素的权重更高。由于分割问题可以使用像素级别问题来处理,我们就可以通过加权损失函数来处理不平衡问题。可以参考这里
model.fit目前不支持3+维度的数据:

try:
  model_history = model.fit(train_batches, epochs=EPOCHS,
                            steps_per_epoch=STEPS_PER_EPOCH,
                            class_weight = {0:2.0, 1:2.0, 2:1.0})
  assert False
except Exception as e:
  print(f"{type(e).__name__}: {e}")

报错:

ValueError: `class_weight` not supported for 3+ dimensional targets.

因此我们需要自己实现加权。可以参考样本权重:model.fit可以接受(data, label)格式之外,还接受(data, label, sample_weight)三维样式。

model.fit可以将sample_weight传给损失函数和矩阵, 然后样本权重会乘以样本。例如:

label = [0,0]
prediction = [[-3., 0], [-3, 0]] 
sample_weight = [1, 10] 

loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True,
                                               reduction=tf.losses.Reduction.NONE)
loss(label, prediction, sample_weight).numpy()

因此为了生成样本权重,我们需要一个函数,这个函数输入(data, label),然后输出(data, label, sample_weight),其中样本权重包含每个像素的权重。

最简单的是将标签作为权重列表的索引:

def add_sample_weights(image, label):
  # The weights for each class, with the constraint that:
  #     sum(class_weights) == 1.0
  class_weights = tf.constant([2.0, 2.0, 1.0])
  class_weights = class_weights/tf.reduce_sum(class_weights)

  # Create an image of `sample_weights` by using the label at each pixel as an 
  # index into the `class weights` .
  sample_weights = tf.gather(class_weights, indices=tf.cast(label, tf.int32))

  return image, label, sample_weights

生成的数据集每个成分都包含了三个元素:

train_batches.map(add_sample_weights).element_spec

现在我们就可以在加权数据集上进行训练模型:

weighted_model = unet_model(OUTPUT_CLASSES)
weighted_model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'])

weighted_model.fit(
    train_batches.map(add_sample_weights),
    epochs=1,
    steps_per_epoch=10)

标签:image,mask,sample,Unet,详解,图像,tf,input,model
来源: https://blog.csdn.net/qq_28531269/article/details/120875597

本站声明: 1. iCode9 技术分享网(下文简称本站)提供的所有内容,仅供技术学习、探讨和分享;
2. 关于本站的所有留言、评论、转载及引用,纯属内容发起人的个人观点,与本站观点和立场无关;
3. 关于本站的所有言论和文字,纯属内容发起人的个人观点,与本站观点和立场无关;
4. 本站文章均是网友提供,不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属;如您发现该文章侵犯了您的权益,可联系我们第一时间进行删除;
5. 本站为非盈利性的个人网站,所有内容不会用来进行牟利,也不会利用任何形式的广告来间接获益,纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

专注分享技术,共同学习,共同进步。侵权联系[81616952@qq.com]

Copyright (C)ICode9.com, All Rights Reserved.

ICode9版权所有