图像分割:深度指南

图像分割:全面指南

目录

  1. 介绍,动机
  2. 提取数据
  3. 可视化图像
  4. 构建简单的U-Net模型
  5. 指标和损失函数
  6. 构建完整的U-Net模型
  7. 总结
  8. 参考资料

相关链接

介绍,动机

图像分割指的是计算机(或更准确地说是存储在计算机上的模型)将图像的每个像素分配给相应的类别的能力。例如,可以通过图像分割器对上面展示的一只站在白色栅栏前的猫的图像进行分割,并得到下面的分割图像:

猫的图像,被分割成‘猫’像素和‘背景’像素。来自DALL·E 3的修改图像。

在这个例子中,我手动分割了图像。这是一项繁琐的操作,我们希望能自动化。在本指南中,我将引导您完成训练一个算法进行图像分割的过程。互联网和教科书上的许多指南在一定程度上有所帮助,但它们都没有详细介绍实现的具体细节。在这里,我将尽量详尽地介绍,以帮助您在自己的数据集上实现图像分割时节省时间。

首先,让我们将我们的任务置于更广泛的机器学习的背景下。机器学习的定义是不言而喻的:我们教会机器学习如何解决我们希望自动化的问题。人类希望自动化许多问题;在本文中,我们专注于计算机视觉中的一部分问题。计算机视觉旨在教会计算机如何看到。给一个六岁孩子一张一只站在白色栅栏前的猫的图像,并要求他将图像分割成‘猫’像素和‘背景’像素是轻而易举的(当然,在向困惑的孩子解释‘分割’是什么意思之后)。然而,几十年来,计算机在这个问题上一直苦苦挣扎。

为什么计算机要费劲地做一件六岁小孩能做的事情呢?我们可以通过想象一个人如何通过触觉来学习读写来理解计算机。想象一下,如果你拿到了一篇用盲文写成的文章,并且你对如何阅读它一无所知,你会怎么做呢?你需要怎样解读盲文并转化为英文呢?

一段用盲文写的小片段。来自Unsplash网站。

你需要的是一种将这个输入转化为可读输出的方法。在数学中,我们称之为映射。我们希望学习一个函数f(x),将不可读的输入x映射为可读的输出y。

通过几个月的练习和一位好导师,任何人都可以学会从盲文到英文的映射。类比地说,计算机处理图像就有点像一个人第一次接触盲文;它看起来像一堆胡言乱语。计算机需要学习必要的映射f(x),将对应于像素的一组数字转化为它可以用于分割图像的东西。不幸的是,计算机模型没有数千年的进化、生物学以及看到世界的经验;它在你启动程序时被“诞生”。这就是我们希望在计算机视觉中教给我们的模型的内容。

为什么我们首先要进行图像分割呢?其中一个更明显的用例是Zoom。许多人在视频会议中喜欢使用虚拟背景,以避免同事们看到他们的狗在客厅里翻跟头。图像分割对于此任务至关重要。另一个强大的用例是医学成像。在对患者器官进行CT扫描时,自动分割图像中的器官可能会有所帮助,以便医疗专业人员可以确定损伤、肿瘤的存在等等。 这是一个关于这一任务的Kaggle竞赛的优秀案例

图像分割有多种类型,从简单到复杂不等。在本文中,我们将处理最简单的图像分割类型:二元分割。这意味着只会有两种不同的对象类型,例如“猫”和“背景”。多说一句,这里提供的代码已经稍作重新排列和编辑,以便更清晰地呈现。如果想运行一些可行的代码,请参考文章顶部的代码链接。本文将使用Kaggle上的Carvana图像遮罩挑战数据集。您需要注册参加此挑战才能访问数据集,并将Kaggle API密钥插入Colab笔记本以使其工作(如果您不想使用Kaggle笔记本的话)。 请参阅此讨论帖,了解如何操作。

