前言

圖像分割在醫(yī)學(xué)成像、自動(dòng)駕駛汽車和衛(wèi)星成像等方面有很多應(yīng)用，本質(zhì)其實(shí)就是圖像像素分類任務(wù)，也就是使用深度學(xué)習(xí)模型為輸入圖像的每個(gè)像素分配一個(gè)標(biāo)簽（或類）。

準(zhǔn)備

本文的準(zhǔn)備如下，使用 pip 安裝如下配置：

• pip install git+github.com/tensorflow/…
• pip install tensorflow == 2.10.1
• pip install tensorflow_datasets == 4.7.8
• pip install ipython == 8.6.0
• pip install matplotlib == 3.6.2

大綱

• 獲取數(shù)據(jù)
• 處理數(shù)據(jù)
• 搭建模型
• 編譯、訓(xùn)練模型
• 預(yù)測

實(shí)現(xiàn)

1. 獲取數(shù)據(jù)

（1）本文使用的數(shù)據(jù)集是 Oxford-IIIT Pet Dataset ,該數(shù)據(jù)集由 37 類寵物的圖像組成，每個(gè)品種有 200 個(gè)圖像（訓(xùn)練集和測試集各有 100 個(gè)），每個(gè)像素都會(huì)被劃入以下三個(gè)類別之一：

• 屬于寵物的像素
• 寵物邊緣的像素
• 其他位置的像素

（2）可以使用 TensorFlow 的內(nèi)置函數(shù)從網(wǎng)絡(luò)上下載本次使用的數(shù)據(jù) oxford_iiit_pet ，一般會(huì)下載到本地目錄 ：C:\Users\【用戶目錄】\tensorflow_datasets\oxford_iiit_pet 。

（3）dataset 中存放是訓(xùn)練集和測試集這兩個(gè)數(shù)據(jù)集，info 中存放的是該數(shù)據(jù)的基本信息，如文件大小，數(shù)據(jù)介紹等基本信息。

import tensorflow as tf
import tensorflow_datasets as tfds
from tensorflow_examples.models.pix2pix import pix2pix
from IPython.display import clear_output
import matplotlib.pyplot as plt
dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)

2. 處理數(shù)據(jù)

（1）normalize 函數(shù)主要是完成將圖像顏色值被歸一化到 [0,1] 范圍，掩碼像素的所屬標(biāo)簽被標(biāo)記為 {1, 2, 3}。為了方便后面的模型計(jì)算，將它們分別減去 1，得到的標(biāo)簽為：{0, 1, 2} 。

（2）load_image 函數(shù)主要是將每個(gè)圖片的輸入和掩碼圖片，使用指定的方法將其大小調(diào)整為指定的 128x128 。

（3）從 dataset 中分理處訓(xùn)練集 train_images 和測試集 test_images 。

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(image):
    input_image = tf.image.resize(image['image'], (128, 128))
    input_mask = tf.image.resize(image['segmentation_mask'], (128, 128))
    input_image, input_mask = normalize(input_image, input_mask)
    return input_image, input_mask
train_images = dataset['train'].map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
test_images = dataset['test'].map(load_image, num_parallel_calls=tf.data.AUTOTUNE)

（4）為了保證在加載數(shù)據(jù)的時(shí)候不會(huì)出現(xiàn) I/O 不會(huì)阻塞，我們?cè)趶拇疟P加載完數(shù)據(jù)之后，使用 cache 會(huì)將數(shù)據(jù)保存在內(nèi)存中，確保在訓(xùn)練模型過程中數(shù)據(jù)的獲取不會(huì)成為訓(xùn)練速度的瓶頸。

如果說要保存的數(shù)據(jù)量太大，可以使用 cache 創(chuàng)建磁盤緩存提高數(shù)據(jù)的讀取效率。另外我們還使用 prefetch 在訓(xùn)練過程中可以并行執(zhí)行數(shù)據(jù)的預(yù)獲取。

