深度学习实战（十）：使用 PyTorch 进行 3D 医学图像分割

MedicalZoo论文： Deep learning in medical image analysis: a comparative analysis of multi-modal brain-MRI segmentation with 3D deep neural networks
代码已开源：MedicalZooPytorch

1. 项目简介

  深度网络在计算机视觉领域的兴起，为经典的图像处理技术表现不佳的问题提供了最先进的解决方案。在泛化的图像识别任务中，包括物体检测、图像分类和分割、活动识别、光流和姿势估计等问题，可以轻松地声称DNN（深度神经网络）已经取得了卓越的性能。

  伴随着计算机视觉的兴起，人们对其在医学成像领域的应用产生了浓厚的兴趣。尽管医学成像数据不是那么容易获得，但DNN似乎是为这种复杂的高维数据建模的理想候选者。

  最近，伦敦帝国学院开设了COVID-19课程。很多研究已经尝试通过三维CT扫描的深度网络自动检测COVID-19。尽管如此，具体的应用数据仍然不可用。很明显，人工智能将通过医学影像对医学的发展产生巨大的影响。

  正如我将看到的，医学图像通常是三或四维的。这个领域吸引大量关注的另一个原因是它对人类生命的直接影响。在美国，医疗事故是继心脏病和癌症之后的第三大死因。因此，很明显，人类死亡的前三个原因都与医学成像有关。这就是为什么估计到2023年，人工智能和深度学习在医学成像领域将创造一个超过10亿美元的全新市场。

  这项工作作为这两个世界的交叉点：深度神经网络和医学成像。在这篇文章中，我将解决医学图像分割的问题，重点是磁共振图像，这是最受欢迎的任务之一，因为它是有人能获得的结构最完善的数据集的任务。由于在线医疗数据收集并不像听起来那么简单；文章的结尾提供了一个链接集，以开始你的旅程。

  本文介绍了一个正在开发的开源库的一些初步结果，该库名为MedicalZoo，可以在这里找到。

2. 3D医学图像分割的需求

  医学图像中的三维体积图像分割对于诊断、监测和治疗计划是强制性的。我将只使用磁共振图像（MRI）。人工操作需要解剖学知识，而且它们既昂贵又费时。另外，由于人为因素，它们可能是不准确的。然而，自动体积分割可以节省医生的时间，并为进一步分析提供准确的可重复的解决方案。

  我将首先描述MR成像的基本原理，因为了解你的输入数据对训练一个深度架构至关重要。然后，为读者提供一个可以有效地用于这项任务的3D-UNET的概述。

3. 医学图像和MRI

  医学成像试图揭示被皮肤和骨骼隐藏的内部结构，以及诊断和治疗疾病。医学磁共振（MR）成像使用氢原子核的信号来生成图像。就氢原子核而言：当它暴露在外部磁场中，表示为B0，磁矩，或自旋，像罗盘针一样与磁场方向对齐。

  所有的恒定磁化被一个额外的射频脉冲旋转到另一个平面，该脉冲足够强，应用时间足够长，以使磁化倾斜。激发之后，磁化就立即在另一个平面上旋转起来。旋转的磁化在接收线圈中产生了MR信号。然而，由于两个独立的过程，MR信号迅速消失，这两个过程减少了磁化，从而导致恢复到激发前的稳定状态，产生所谓的T1图像和T2磁共振图像。T1松弛与核子对其周围的能量过剩有关，而T2松弛指的是个别磁化矢量开始相互抵消的现象。上述现象是完全独立的。因此，不同的强度代表不同的组织，如下图所示。

4. 三维医学图像表示

  由于医学图像表示的是三维结构，所以可以通过使用三维体积的切片来处理它们，并进行常规的二维滑动卷积，如下图所示。我们假设红色矩形是一个5x5的图像片，可以用一个包含强度值的矩阵来表示。体素强度和内核用一个3x3的卷积内核进行卷积，如下图所示。以同样的模式，核在整个二维网格（医学图像切片）上滑动，每次我们都进行交叉相关。一个5x5的卷积补丁的结果被存储在一个3x3的矩阵中（为说明起见没有填充），并在网络的下一层传播。  或者，你可以把它们表示成类似于中间层的输出。在深度架构中，我们通常有多个特征图，这实际上是一个三维张量。如果有理由相信在额外的维度中存在模式，那么执行三维滑动卷积是最佳选择。在医学图像中就是这样的情况。与二维卷积类似，二维卷积是对二维领域中物体的空间关系进行编码，三维卷积可以描述三维空间中物体的空间关系。由于二维表示法对医学图像来说是次优的，我们将在这篇文章中选择使用三维卷积网络。

  医学图像切片可以被看作是一个中间层的多个特征图，不同的是，它们有很强的空间关系。