还有一件事;尽管我很想详细介绍代码中的所有概念,但是我假设您对卷积神经网络、最大池化层、全连接层、随机失活层和残差连接器具有一些工作知识。不幸的是,对这些概念进行详细讨论将需要一篇新文章,超出了本文所要涵盖的范围,我们将重点关注实现的实质。

提取数据

本文相关数据将存放在以下文件夹中:

  • train_hq.zip:包含高清车辆训练图像的文件夹
  • test_hq.zip:包含高清车辆测试图像的文件夹
  • train_masks.zip:包含训练集的掩码文件夹

在图像分割的语境中,掩码就是分割后的图像。我们试图让我们的模型学会如何将输入图像映射到输出的分割掩码。通常假设真正的掩码(即地面实况)是由人类专家手工绘制的。

一个由人类手绘的图像示例及其相应的真实蒙版,来自Carvana Image Masking Challenge数据集。

您的第一步将是从/kaggle/input来源解压缩文件夹:

def getZippedFilePaths():
    zip_file_names = []
    for dirname, _, filenames in os.walk('/kaggle/input'):
        for filename in filenames:
            if filename.split('.')[-1] == 'zip':
                zip_file_names.append((os.path.join(dirname, filename)))
    return zip_file_names

zip_file_names = getZippedFilePaths()
items_to_remove = ['/kaggle/input/carvana-image-masking-challenge/train.zip',
                    '/kaggle/input/carvana-image-masking-challenge/test.zip']
zip_file_names = [item for item in zip_file_names if item not in items_to_remove]

for zip_file_path in zip_file_names:
    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
        zip_ref.extractall()

此代码获取输入中所有.zip文件的文件路径,并将其解压缩到/kaggle/output目录中。注意,我故意不提取非高质量的照片;Kaggle存储库只能容纳20 GB的数据,这一步是为了防止超过此限制。

可视化图像

在大多数计算机视觉问题中,第一步是检查数据集。我们具体要处理什么?我们首先需要将图像组合成有组织的数据集进行查看。(本指南将使用TensorFlow;转换到PyTorch不应太困难。)

# 添加所有路径名到排序列表中
train_hq_dir = '/kaggle/working/train_hq/'
train_masks_dir = '/kaggle/working/train_masks/'
test_hq_dir = '/kaggle/working/test_hq/'
X_train_id = sorted([os.path.join(train_hq_dir, filename) for filename in os.listdir(train_hq_dir)], key=lambda x: x.split('/')[-1].split('.')[0])
y_train = sorted([os.path.join(train_masks_dir, filename) for filename in os.listdir(train_masks_dir)], key=lambda x: x.split('/')[-1].split('.')[0])
X_test_id = sorted([os.path.join(test_hq_dir, filename) for filename in os.listdir(test_hq_dir)], key=lambda x: x.split('/')[-1].split('.')[0])

X_train_id = X_train_id[:1000]
y_train = y_train[:1000]

X_train = tf.data.Dataset.from_tensor_slices(X_train)
y_train = tf.data.Dataset.from_tensor_slices(y_train)
X_val = tf.data.Dataset.from_tensor_slices(X_val)
y_val = tf.data.Dataset.from_tensor_slices(y_val)
X_test = tf.data.Dataset.from_tensor_slices(X_test_id)

img_height = 96
img_width = 128
num_channels = 3
img_size = (img_height, img_width)

X_train = X_train.map(preprocess_image)
y_train = y_train.map(preprocess_target)
X_val = X_val.map(preprocess_image)
y_val = y_val.map(preprocess_target)
X_test = X_test.map(preprocess_image)

train_dataset = tf.data.Dataset.zip((X_train, y_train))
val_dataset = tf.data.Dataset.zip((X_val, y_val))

BATCH_SIZE = 32
batched_train_dataset = train_dataset.batch(BATCH_SIZE)
batched_val_dataset = val_dataset.batch(BATCH_SIZE)
batched_test_dataset = X_test.batch(BATCH_SIZE)