TRAIN_LENGTH = info.splits['train'].num_examples
BATCH_SIZE = 32
BUFFER_SIZE = 1000
train_batches = (train_images.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat().prefetch(buffer_size=tf.data.AUTOTUNE))
test_batches = test_images.batch(BATCH_SIZE)

（5）這里的 display 函數(shù)主要是將每個(gè)樣本的寵物圖像、對(duì)應(yīng)的掩碼圖像、預(yù)測的掩碼圖像繪制出來，在這里我們只隨機(jī)挑選了一個(gè)樣本進(jìn)行顯示。因?yàn)檫@里還沒有預(yù)測的掩碼圖像，所以沒有將其繪制出來。

（6）我們可以看到左側(cè)是一張寵物的生活照，右邊是一張?jiān)搶櫸镌谡掌械妮喞€圖，寵物的樣子所處的像素為紫色，寵物的輪廓邊緣線的像素是黃色，背景的像素是墨綠色，這其實(shí)對(duì)應(yīng)了圖片中的像素會(huì)分成三個(gè)類別。

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.utils.array_to_img(display_list[i]))
        plt.axis('off')
    plt.show()
for images, masks in train_batches.take(1):
    sample_image, sample_mask = images[0], masks[0]
    display([sample_image, sample_mask])

3. 搭建模型

（1）這里使用的模型是修改后的 U-Net ，詳細(xì)內(nèi)容可看鏈接。U-Net 由編碼器（下采樣器）和解碼器（上采樣器）組成。為了學(xué)習(xí)穩(wěn)健的特征并減少可訓(xùn)練參數(shù)的數(shù)量，請(qǐng)使用預(yù)訓(xùn)練模型 MobileNetV2 作為編碼器。對(duì)于解碼器，您將使用上采樣塊，該塊已在 TensorFlow Examples 倉庫的 pix2pix 示例中實(shí)現(xiàn)。

（2）如前所述，編碼器是一個(gè)預(yù)訓(xùn)練的 MobileNetV2 模型。您將使用來自 tf.keras.applications 的模型。編碼器由模型中中間層的特定輸出組成。請(qǐng)注意，在訓(xùn)練過程中不會(huì)訓(xùn)練編碼器。

（3）我們這里使用模型由兩部分組成， 一個(gè)是編碼器 down_stack（也就是下采樣器），另一個(gè)是解碼器 up_stack （也就是上采樣器）。我們這里使用預(yù)訓(xùn)練的模型 MobileNetV2 作為編碼器， MobileNetV2 模型可以直接從網(wǎng)絡(luò)上下載到本地使用，使用它來進(jìn)行圖片的特征抽取，需要注意的是我們這里選取了模型中的若干中間層，將其作為模型的輸出，而且在訓(xùn)練過程中我們?cè)O(shè)置了不會(huì)去訓(xùn)練編碼器模型中的權(quán)重。對(duì)于解碼器，我們使用已經(jīng)在倉庫實(shí)現(xiàn)了的 pix2pix 。

（4）我們的 U-Net 網(wǎng)絡(luò)接收的每張圖片大小為 [128, 128, 3] ，先通過模型進(jìn)行下采樣，然后計(jì)算上采樣和 skip 的特征連接，最后經(jīng)過一層 Conv2DTranspose 輸出一個(gè)大小為 [batch_size, 128, 128, 3] 的向量結(jié)果。

base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)
layer_names = [ 'block_1_expand_relu', 'block_3_expand_relu', 'block_6_expand_relu', 'block_13_expand_relu', 'block_16_project']
base_model_outputs = [base_model.get_layer(name).output for name in layer_names]
down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)
down_stack.trainable = False
up_stack = [  pix2pix.upsample(512, 3),  pix2pix.upsample(256, 3),   pix2pix.upsample(128, 3),   pix2pix.upsample(64, 3)]
def unet_model(output_channels:int):
    inputs = tf.keras.layers.Input(shape=[128, 128, 3])
    skips = down_stack(inputs)
    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( filters=output_channels, kernel_size=3, strides=2, padding='same')   
    x = last(x)
    return tf.keras.Model(inputs=inputs, outputs=x)