5. 3D-Unet模型

  在我们的例子中，我们将使用已被接受的3D U形网络。后者（代码）扩展了对称的U形二维Unet网络的连续理念，在RGB相关的任务中产生了令人印象深刻的结果，如语义分割。该模型有一个编码器（收缩路径）和一个解码器（合成路径）路径，每个路径有四个解析步骤。在编码器路径中，每层包含两个3×3×3的卷积，每个卷积后面都有一个整流线性单元（ReLu），然后是一个2×2×2的最大集合，每个维度的步长为2。在解码路径中，每层由2×2×2的转置卷积组成，每个维度的步长为2，然后是两个3×3×3的卷积，每个卷积后面是一个ReLu。分析路径中同等分辨率层的快捷跳转连接为合成路径提供了基本的高分辨率特征。在最后一层，一个1×1×1的卷积将输出通道的数量减少到标签的数量。通过在最大池化之前将通道数量增加一倍，避免了瓶颈问题。在每个ReLU之前，都会引入三维批次归一化。在训练过程中，每个批次的平均值和标准差都被归一化，并且使用这些值更新全局统计数据。接下来是一个明确学习比例和偏置的层。下图说明了网络结构。

5.1损失函数：Dice Loss

  由于固有的任务不平衡性，交叉熵不能总是为这项任务提供良好的解决方案。具体来说，交叉熵损失单独检查每个像素，将类预测（深度方向的像素向量）与我们的一热编码目标向量进行比较。因为交叉熵损失单独评估了每个像素向量的类别预测，然后对所有像素进行平均，所以我们基本上是在断言对图像中的每个像素进行平等的学习。如果你的各种类别在图像中的代表性不平衡，这可能是一个问题，因为最普遍的类别可能主导训练。

  我们将试图在大脑核磁共振成像中区分的4个类别在图像中具有不同的频率（即空气比其他组织的实例多得多）。这就是为什么采用Dice Loss指标的原因。它以骰子系数为基础，本质上是对两个样本之间重叠的测量。这个度量的范围从0到1，其中Dice系数为1表示完美和完全重叠。Dice损失最初是为二元分类而开发的，但它可以被推广到多类工作中。请自由使用我们的Dice loss的多类实现。

5.2医学成像数据

  深度架构需要大量的训练样本，然后才能产生任何有用的泛化表征，而标记的训练数据通常既昂贵又难以产生。这就是为什么我们每天都能看到使用生成学习的新技术来产生越来越多的医学成像数据。此外，训练数据必须能代表网络在未来将遇到的数据。如果训练样本来自与现实世界中会遇到的数据分布不同，那么网络的泛化性能将低于预期。

  由于我们关注的是大脑核磁共振自动分割，所以有必要简单介绍一下DNN试图区分的大脑基本结构：a）白质（WM），b）灰质（GM），c）脑脊液（CSF）。下图说明了大脑MRI切片中的分割组织。

5.2.12017年I-Seg医学图像数据挑战赛

  在这个关键时期，将婴儿大脑MRI图像准确分割成白质（WM）、灰质（GM）和脑脊液（CSF），对研究正常和异常的早期大脑发育具有根本性的意义。生命的第一年是产后人脑发育最活跃的阶段，伴随着快速的组织生长和各种认知和运动功能的发展。这个早期阶段在许多神经发育和神经精神疾病中至关重要，如精神分裂症和自闭症。人们对这一关键时期给予了越来越多的关注。

  这个数据集的目的是促进6个月婴儿大脑核磁共振成像的自动分割算法。这项挑战与MICCAI 2017同时进行，共有21个国际团队参加。该数据集包含10张来自专家的密集注释的图像和13张用于测试的成像。不提供测试标签，只有在官方网站上上传结果后才能看到自己的分数。每个受试者都有一张T1加权和T2加权的图像。

  第一个受试者将用于测试。原始MR卷的大小为256x192x144。在3D-Unet中，所使用的采样子卷的大小为128x128x64。产生的训练数据集包括500个子卷。对于验证集，使用了来自一个主体的10个随机样本。

  数据集下载

6. Medical Zoo

  我们的目标是在PyTorch中实现一个由最先进的三维深度神经网络组成的开源医学图像分割库，以及最常见的医学数据集的数据加载器。我们的资源库的第一个稳定版本预计将很快发布。

  我们坚信开放和可重复的深度学习研究。为了重现我们的成果，这项工作的代码和材料都可以在这个资源库中找到。这个项目是作为硕士论文开始的，目前正在进一步开发中。

6.1实施细节

  我们使用PyTorch框架，这被认为是最广泛接受的深度学习研究工具。所有的实验都采用了随机梯度下降法，单批大小，学习率为1e-3，权重衰减为1e-8。我们在资源库中提供了测试，你可以很容易地重现我们的结果，这样你就可以使用代码、模型和数据加载器。

  最近，我们用Pytorch增加了Tensorboard的可视化功能。这个惊人的功能让你的理智保持在原地，让你跟踪模型的训练过程。下面你可以看到一个保持训练统计的例子，dice coeff.和损失，以及每类分数，以了解模型的行为。

6.2代码

  让我们把所有描述的模块放在一起，用MedicalZoo的一个简短的脚本（为了说明问题）建立一个实验。

# Python libraries
import argparse
import os

# Lib files
import lib.medloaders as medical_loaders
import lib.medzoo as medzoo
import lib.train as train
import lib.utils as utils
from lib.losses3D import DiceLoss