AUTOTUNE = tf.data.experimental.AUTOTUNE
batched_train_dataset = batched_train_dataset.prefetch(buffer_size=AUTOTUNE)
batched_val_dataset = batched_val_dataset.prefetch(buffer_size=AUTOTUNE)
batched_test_dataset = batched_test_dataset.prefetch(buffer_size=AUTOTUNE)

让我们逐步来看一下:

  • 我们首先创建一个对训练集、测试集和真实蒙版的所有图像的文件路径进行排序的列表。请注意,这些还不是图像;此时我们只查看文件路径而已。
  • 然后,我们只取第一个1000个图像/蒙版的文件路径。这是为了减少计算负载并加快训练速度。如果您有多个强大的GPU(真幸运!),可以使用所有图像以获得更好的性能。我们还创建了80/20的训练/验证拆分。包含的数据(图像)越多,此拆分就越倾向于训练集。处理非常大型数据集时,训练/验证/测试拆分通常为98/1/1是再常见不过了。训练集中有更多数据(图像),您的模型在一般情况下会更好。
  • 接下来,我们使用tf.data.Dataset.from_tensor_slices()方法创建TensorFlow(TF)Dataset对象。使用Dataset对象处理训练、验证和测试集是一种常见的方法,而不是将它们保持为Numpy数组。根据我的经验,使用Dataset对象进行数据预处理速度更快、更容易。有关文档,请参阅此链接
  • 接下来,我们为输入图像指定图像高度、宽度和通道数。实际的高质量图像远远大于96像素乘以128像素;这种图像的降采样是为了减少计算负载(较大的图像需要更多训练时间)。如果您有必要的计算资源(GPU),我不建议降采样。
  • 然后,我们使用Dataset对象的.map()函数对图像进行预处理。这将文件路径转换为图像并进行适当的预处理。稍后将详细讨论这个。
  • 在预处理原始训练图像和真实蒙版之后,我们需要一种方法将图像与蒙版配对。为此,我们使用Dataset对象的.zip()函数。它接受两个数据列表,并将每个列表的第一个元素并入元组中。对于第二个元素、第三个元素等,也是如此。最终结果是一个由(image, mask)形式的元组组成的列表。
  • 然后,我们使用.batch()函数来从我们的一千个图像中创建包含32个图像的批次。批处理是机器学习流程的重要部分,它允许我们一次处理多个图像,而不是一个个处理。这加快了训练速度。
  • 最后,我们使用.prefetch()函数。这是另一个可以加速训练的步骤。加载和预处理数据可能是训练过程中的瓶颈。这可能会导致GPU或CPU空闲时间,这是没有人想要的。当您的模型执行前向和后向传播时,.prefetch()函数可以准备好下一个批次。TensorFlow中的AUTOTUNE变量会根据您的系统资源动态计算要预取多少批次;这通常是推荐的做法。

让我们更详细地了解预处理步骤:

def preprocess_image(file_path):    # 加载并解码图像    img = tf.io.read_file(file_path)    # 根据您的图像调整通道(RGB为3个通道)    img = tf.image.decode_jpeg(img, channels=3) # 返回为uint8类型    # 将像素值归一化到[0,1]    img = tf.image.convert_image_dtype(img, tf.float32)    # 将图像调整为所需的尺寸    img = tf.image.resize(img, [96, 128], method = 'nearest')    return imgdef preprocess_target(file_path):    # 加载并解码图像    mask = tf.io.read_file(file_path)    # 将值归一化到0和1之间(仅两个类别)    mask = tf.image.decode_image(mask, expand_animations=False, dtype=tf.float32)    # 仅获取第3个通道的一个值    mask = tf.math.reduce_max(mask, axis=-1, keepdims=True)    # 将图像调整为所需的尺寸    mask = tf.image.resize(mask, [96, 128], method = 'nearest')    return mask

