在 TensorFlow 之中進(jìn)行圖像分割

在之前的學(xué)習(xí)之中，對于圖像數(shù)據(jù)，我們進(jìn)行過分類等一些常見的任務(wù)；這節(jié)課我們便來學(xué)習(xí)一下對于圖像數(shù)據(jù)的另外一種任務(wù)：圖像分割。

1. 什么是圖像分割

圖像分割，顧名思義，就是對圖像數(shù)據(jù)進(jìn)行分割，而分類的物體一般是我們認(rèn)為進(jìn)行指定的。比如物品分割、人臉分割、醫(yī)學(xué)病灶分割等。

舉個(gè)例子，如下圖所示，原來的圖像是一個(gè)馬路的圖片，通過圖像分割，我們會按照不同的物體進(jìn)行不同的分割，比如車分為一類、人分為一類、建筑分為一類、馬路分為一類等。

圖像分割是很多任務(wù)的前提，有很多的任務(wù)只有進(jìn)行了有效的分割之后才能進(jìn)行有效的處理，比如：

• 醫(yī)學(xué)病灶識別；
• 人臉情緒識別；
• 路況檢測；
• 自動駕駛；
• 等等。

2. 如何進(jìn)行圖像分割

圖像分割看上去是一個(gè)很復(fù)雜的任務(wù)，但是實(shí)現(xiàn)起來的原理卻是非常簡單，具體來說分為以下幾步：

• 確定要分類的類別，比如，我們可以將圖片中所有的物體分割為 10 類，包括車、人等；
• 對于每個(gè)像素點(diǎn)進(jìn)行數(shù)字分類，數(shù)字分類的類別數(shù)量對應(yīng)于上述的類別，這里是 10 ；
• 將每個(gè)數(shù)字類別對應(yīng)于分類的類別，比如 0 代表車、1 代表人。

可以看出，圖像分割任務(wù)其實(shí)就是一個(gè)分類任務(wù)，只不過是對于每個(gè)像素點(diǎn)進(jìn)行分類，也就是確定每個(gè)像素點(diǎn)所對應(yīng)的類別。

在這節(jié)課之中，我們會使用圖像分割的基礎(chǔ)數(shù)據(jù)集：oxford_iiit_pet 圖像分割數(shù)據(jù)集來進(jìn)行演示。與此同時(shí)，我們也會采用之前學(xué)習(xí)到的遷移學(xué)習(xí)的方式來進(jìn)行模型的構(gòu)建，從而完成圖像分割的任務(wù)。

3. 使用 TensorFlow 進(jìn)行圖像分割的程序示例

在 oxford_iiit_pet 之中，所有的圖片都是寵物，我們的任務(wù)是將圖片中的寵物分割出來，所有的像素點(diǎn)都被分為三類：

• 1： 對應(yīng)于寵物的一部分；
• 2： 對應(yīng)于寵物的邊界；
• 3： 不屬于寵物的一部分。

在這里，我們使用代碼有一部分來自 TensorFlow 官方的一個(gè)例子，這個(gè)例子非常的簡單易懂，作為圖像分割任務(wù)的入門是再適合不過的了。

我們會逐步進(jìn)行代碼的解釋與理解，從而幫助大家學(xué)習(xí)圖像分割的任務(wù)的特點(diǎn)。

1. 首先我們獲取數(shù)據(jù)集

import tensorflow as tf
from tensorflow_examples.models.pix2pix import pix2pix
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt

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

這里會下載數(shù)據(jù)集，因?yàn)槭菆D片數(shù)據(jù)，因此數(shù)據(jù)集相對比較大。

2. 定義歸一化處理函數(shù)

def normalize(input_image, input_mask):
  input_image = tf.cast(input_image, tf.float32) / 255.0
  return input_image, input_mask

它接收兩個(gè)參數(shù)，第一個(gè)參數(shù)是圖片，我們會將其歸一化到 [0, 1] ，第二個(gè)參數(shù)是圖像的標(biāo)簽。

3. 構(gòu)建數(shù)據(jù)集

def load_image_train(data):
  input_image = tf.image.resize(data['image'], (128, 128))
  input_mask = tf.image.resize(data['segmentation_mask'], (128, 128))

  input_image, input_mask = normalize(input_image, input_mask)

  return input_image, input_mask

