基于 U-Net 的图像分割

点击下方卡片，关注“小白玩转Python”公众号

图像分割是一种将图像划分为不同区域或对象的过程。它通常在像素级别进行，通过将图像中具有相似特征的区域分组或定义对象的边界来完成。这是一种识别和解析图像中不同对象或特征的方法。

假设一位医学专业人士正在检查脑部扫描图像，试图找到潜在的癌性病变。这就是图像分割发挥作用的地方。分割过程用于识别图像中的不同组织和结构，在区分癌细胞和其他正常组织方面发挥着重要作用。

例如，在下面的脑部扫描图像中，分割已经识别出癌性肿瘤并以不同的颜色显示。

尽管 U-Net 专注于生物医学图像，但其灵活的架构允许它有效地用于其他类型的图像数据。

U-Net 的命名是因为它的结构类似于字母 U，如图所示。我们在输出端得到分割后的输入图像。U-Net 的架构是独特的，因为它由收缩路径和扩展路径组成。

收缩路径（编码器）从输入图像中提取属性图，而扩展路径（解码器）将这些属性转换回更高分辨率的形式。跳跃连接允许低级和高级属性结合，从而实现更好的分割性能。

今天，我们将使用 U-Net 架构对 Kvasir 数据集中的图像进行分割。该数据集是在挪威 Vestre Viken 健康基金会（VV）使用内窥镜设备收集的。数据集已被经验丰富的内窥镜医生标记。您可以在此处访问数据集。我们首先导入所需的库。

import os
import cv2
import random
import numpy as np


import matplotlib.pyplot as plt
%matplotlib inline


from sklearn.model_selection import train_test_split
from skimage.io import imread, imshow
from skimage.transform import resize
from skimage.color import rgb2gray




import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import backend as K
from tensorflow.keras.models import Model, load_model, save_model
from tensorflow.keras.optimizers import Adam, Adamax
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
from tensorflow.keras.layers import Input, Activation, BatchNormalization, Dropout, Lambda, Conv2D, Conv2DTranspose, MaxPooling2D, concatenate


# Ignore Warnings
import warnings
warnings.filterwarnings("ignore")

然后我们为 U-Net 输入图像定义特定大小的变量。

IMG_CHANNELS = 3
IMG_WIDTH = 256
IMG_HEIGHT = 256

由于我们在 kaggle 上运行我们的笔记本，我们以以下方式访问我们的数据并获取它们的 id，即图像的名称。

images_path = "/kaggle/input/kvasir-dataset-for-classification-and-segmentation/kvasir-seg/Kvasir-SEG/images"
mask_path = "/kaggle/input/kvasir-dataset-for-classification-and-segmentation/kvasir-seg/Kvasir-SEG/masks"


img_ids = next(os.walk(images_path))[2]
mask_ids = next(os.walk(mask_path))[2]

对于变量 img_ids 和 mask_ids，您应该得到以下输出：

X = np.zeros((len(img_ids), 256, 256, 3), dtype=np.uint8)
y = np.zeros((len(mask_ids), 256, 256, 1), dtype=np.bool_)

以这些零开头的数据结构是数据预处理步骤的一部分。实际的图像和标签数据随后将填充这些结构，并用于模型训练。这个预处理步骤是清洁、组织和准备数据集以训练模型的过程的一部分。

输入图像（X）：创建一个数组来存储输入图像。在这个数组中，每张图像的大小为 256x256 像素，并且有 3 个颜色通道（RGB）。然而，最初这些图像的内容未指定，它们都填充了零。
目标标签（y）：目标标签代表模型应该学习的正确的输出。这个数组也包含具有 256x256 像素大小的掩码，但只有一个颜色通道（黑白）。最初，这些掩码也填充了零。

现在我们将创建我们的输入图像以提供给我们的模型。

for n,img in enumerate(os.listdir(images_path)):
    file_path = os.path.join(images_path, img)
    image = imread(file_path)
    image = resize(image, (256, 256), mode="constant", preserve_range=True)
    X[n] = image


for n,mask in enumerate(os.listdir(mask_path)):
    file_path = os.path.join(mask_path, mask)
    mask = imread(file_path)
    mask = rgb2gray(mask)
    mask = np.expand_dims(resize(mask, (256, 256), mode="constant", preserve_range=True), axis=-1)
    y[n] = mask

这个脚本从给定目录（images_path）获取图像文件。然后它将每张图像调整到给定的大小（256x256 像素）并保存在 numpy 数组（X）中。掩码也执行相同的过程。掩码也被调整到相同的大小并保存在 numpy 数组（y）中。结果，数组 X 包含处理后的图像，而数组 y 包含处理后的掩码。让我们看一个示例图像及其掩码：

使用 train_test_split 函数，我们自动化了分割和拆分用于训练模型和评估其性能的数据集的过程。这允许模型使用训练数据进行训练，然后使用拆分的测试数据评估其性能。

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)


print(f"X_train.shape: {X_train.shape}\nY_train.shape: {y_train.shape}")

我们正在编写我们的 U-Net 模型架构。

input = tf.keras.layers.Input((256, 256, 3))


# reduce image config
s = tf.keras.layers.Lambda(lambda x: x / 256)(input)


# block 1
c1 = tf.keras.layers.Conv2D(16, (3, 3), activation="relu", kernel_initializer='he_normal', padding='same')(s)
c1 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c1)
c1 = tf.keras.layers.Dropout(0.1)(c1)
c1 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c1)
p1 = tf.keras.layers.MaxPooling2D((2, 2))(c1)