这些函数的作用如下:

  • 首先,我们使用tf.io.read_file()将文件路径转换为‘字符串’类型的张量。张量是TensorFlow中的一种特殊数据结构,类似于其他数学库中的多维数组,但具有对深度学习有用的特殊属性和方法。引用TensorFlow文档的话来说:tf.io.read_file()“不进行任何解析,它只返回原样内容。”这基本上意味着它以字符串类型返回包含图像信息的二进制文件(1和0)。
  • 其次,我们需要解码二进制数据。为此,我们需要在TensorFlow中使用适当的方法。由于原始图像数据采用.jpeg格式,因此我们使用tf.image.decode_jpeg()方法。由于掩码是GIF格式,我们可以使用tf.io.decode_gif(),或者使用更常规的tf.image.decode_image(),它可以处理任何文件类型。您选择哪个都无关紧要。我们设置expand_animations=False,因为这些实际上不是动画,只是图像。
  • 然后我们使用convert_image_dtype()将图像数据转换为float32类型。这仅适用于图像,而不适用于掩码,因为掩码已经被解码为float32类型。在图像处理中使用的两种常见数据类型是float32和uint8。Float32代表在计算机内存中占用32位空间的浮点数(十进制数)。它们是有符号的(意味着数字可以是负数),其值的范围可从0到2³²=4294967296,尽管根据图像处理的约定,我们会将这些值归一化为0到1之间,其中1是颜色的最大值。Uint8代表无符号(正数)整数,取值范围在0到255之间,且在内存中仅占用8位。例如,我们可以将燃橙色表示为Uint8为(红:204,绿:85,蓝:0),或者为float32为(红:0.8,绿:0.33,蓝:0)。由于float32提供更高的精度并已被归一化,通常它是更好的选择。然而,uint8节省内存,根据内存限制,这可能是更好的选择。在convert_image_dtype中使用float32会自动归一化值。
  • 在二进制分割中,我们希望掩码的形状为(批次、高度、宽度、通道),其中通道为1。换句话说,我们希望数字1表示一个类(汽车),数字0表示另一个类(背景)。没有必要将通道数设置为3,如RGB图像那样。不幸的是,在解码之后,掩码具有三个通道,且类别号重复三次。为了解决这个问题,我们使用tf.math.reduce_max(mask, axis=-1, keepdims=True)来取得三个通道中的最大值,并去除其余部分。因此,通道值(1,1,1)被缩减为(1),通道值(0,0,0)被缩减为(0)。
  • 最后,我们将图像/掩码调整为所需的尺寸(较小)。请注意,我之前展示的包含汽车和真实掩码的图像看起来模糊;这是有意为之,目的是减少计算负载并使训练相对较快。使用method=’nearest’作为默认值是个好主意;否则,该函数将始终返回float32类型,如果您希望它是uint8类型,则不好。
颜色“燃烧的橙色”可以用float32或uint8格式表示。作者提供图片。

我们整理好数据集之后,现在可以查看我们的图片:

# 查看图片以及对应的标签for images, masks in batched_val_dataset.take(1):    car_number = 0    for image_slot in range(16):        ax = plt.subplot(4, 4, image_slot + 1)        if image_slot % 2 == 0:            plt.imshow((images[car_number]))             class_name = '图片'        else:            plt.imshow(masks[car_number], cmap = 'gray')            plt.colorbar()            class_name = '掩膜'            car_number += 1                    plt.title(class_name)        plt.axis("off")
我们的汽车图片与相应的掩膜。

在这里,我们使用.take()方法来查看我们的批次数据集中的第一个批次。由于我们进行的是二进制分割,我们希望我们的掩膜只包含两个值:0和1。通过在掩膜imshow()中添加参数cmap =’gray’,我们可以确认我们设置正确。这表示我们希望以灰度形式呈现这些图片。

建立一个简单的U-Net模型

在1675年2月5日的信中,艾萨克·牛顿对他的竞争对手罗伯特·胡克说:

“如果我看得更远,那是因为我站在巨人的肩膀上。”