def main():
    args = get_arguments()
    utils.make_dirs(args.save)

    training_generator, val_generator, full_volume, affine = medical_loaders.generate_datasets(args,
                                                                                               path='.././datasets')
    model, optimizer = medzoo.create_model(args)
    criterion = DiceLoss(classes=args.classes)

    if args.cuda:
        model = model.cuda()
        print("Model transferred in GPU.....")

    trainer = train.Trainer(args, model, criterion, optimizer, train_data_loader=training_generator,
                            valid_data_loader=val_generator, lr_scheduler=None)
    print("START TRAINING...")
    trainer.training()


def get_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--batchSz', type=int, default=4)
    parser.add_argument('--dataset_name', type=str, default="iseg2017")
    parser.add_argument('--dim', nargs="+", type=int, default=(64, 64, 64))
    parser.add_argument('--nEpochs', type=int, default=200)
    parser.add_argument('--classes', type=int, default=4)
    parser.add_argument('--samples_train', type=int, default=1024)
    parser.add_argument('--samples_val', type=int, default=128)
    parser.add_argument('--inChannels', type=int, default=2)
    parser.add_argument('--inModalities', type=int, default=2)
    parser.add_argument('--threshold', default=0.1, type=float)
    parser.add_argument('--terminal_show_freq', default=50)
    parser.add_argument('--augmentation', action='store_true', default=False)
    parser.add_argument('--normalization', default='full_volume_mean', type=str,
                        help='Tensor normalization: options ,max_min,',
                        choices=('max_min', 'full_volume_mean', 'brats', 'max', 'mean'))
    parser.add_argument('--split', default=0.8, type=float, help='Select percentage of training data(default: 0.8)')
    parser.add_argument('--lr', default=1e-2, type=float,
                        help='learning rate (default: 1e-3)')
    parser.add_argument('--cuda', action='store_true', default=True)
    parser.add_argument('--loadData', default=True)
    parser.add_argument('--resume', default='', type=str, metavar='PATH',
                        help='path to latest checkpoint (default: none)')
    parser.add_argument('--model', type=str, default='VNET',
                        choices=('VNET', 'VNET2', 'UNET3D', 'DENSENET1', 'DENSENET2', 'DENSENET3', 'HYPERDENSENET'))
    parser.add_argument('--opt', type=str, default='sgd',
                        choices=('sgd', 'adam', 'rmsprop'))
    parser.add_argument('--log_dir', type=str,
                        default='../runs/')

    args = parser.parse_args()

    args.save = '../saved_models/' + args.model + '_checkpoints/' + args.model + '_{}_{}_'.format(
        utils.datestr(), args.dataset_name)
    return args


if __name__ == '__main__':
    main()

6.3实验结果

  下面你可以看到模型的训练和验证Dice LOSS曲线。监测你的模型性能并调整参数以获得如此平滑的训练曲线是很重要的。这很容易理解这个模型的效率。


  令人惊讶的是，该模型在子卷的验证集中达到了大约93%的骰子系数得分。最后但同样重要的是，让我们看看3D-Unet在验证集中的一些可视化预测。我们在这里只展示了一个有代表性的切片，尽管预测的是一个三维卷。通过提取核磁共振的多个子容积，我们可以将它们结合起来形成一个完整的三维核磁共振分割。请注意，我们使用子卷取样的事实是作为数据的增强。

来自训练过的3D-Unet的未规范化的最后一层预激活。该网络学习高度语义的任务相关内容，对应于与输入相似的大脑结构。

  我们的预测VS地面真相。你认为哪一个预测是地面真相？在你决定之前，请仔细观察 值得注意的是，我们在这里只展示了中轴切片，但预测是一个三维体积。人们可以观察到，该网络完美地预测了空气体素，而它在区分组织边界方面有困难。但是，让我们再检查一下，找出真正的!

  现在，我相信你能分辨出地面的真相。如果你不确定，请查看文章的结尾

  最近我们还增加了Pytorch的Tensorboard可视化功能。这个惊人的功能让你的理智保持在原地，让你跟踪模型的训练过程。下面你可以看到一个保持训练统计的例子，骰子系数和损失，以及每类分数，以了解模型的行为。

  很明显，不同的组织有不同的准确度，甚至从训练的开始。例如，看一下验证集中的空气体素，它们从高值开始，因为它是不平衡数据集中最主要的一类。另一方面，灰质从最低值开始，因为它是最难区分的，而且训练实例较少。

7. 总结

  这篇文章部分地说明了MedicalZoo Pytorch库的一些特点。深度学习模型将为社会提供沉浸式医疗图像解决方案。

  在这篇文章中，回顾了医学成像和MRI的基本概念，以及它们如何在深度学习架构中被表示和使用。然后，描述了一个高效的被广泛接受的三维架构（Unet）和处理类不平衡的骰子损失函数。最后，结合了所有上述描述的特征，并使用库脚本提供了我们在大脑MRI中的实验分析的初步结果。这些结果证明了3D架构的效率和深度学习在医学图像分析中的潜力。