# block 2
c2 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p1)
c2 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c2)
c2 = tf.keras.layers.Dropout(0.1)(c2)
c2 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c2)
p2 = tf.keras.layers.MaxPooling2D((2, 2))(c2)


# block 3
c3 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p2)
c3 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c3)
c3  = tf.keras.layers.Dropout(0.2)(c3)
c3 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c3)
p3 = tf.keras.layers.MaxPooling2D((2, 2))(c3)


# block 4
c4 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p3)
c4 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c4)
c4 = tf.keras.layers.Dropout(0.2)(c4)
c4 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c4)
p4 = tf.keras.layers.MaxPooling2D((2, 2))(c4)


# block 5
c5 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p4)
c5 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c5)
c5 = tf.keras.layers.Dropout(0.3)(c5)
c5 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c5)
p5 = tf.keras.layers.MaxPooling2D((2, 2))(c5)


# block 6
c6 = tf.keras.layers.Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(p5)
c6 = tf.keras.layers.Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c6)
c6 = tf.keras.layers.Dropout(0.3)(c6)
c6 = tf.keras.layers.Conv2D(512, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c6)


# block 7 - back 1
u7 = tf.keras.layers.Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(c6)
u7 = tf.keras.layers.concatenate([u7, c5])
c7 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u7)
c7 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c7)
c7 = tf.keras.layers.Dropout(0.3)(c7)
c7 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c7)


# block 8 - back 2
u8 = tf.keras.layers.Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c7)
u8 = tf.keras.layers.concatenate([u8, c4])
c8 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u8)
c8 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c8)
c8 = tf.keras.layers.Dropout(0.2)(c8)
c8 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c8)


# block 9 - back 3
u9 = tf.keras.layers.Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(c8)
u9 = tf.keras.layers.concatenate([u9, c3])
c9 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u9)
c9 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c9)
c9 = tf.keras.layers.Dropout(0.2)(c9)
c9 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c9)


# block 10 - back 4
u10 = tf.keras.layers.Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(c9)
u10 = tf.keras.layers.concatenate([u10, c2])
c10 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u10)
c10 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c10)
c10 = tf.keras.layers.Dropout(0.1)(c10)
c10 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c10)


# block 11 - back 5
u11 = tf.keras.layers.Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same')(c10)
u11 = tf.keras.layers.concatenate([u11, c1])
c11 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u11)
c11 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c11)
c11 = tf.keras.layers.Dropout(0.1)(c11)
c11 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c11)


outputs = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid')(c11)


model = tf.keras.Model(inputs=input, outputs=outputs, name='U-NET')


model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy','iou_coef'])

Adam 优化器是一种在训练期间更新权重的优化算法。Adam 使用自适应矩估计，这在处理大型数据集和复杂模型时特别有效。

我们使用 binary_crossentropy 作为损失函数。binary_crossentropy 是一种常用于二元分类问题的损失函数。这个损失函数通过计算两个类别之间的差异，并通过比较模型的输出与实际标签来计算误差，优化了模型的学习过程。在这里我们执行像素级分割。我们检查图像中的每个像素是否属于一个类别。在我们的问题中，有两个类别：背景和对象。因此，我们更喜欢使用 binary_crossentropy 损失函数。

在我们的评估指标中，我们看到 iou_coef 指标。这个指标帮助我们评估我们在分割问题中的成功程度。iou_coef 也被称为交集比并集（IoU）系数，衡量模型预测的分割结果与实际分割结果的匹配程度。这个指标越高，模型的分割性能就越好。因此，IoU 系数是确定模型执行分割的准确性的重要指标。

现在我们在代码中包括 iou_score：

def iou_coef(y_true, y_pred, smooth=100):
    intersection = K.sum(y_true * y_pred)
    sum = K.sum(y_true + y_pred)
    iou = (intersection + smooth) / (sum - intersection + smooth)
    return iou

我们开始训练。

model.fit(X_train, y_train, validation_split=0.1, batch_size=8, epochs=100)

作为模型的结果，我们看到训练阶段的评估指标：

Epoch 100/100
76/76 [==========================] - 5s 60ms/step - loss: 0.0473 - accuracy: 0.9813 - iou_coef: 0.8491 -val_loss: 0.3569 - val_accuracy: 0.9017 - val_iou_coef: 0.5388
<keras.callbacks.History at 0x7e9c514b6290>

iou_coef 的值在 0 和 1 之间，越接近 1，模型的性能越好。0.8496 的 IOU 值意味着模型的预测与实际分割掩码很好地重叠。这意味着模型的预测相当准确地识别了真实图像中的对象。现在，我们进入模型将在测试数据上进行预测的阶段：

ind = random.randint(0, len(X_test))
img = X_test[ind]
predictions = model.predict(np.expand_dims(img, axis=0), verbose=0)
plt.figure(figsize=(15, 12))
plt.subplot(1, 3, 1)
plt.title("original image")
plt.imshow(np.squeeze(img))
plt.subplot(1, 3, 2)
plt.title("predicted mask")
plt.imshow(np.squeeze(predictions))
plt.subplot(1, 3, 3)
plt.imshow(np.squeeze(img))
# plt.imshow(msk,alpha=0.5)
plt.imshow(np.squeeze(predictions), alpha=0.5)
plt.show()

让我们看看我们的结果：