同样,我们将借鉴先前的机器学习研究人员的经验,他们发现了哪种架构对于图像分割任务最有效。尝试使用自己的架构并没有错,但是在我们之前的研究人员中,已经走了很多弯路才发现有效的模型。这些架构并不一定是铁板钉钉的最终选择,因为研究仍在进行中,可能还会找到更好的架构。

U-Net的可视化,见[1]。作者提供图片。

其中一个较为著名的架构称为U-Net,因为网络的下采样和上采样部分可以可视化为一个U形(如下图所示)。在Ronneberger、Fisher和Brox的一篇名为《U-Net: 用于生物医学图像分割的卷积网络》的论文[1]中,作者描述了如何创建一个有效的全卷积网络(FCN),用于图像分割。全卷积意味着没有密集连接的层,所有层都是卷积层。

有几点需要注意:

  • 网络由一系列重复的两个卷积层组成,填充为“same”,步幅为1,以便区块内的卷积输出不会缩小。
  • 每个区块后面是一个最大池化层,将特征图的宽度和高度减半。
  • 下一个区块将过滤器的数量加倍。这种减小特征空间的同时增加过滤器数量的模式,如果您学过CNN,应该很熟悉。这完成了作者所称的“收缩路径(contracting path)”。
  • “瓶颈层”位于‘U’形的底部。这一层捕获高度抽象的特征(线条、曲线、窗户、门等),但在显著降低的空间分辨率下。
  • 接下来开始“扩展路径(expanding path)”。简单来说,这相当于反向进行了收缩操作,每个区块再次由两个卷积层组成。每个区块后面是一个上采样层,我们在TensorFlow中称之为Conv2DTranspose层。这将较小的特征图的高度和宽度加倍。
  • 下一个区块将过滤器数量减半。重复此过程,直到最终获得与起始图像相同的高度和宽度。最后,使用1×1卷积层将通道数量减少为1。我们希望最终只有一个通道,因为这是二进制分割,所以我们希望一个像素值对应于我们的两个类别。我们使用sigmoid激活函数将像素值压缩在0和1之间。
  • U-Net架构还包括“跳跃连接(skip connections)”,允许网络在下采样和上采样后保留细粒度的空间信息。通常在这个过程中会丢失大量信息。通过将来自收缩块的信息传递到相应的扩展块,可以保留这些空间信息。架构具有良好的对称性。

我们将从一个简单的U-Net版本开始。这将是一个FCN,但没有残差连接和最大池化层。

data_augmentation = tf.keras.Sequential([        tfl.RandomFlip(mode="horizontal", seed=42),        tfl.RandomRotation(factor=0.01, seed=42),        tfl.RandomContrast(factor=0.2, seed=42)])def get_model(img_size):    inputs = Input(shape=img_size + (3,))    x = data_augmentation(inputs)        # 收缩路径    x = tfl.Conv2D(64, 3, strides=2, activation="relu", padding="same", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(64, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(128, 3, strides=2, activation="relu", padding="same", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(128, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(256, 3, strides=2, padding="same", activation="relu", kernel_initializer='he_normal')(x)     x = tfl.Conv2D(256, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)        # 扩展路径    x = tfl.Conv2DTranspose(256, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)    x = tfl.Conv2DTranspose(256, 3, activation="relu", padding="same", kernel_initializer='he_normal', strides=2)(x)    x = tfl.Conv2DTranspose(128, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)    x = tfl.Conv2DTranspose(128, 3, activation="relu", padding="same", kernel_initializer='he_normal', strides=2)(x)    x = tfl.Conv2DTranspose(64, 3, activation="relu", padding="same", kernel_initializer='he_normal')(x)    x = tfl.Conv2DTranspose(64, 3, activation="relu", padding="same", kernel_initializer='he_normal', strides=2)(x)    outputs = tfl.Conv2D(1, 3, activation="sigmoid", padding="same")(x)    model = keras.Model(inputs, outputs)         return modelcustom_model = get_model(img_size=img_size)