def load_image_test(data):
  input_image = tf.image.resize(data['image'], (128, 128))
  input_mask = tf.image.resize(data['segmentation_mask'], (128, 128))

  input_image, input_mask = normalize(input_image, input_mask)

  return input_image, input_mask

num_examples = info.splits['train'].num_examples
BATCH = 64
step_per_epch = num_examples // BATCH

train = dataset['train'].map(load_image_train)
test = dataset['test'].map(load_image_test)

train_dataset = train.cache().shuffle(1000).batch(BATCH).repeat()
test_dataset = test.batch(BATCH)

在構(gòu)建數(shù)據(jù)集函數(shù)之中，我們做了兩件事情：

• 將圖像與標(biāo)簽重新調(diào)整大小到 [128, 128] ；
• 將數(shù)據(jù)歸一化。

然后我們進(jìn)行了分批的處理，這里取批次的大小為 64 ，大家可以根據(jù)自己的內(nèi)存或現(xiàn)存大小靈活調(diào)整。

4. 構(gòu)建網(wǎng)絡(luò)模型

output_channels = 3

# 獲取基礎(chǔ)模型
base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)

# 定義要使用其輸出的基礎(chǔ)模型網(wǎng)絡(luò)層
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]

# 創(chuàng)建特征提取模型
down_stack = tf.keras.Model(inputs=base_model.input, outputs=layers)

down_stack.trainable = False

# 進(jìn)行降頻采樣
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
]

# 定義UNet網(wǎng)絡(luò)模型
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'])

在這里，我們首先得到了一個(gè)預(yù)訓(xùn)練的 MobileNetV2 用于特征提取，在這里我們并沒有包含它的輸出層，因?yàn)槲覀円鶕?jù)自己的任務(wù)靈活調(diào)節(jié)。

然后定義了我們要使用的 MobileNetV2 的網(wǎng)絡(luò)層的輸出，我們使用這些輸出來作為我們提取的特征。

然后我們定義了我們的網(wǎng)絡(luò)模型，這個(gè)模型的理解有些困難，大家可能不用詳細(xì)了解網(wǎng)絡(luò)的具體原理。大家只需要知道，這個(gè)網(wǎng)絡(luò)大致經(jīng)過的步驟包括：

• 先將數(shù)據(jù)壓縮（便于數(shù)據(jù)的處理）；
• 然后進(jìn)行數(shù)據(jù)的處理；
• 最后將數(shù)據(jù)解壓返回到原來的大小，從而完成網(wǎng)絡(luò)的任務(wù)。

最后我們編譯該模型，我們使用 adam 優(yōu)化器，交叉熵?fù)p失函數(shù)（因?yàn)閳D像分割是個(gè)分類任務(wù)）。

5. 模型的訓(xùn)練

epoch = 20
valid_steps = info.splits['test'].num_examples//BATCH

model_history = model.fit(train_dataset, epochs=epoch,
                          steps_per_epoch=step_per_epch,
                          validation_steps=valid_steps,
                          validation_data=test_dataset)

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

這邊就是一個(gè)簡單的訓(xùn)練過程，我們可以得到如下輸出：

Epoch 1/20
57/57 [==============================] - 296s 5s/step - loss: 0.4928 - accuracy: 0.7995 - val_loss: 0.6747 - val_accuracy: 0.7758
......
Epoch 20/20
57/57 [==============================] - 276s 5s/step - loss: 0.2586 - accuracy: 0.9218 - val_loss: 0.2821 - val_accuracy: 0.9148

我們可以看到我們最后達(dá)到了 91% 的準(zhǔn)確率，還是一個(gè)可以接受的結(jié)果。

感興趣的同學(xué)可以嘗試一下進(jìn)行結(jié)果的可視化，從而更加直觀的查看到結(jié)果。

4. 小結(jié)

在這節(jié)課之中，我們學(xué)習(xí)了什么是圖像分割，同時(shí)了解了圖像分割的簡單的實(shí)現(xiàn)方式，最終我們通過一個(gè)示例來了解了如何在 TensorFlow 之中進(jìn)行圖像分割。