4. 編譯、訓(xùn)練模型

（1）因?yàn)槊總€(gè)像素面臨的是一個(gè)多類分類問題，所以我們使用 SparseCategoricalCrossentropy 作為損失函數(shù)，計(jì)算多分類問題的交叉熵，并將 from_logits 參數(shù)設(shè)置為 True，因?yàn)闃?biāo)簽是用 0、1、2 三個(gè)整數(shù)表示。SparseCategoricalCrossentropy 函數(shù)中當(dāng) from_logits=true 時(shí)，會(huì)先對(duì)預(yù)測值進(jìn)行 Softmax 概率化，就無須在模型最后添加 Softmax 層，我們只需要使用經(jīng)過 Softmax 輸出的小數(shù)和真實(shí)整數(shù)標(biāo)簽來計(jì)算損失即可。reduction 默認(rèn)設(shè)置為 auto 時(shí)，會(huì)對(duì)一個(gè) batch 的樣本損失值求平均。

舉例：

y_true = [0,1,2]
y_pred = [[0.2,0.5,0.3],[0.6,0.1,0.3],[0.4,0.4,0.2]]
使用函數(shù)結(jié)果：
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False,name='sparse_categorical_crossentropy')
loss_val = loss_fn(y_true,y_pred).numpy()
loss_val
1.840487
手動(dòng)計(jì)算 SparseCategoricalCrossentropy 結(jié)果：
(-np.log(0.2)-np.log(0.1)-np.log(0.2))/3
 1.8404869726207487

（2）使用 Adam 作為優(yōu)化器，使用 accuracy 作為評(píng)估指標(biāo)。

OUTPUT_CLASSES = 3
EPOCHS = 20
VAL_SUBSPLITS = 5
STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE
VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS
model = unet_model(output_channels=OUTPUT_CLASSES)
model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),  metrics=['accuracy'])
model_history = model.fit(train_batches, epochs=EPOCHS, steps_per_epoch=STEPS_PER_EPOCH, validation_steps=VALIDATION_STEPS, validation_data=test_batches)

訓(xùn)練結(jié)果輸出：

115/115 [==============================] - 110s 961ms/step - loss: 0.1126 - accuracy: 0.9473 - val_loss: 0.3694 - val_accuracy: 0.8897

5. 預(yù)測

（1）使用 create_mask 我們會(huì)將對(duì)該批次的第一張圖片的預(yù)測掩碼圖像進(jìn)行展示，結(jié)果是一個(gè)大小為 (128, 128, 1) 的向量，其實(shí)就是給出了該圖片每個(gè)像素點(diǎn)的預(yù)測標(biāo)簽。

（2）在這里我們使用了上面的一個(gè)樣本 sample_image ，使用訓(xùn)練好的模型進(jìn)行預(yù)測，因?yàn)檫@里的樣本 sample_image 是的大小是 (128, 128, 3) ，我們的模型需要加入 batch_size 維度，所以在第一維擴(kuò)展了一個(gè)維度，大小變?yōu)?(1, 128, 128, 3) 才能輸入模型。

（3）從繪制的預(yù)測掩碼圖像結(jié)果看，預(yù)測寵物邊界線已經(jīng)相當(dāng)清晰了，如果進(jìn)一步調(diào)整模型結(jié)果和訓(xùn)練的迭代次數(shù)，效果會(huì)更加好。

def create_mask(pred_mask):
    pred_mask = tf.math.argmax(pred_mask, axis=-1)
    pred_mask = pred_mask[..., tf.newaxis]
    return pred_mask[0]
display([sample_image, sample_mask,  create_mask(model.predict(sample_image[tf.newaxis, ...]))])

以上就是Tensorflow2.10實(shí)現(xiàn)圖像分割任務(wù)示例詳解的詳細(xì)內(nèi)容，更多關(guān)于Tensorflow 圖像分割的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章！