这里我们有与U-Net相同的基本结构,包括一个收缩路径和一个扩展路径。一个有趣的事情要注意的是,我们使用strides=2的卷积层来将特征空间减半,而不是使用最大池化层。根据Chollet [2]的说法,与破坏性最大池化层相比,使用带有strided convolutions可以保留更多的空间信息。他指出,在位置信息重要的时候(如图像分割),避免破坏性最大池化层并坚持使用strided convolutions是一个好主意(这很奇怪,因为著名的U-Net架构确实使用了最大池化层)。还要注意我们正在进行一些数据增强,以帮助模型推广到看不见的示例。

一些重要的细节:将核初始化器设置为’ReLU’激活函数的’he_normal’设置在训练稳定性方面产生了令人惊讶的大差异。我最初低估了核初始化的能力。he_normalization不像随机初始化权重那样,它将权重初始化为均值为0,标准差为(2 /上一层的输入单元数的平方根)。对于CNN,输入单元的数量是指上一层特征映射中的通道数。已经发现这样做可以加快收敛速度,减轻梯度消失问题,并改善学习效果。更多详情请参见参考[3]。

指标和损失函数

对于二进制分割,有几种常见的指标和损失函数可以使用。这里,我们将使用 dice系数作为指标,并使用相应的 dice损失进行训练,因为这是竞赛要求的。

让我们首先看看dice系数背后的数学原理:

通用形式下的dice系数。

骰子系数被定义为两个集合(X和Y)的交集除以每个集合的总和乘以2。骰子系数的取值范围在0(如果集合没有交集)和1(如果集合完全重叠)之间。现在我们知道为什么它成为了图像分割的一个绝佳度量。

两个遮罩图像重叠的示例。为了清晰起见,使用了橙色。图片由作者提供。

上述方程是骰子系数的一般定义;当应用于矢量数量(而不是集合)时,我们使用更具体的定义:

矢量形式的骰子系数。

这里,我们对每个遮罩中的每个元素(像素)进行迭代。x代表预测遮罩中的第i个像素,y代表对应的地面实况遮罩中的像素。在顶部,我们进行逐元素乘法,在底部,我们分别对每个遮罩中的所有元素求和。N代表像素的总数(对于预测和目标遮罩应该是相同的)。请记住,在我们的遮罩中,数字要么是0,要么是1,因此地面实况遮罩中值为1的像素与预测遮罩中对应的值为0的像素不会对骰子分数产生贡献,正如预期的那样(1 x 0 = 0)。

骰子损失将被简单定义为1 – 骰子分数。由于骰子分数介于0和1之间,骰子损失也介于0和1之间。事实上,骰子分数和骰子损失的和必须等于1。它们是反相关的。

让我们来看看如何在代码中实现:

from tensorflow.keras import backend as Kdef dice_coef(y_true, y_pred, smooth=10e-6):    y_true_f = K.flatten(y_true)    y_pred_f = K.flatten(y_pred)    intersection = K.sum(y_true_f * y_pred_f)    dice = (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)    return dicedef dice_loss(y_true, y_pred):    return 1 - dice_coef(y_true, y_pred)

在这里,我们将两个4-D遮罩(批次、高度、宽度、通道=1)展平为1维向量,并计算批次中所有图像的骰子分数。请注意,我们在分子和分母中添加了平滑值,以防止两个遮罩没有重叠时出现0/0的问题。

最后,我们开始训练。我们正在采用早停法来防止过拟合,并保存最佳模型。

custom_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001,                                                        epsilon=1e-06),                                                         loss=[dice_loss],                                                         metrics=[dice_coef])callbacks_list = [    keras.callbacks.EarlyStopping(        monitor="val_loss",        patience=2,    ),    keras.callbacks.ModelCheckpoint(        filepath="best-custom-model",        monitor="val_loss",        save_best_only=True,    )]history = custom_model.fit(batched_train_dataset, epochs=20,                    callbacks=callbacks_list,                    validation_data=batched_val_dataset)

