Tensorflow2.10實現(xiàn)圖像分割任務示例詳解

更新時間：2023年01月05日 15:11:34 作者：我是王大你是誰

這篇文章主要為大家介紹了Tensorflow2.10實現(xiàn)圖像分割任務示例詳解，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進步，早日升職加薪

前言

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

準備

本文的準備如下，使用 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ù)
搭建模型
編譯、訓練模型
預測

實現(xiàn)

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

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

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

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

（3）dataset 中存放是訓練集和測試集這兩個數(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] 范圍，掩碼像素的所屬標簽被標記為 {1, 2, 3}。為了方便后面的模型計算，將它們分別減去 1，得到的標簽為：{0, 1, 2} 。

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

（3）從 dataset 中分理處訓練集 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ù)的時候不會出現(xiàn) I/O 不會阻塞，我們在從磁盤加載完數(shù)據(jù)之后，使用 cache 會將數(shù)據(jù)保存在內存中，確保在訓練模型過程中數(shù)據(jù)的獲取不會成為訓練速度的瓶頸。

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

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ù)主要是將每個樣本的寵物圖像、對應的掩碼圖像、預測的掩碼圖像繪制出來，在這里我們只隨機挑選了一個樣本進行顯示。因為這里還沒有預測的掩碼圖像，所以沒有將其繪制出來。

（6）我們可以看到左側是一張寵物的生活照，右邊是一張該寵物在照片中的輪廓線圖，寵物的樣子所處的像素為紫色，寵物的輪廓邊緣線的像素是黃色，背景的像素是墨綠色，這其實對應了圖片中的像素會分成三個類別。

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 ，詳細內容可看鏈接。U-Net 由編碼器（下采樣器）和解碼器（上采樣器）組成。為了學習穩(wěn)健的特征并減少可訓練參數(shù)的數(shù)量，請使用預訓練模型 MobileNetV2 作為編碼器。對于解碼器，您將使用上采樣塊，該塊已在 TensorFlow Examples 倉庫的 pix2pix 示例中實現(xiàn)。

（2）如前所述，編碼器是一個預訓練的 MobileNetV2 模型。您將使用來自 tf.keras.applications 的模型。編碼器由模型中中間層的特定輸出組成。請注意，在訓練過程中不會訓練編碼器。

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

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

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. 編譯、訓練模型

（1）因為每個像素面臨的是一個多類分類問題，所以我們使用 SparseCategoricalCrossentropy 作為損失函數(shù)，計算多分類問題的交叉熵，并將 from_logits 參數(shù)設置為 True，因為標簽是用 0、1、2 三個整數(shù)表示。SparseCategoricalCrossentropy 函數(shù)中當 from_logits=true 時，會先對預測值進行 Softmax 概率化，就無須在模型最后添加 Softmax 層，我們只需要使用經過 Softmax 輸出的小數(shù)和真實整數(shù)標簽來計算損失即可。reduction 默認設置為 auto 時，會對一個 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ù)結果：
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
手動計算 SparseCategoricalCrossentropy 結果：
(-np.log(0.2)-np.log(0.1)-np.log(0.2))/3
1.8404869726207487

（2）使用 Adam 作為優(yōu)化器，使用 accuracy 作為評估指標。

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)

訓練結果輸出：

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

5. 預測

（1）使用 create_mask 我們會將對該批次的第一張圖片的預測掩碼圖像進行展示，結果是一個大小為 (128, 128, 1) 的向量，其實就是給出了該圖片每個像素點的預測標簽。

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

（3）從繪制的預測掩碼圖像結果看，預測寵物邊界線已經相當清晰了，如果進一步調整模型結果和訓練的迭代次數(shù)，效果會更加好。

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, ...]))])