我们可以使用以下代码确定训练的结果:

def display(display_list):    plt.figure(figsize=(15, 15))    title = ['输入图像', '真实遮罩', '预测遮罩']    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()    def create_mask(pred_mask):    mask = pred_mask[..., -1] >= 0.5    pred_mask[..., -1] = tf.where(mask, 1, 0)    # 仅返回批次的第一个遮罩    return pred_mask[0]def show_predictions(model, dataset=None, num=1):    """    显示前num个批次中的第一张图像    """    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, ...]))])custom_model = keras.models.load_model("/kaggle/working/best-custom-model", custom_objects={'dice_coef': dice_coef, 'dice_loss': dice_loss})show_predictions(model = custom_model, dataset = batched_train_dataset, num = 6)

经过10个epoch,我们得到了一个最高验证骰子得分为0.8788。不算太糟糕,但也不算很好。在P100 GPU上,我大约花了20分钟。以下是我们的复习样本的掩码:

输入图像,真实掩码和预测掩码的比较。作者提供。

强调一些有趣的点:

  • 请注意create_mask是将像素值推送到0或1的函数。像素值小于0.5将被截断为0,并将该像素分配给“背景”类别。值≥0.5将增加到1,并将该像素分配给“汽车”类别。
  • 为什么掩码变成了黄色和紫色,而不是黑色和白色?我们使用:tf.keras.preprocessing.image.array_to_img()将掩码的输出从张量转换为PIL图像。然后将图像传递给plt.imshow()。来自文档我们可以看到单通道图像的默认色彩映射是“viridis”(3通道RGB图像直接输出)。viridis色彩映射将低值转换为深紫色,高值转换为黄色。这种色彩映射显然可以帮助色盲人准确地观察图像中的颜色。我们可以通过添加cmap=“grayscale”来修复这个问题,但这会破坏我们的输入图像。在此链接上了解更多信息。
viridis色彩映射,从低值(紫色)到高值(黄色)。作者提供。

构建完整的U-Net

现在我们转向使用完整的U-Net架构,包括残差连接、最大池化层和包括丢弃层进行正则化。请注意收缩路径、瓶颈层和扩展路径。丢弃层可以添加在收缩路径的每个块的末尾。

def conv_block(inputs=None, n_filters=64, dropout_prob=0, max_pooling=True):    conv = Conv2D(n_filters,                    3,                     activation='relu',                  padding='same',                  kernel_initializer='he_normal')(inputs)    conv = Conv2D(n_filters,                    3,                     activation='relu',                  padding='same',                  kernel_initializer='he_normal')(conv)    if dropout_prob > 0:        conv = Dropout(dropout_prob)(conv)    if max_pooling:        next_layer = MaxPooling2D(pool_size=(2, 2))(conv)    else:        next_layer = conv    skip_connection = conv    return next_layer, skip_connectiondef upsampling_block(expansive_input, contractive_input, n_filters=64):    up = Conv2DTranspose(        n_filters,            3,            strides=(2, 2),        padding='same',        kernel_initializer='he_normal')(expansive_input)    #合并前一个输出和收缩输入    merge = concatenate([up, contractive_input], axis=3)    conv = Conv2D(n_filters,                     3,                       activation='relu',                  padding='same',                  kernel_initializer='he_normal')(merge)    conv = Conv2D(n_filters,                     3,                       activation='relu',                  padding='same',                  kernel_initializer='he_normal')(conv)    return convdef unet_model(input_size=(96, 128, 3), n_filters=64, n_classes=1):    inputs = Input(input_size)        inputs = data_augmentation(inputs)    #收缩路径(编码)    cblock1 = conv_block(inputs, n_filters)    cblock2 = conv_block(cblock1[0], n_filters*2)    cblock3 = conv_block(cblock2[0], n_filters*4)    cblock4 = conv_block(cblock3[0], n_filters*8, dropout_prob=0.3)    #瓶颈层    cblock5 = conv_block(cblock4[0], n_filters*16, dropout_prob=0.3, max_pooling=False)        #扩展路径(解码)    ublock6 = upsampling_block(cblock5[0], cblock4[1],  n_filters*8)    ublock7 = upsampling_block(ublock6, cblock3[1],  n_filters*4)    ublock8 = upsampling_block(ublock7, cblock2[1],  n_filters*2)    ublock9 = upsampling_block(ublock8, cblock1[1],  n_filters)    conv9 = Conv2D(n_filters,                   3,                   activation='relu',                   padding='same',                   kernel_initializer='he_normal')(ublock9)    conv10 = Conv2D(n_classes, 1, padding='same', activation="sigmoid")(conv9)    model = tf.keras.Model(inputs=inputs, outputs=conv10)    return model

然后我们开始编译U-Net。我在第一个卷积块中使用了64个滤波器。这是一个超参数,您会希望调整以获得最佳结果。

unet = unet_model(input_size=(img_height, img_width, num_channels), n_filters=64, n_classes=1)unet.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001, epsilon=1e-06),             loss=[dice_loss],              metrics=[dice_coef])callbacks_list = [    keras.callbacks.EarlyStopping(        monitor="val_loss",        patience=2,    ),    keras.callbacks.ModelCheckpoint(        filepath="best-u_net-model",        monitor="val_loss",        save_best_only=True,    )]history = unet.fit(batched_train_dataset, epochs=20,                    callbacks=callbacks_list,                    validation_data=batched_val_dataset)

经过16个epochs,我得到了0.9416的验证dice分数,比简单的U-Net要好得多。这应该不会太令人惊讶;从参数计数方面来看,我们从简单的U-Net到完整的U-Net有数量级的增加。在P100 GPU上,我花费了大约32分钟。然后我们来看一下预测结果:

unet = keras.models.load_model("/kaggle/working/best-u_net-model", custom_objects={'dice_coef': dice_coef, 'dice_loss': dice_loss})show_predictions(model = unet, dataset = batched_train_dataset, num = 6)
Predicted mask for the complete U-Net. Much better! By the author.

这些预测结果要好得多。从观察多个预测结果来看,车辆中突出的天线对网络来说是个难题。考虑到这些图像的像素化程度很高,我不能怪网络错过这个。

为了提高性能,我们需要调整一些超参数,包括:

  • 下采样和上采样块的数量
  • 滤波器的数量
  • 图像分辨率
  • 训练集的大小
  • 损失函数(可能将dice loss与二元交叉熵损失相结合)
  • 调整优化器参数。训练稳定性似乎是两个模型都存在的问题。根据Adam优化器的文档:“ε的默认值1e-7在一般情况下可能不是一个好的默认值。”将ε增加一个数量级或更多可能有助于训练的稳定性。

我们已经可以看到在Carvana挑战中取得出色成绩的道路。可惜活动已经结束了!

总结

本文对图像分割的主题进行了深入探讨,特别是二值分割。如果您想从中获得一些收获,请记住以下内容:

  • 图像分割的目标是找到从输入图像的像素值到模型可以使用以为每个像素分配类别的输出数字的映射。
  • 首先的步骤之一是将图像整理成TensorFlow Dataset对象,并查看图像及其对应的掩码。
  • 在模型架构方面,没有必要重新发明轮子:我们知道U-Net效果很好。
  • dice分数是一种常用的用于监测模型预测成功的度量标准。我们还可以从中获取我们的损失函数。

未来的工作可将规范U-Net架构中的最大池化层转换为步幅卷积层。

祝您在图像分割问题上好运!

参考文献

[1] O. Ronneberger, P. Fischer, and T. Brox, U-Net: Convolutional Networks for Biomedical Image Segmentation (2015), MICCAI 2015 International Conference

[2] F. Chollet, Deep Learning with Python (2021), Manning Publications Co.

[3] K. He, X. Zhang, S. Ren, J. Sun, Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification (2015), International Conference on Computer Vision (ICCV)