Swin-PSAxialNet：一种高效的多器官分割技术

摘要

本方案描述了一种高效的多器官分割方法，称为Swin-PSAxialNet，该方法相较于以往的分割方法具有优异的准确性。该流程的关键步骤包括数据集收集、环境配置、数据预处理、模型训练与比较，以及消融实验。

摘要

Abdominal multi-organ segmentation is one of the most important topics in the field of medical image analysis, and it plays an important role in supporting clinical workflows such as disease diagnosis and treatment planning. In this study, an efficient multi-organ segmentation method called Swin-PSAxialNet based on the nnU-Net architecture is proposed. It was designed specifically for the precise segmentation of 11 abdominal organs in CT images. The proposed network has made the following improvements compared to nnU-Net. Firstly, Space-to-depth (SPD) modules and parameter-shared axial attention (PSAA) feature extraction blocks were introduced, enhancing the capability of 3D image feature extraction. Secondly, a multi-scale image fusion approach was employed to capture detailed information and spatial features, improving the capability of extracting subtle features and edge features. Lastly, a parameter-sharing method was introduced to reduce the model's computational cost and training speed. The proposed network achieves an average Dice coefficient of 0.93342 for the segmentation task involving 11 organs. Experimental results indicate the notable superiority of Swin-PSAxialNet over previous mainstream segmentation methods. The method shows excellent accuracy and low computational costs in segmenting major abdominal organs.

引言

当前的临床干预，包括疾病诊断、治疗方案制定以及治疗效果追踪，均依赖于对医学图像的精确分割¹。然而，腹部器官之间复杂的结构关系²使得实现多个腹部器官的精确分割成为一项具有挑战性的任务³。在过去的几十年中，医学影像与计算机视觉技术的蓬勃发展，为腹部多器官分割领域带来了新的机遇与挑战。先进的磁共振成像（Magnetic Resonance Imaging, MRI）⁴和计算机断层扫描（Computed Tomography, CT）技术⁵使我们能够获取高分辨率的腹部图像。从CT图像中对多个器官进行精确分割，对于肝脏、肾脏、脾脏、胰腺等重要器官的评估与治疗具有重要的临床价值^6,7,8,9,10。然而，对这些解剖结构的手动标注，尤其是需要放射科医生或放射肿瘤科医生参与的操作，既耗时又容易受到主观因素的影响¹¹。因此，迫切需要开发自动化且精确的腹部多器官分割方法。

以往的图像分割研究主要依赖卷积神经网络（CNN），通过堆叠网络层并引入 ResNet¹² 来提高分割效率。2020 年，谷歌研究团队提出了视觉 Transformer（Vision Transformer, VIT）模型¹³，首次成功将 Transformer 架构引入传统视觉领域，用于多种视觉任务¹⁴。卷积操作只能关注局部特征信息，而 Transformer 中的注意力机制则能够全面考虑全局特征信息。

鉴于基于Transformer的架构在性能上优于传统的卷积网络¹⁵，许多研究团队已广泛开展研究，致力于优化Transformer与卷积网络各自优势之间的协同作用^16,17,18,19。Chen等人提出了用于医学图像分割任务的TransUNet模型¹⁶，该模型利用Transformer从图像中提取全局特征。然而，由于网络训练成本较高，且未能有效利用特征提取的层次化结构概念，Transformer的优势尚未得到充分发挥。

为了解决这些问题，许多研究人员开始尝试将 Transformer 作为训练分割网络的主干架构。Liu 等人¹⁷提出了 Swin Transformer，该模型采用分层构建方法进行多层级特征提取。该方法引入了窗口多头自注意力机制（Windows Multi-Head Self-Attention, W-MSA），显著降低了计算成本，尤其是在存在较大浅层特征图的情况下。尽管该方法减少了计算需求，但也导致了不同窗口之间的信息传递被隔离。为解决这一问题，作者进一步提出了移位窗口多头自注意力机制（Shifted Windows Multi-Head Self-Attention, SW-MSA），实现了相邻窗口之间的信息传播。在此基础上，Cao 等人提出了 Swin-UNet¹⁸，将 U-Net 中的二维卷积替换为 Swin 模块，并在编码和解码过程中引入 W-MSA 和 SW-MSA，取得了令人满意的分割效果。

相反，Zhou 等人强调，在处理高分辨率图像时，卷积操作的优势不容忽视¹⁹。他们提出的 nnFormer 采用基于局部三维图像块的自注意力计算方法，构成了一种具有十字形结构的 Transformer 模型。基于局部三维块的注意力机制显著降低了网络的训练负担。

鉴于上述研究存在的问题，本文提出了一种用于三维医学图像分割的高效混合分层结构，称为 Swin-PSAxialNet。该方法引入了一种下采样模块——空间到深度（Space-to-depth, SPD）²⁰ 模块，能够提取全局信息²¹。此外，还加入了一种参数共享的轴向注意力（parameter shared axial attention, PSAA）模块，可将学习参数数量从二次降低至线性级别，从而有效提升网络训练的准确性，并降低训练模型的复杂度²²。

Swin-PSAxialNet 网络
该网络的整体架构采用 nnU-Net²³ 的 U 形结构，包含编码器和解码器结构。这些结构负责局部特征提取，并对来自大尺度和小尺度图像的特征进行拼接，如图 1所示。

图 1：Swin-PSAxialNet 网络架构示意图。 请点击此处查看此图的放大版本。

在编码器结构中，传统的卷积模块（Conv block）与SPD模块²⁰相结合，构成一个下采样单元。编码器的第一层引入了“Patch Embedding”模块，该模块将三维数据划分为三维图像块 ，其中 (P₁, P₂, P₃) 表示此处的非重叠图像块， 表示三维图像块的序列长度。在嵌入层之后，下一步是一个非重叠的卷积下采样单元，该单元由一个卷积模块和一个SPD模块共同组成。在此结构中，卷积模块的步幅设为1，而SPD模块用于图像缩放，从而实现分辨率降低为原来的四分之一，同时通道数增加为原来的两倍。

在解码器结构中，瓶颈特征层之后的每个上采样模块均由一个上采样模块和一个PSAA模块组合而成。在每两个解码阶段之间，特征图的分辨率提高两倍，通道数减半。为了恢复空间信息并增强特征表示，在上采样模块之间执行大尺度与小尺度图像之间的特征融合。最终，上采样结果被输入至头部层，以恢复原始图像尺寸，输出尺寸为（H × W × D × C，其中C = 3）。

SPD 模块架构
在传统方法中，下采样部分采用步长为 2 的单一卷积操作。这种方法涉及在图像的局部位置进行卷积池化，限制了感受野，使模型只能从小块图像区域中提取特征。本文方法采用 SPD 模块，将原始图像精细地划分为三个维度。原始的三维图像沿 x、y 和 z 轴被均匀分割，生成四个子体块。（图 2）随后，通过“cat”操作将这四个体块进行拼接，并对所得图像进行 1 × 1 × 1 卷积，从而获得下采样后的图像²⁰。

图 2：SPD 模块框图。 请点击此处查看此图的放大版本。

PSAA 模块架构
与传统的卷积神经网络（CNN）相比，所提出的 PSAA 模块在实现全局信息聚焦方面更为有效，在网络学习与训练中也更加高效，从而能够捕获更丰富的图像与空间特征。PSAA 模块包含基于参数共享的三维轴向注意力学习机制，这三个维度分别为：高度（Height）、宽度（Width）和深度（Depth）。相较于传统注意力机制对图像中每个像素进行注意力学习的方式，该方法在三个维度上分别独立进行注意力学习，将自注意力的复杂度从二次级降低至线性级。此外，该模块采用可学习的键-查询（keys-queries）参数共享机制，使网络能够在三个维度上并行执行注意力机制操作，从而实现更快、更优且更有效的特征表示。

方案

本方案已获得南通大学伦理委员会批准。该方案涉及利用人工智能技术对获得的无创或微创多模态数据进行智能评估与研究，包括人体医学影像、肢体运动和血管成像数据。图3展示了多器官分割的整体流程图。所有必要的网络链接均列在材料表中。

图3：多器官分割的整体流程图。 请点击此处查看此图的放大版本。

1. 数据集收集

1. 下载 AMOS2022 数据集²⁴（参见材料表）。
2. 选择 160 张图像作为训练数据集，40 张图像作为测试数据集。
3. 终端执行 mkdir ~/PaddleSeg/contrib/MedicalSeg/data/raw_data。
4. 打开 raw_data 路径。
5. 将数据集放置在 raw_data 目录中，并确保与“imagesTr”和“labelsTr”目录中的文件名正确对应。
6. 创建 dataset.json 文件。
7. 将数据集信息（如 numTraining、numTest 和 labels name）添加到“dataset.json”文件中。
   注意：本研究使用的数据集由从 AMOS2022 官方 500 例 CT 数据集中筛选出的 200 例高质量 CT 数据组成，用于训练和测试。该数据集包含 11 个器官，分别为脾脏、右肾、左肾、胆囊、食管、肝脏、胃、主动脉、下腔静脉、胰腺和膀胱。

2. 环境配置

1. 安装 Paddlepaddle 环境（参见材料表）及相应版本的 CUDA²⁵。
   注意：建议使用 Paddlepaddle 2.5.0 版本和 CUDA 11.7 版本。
2. 通过 Git 克隆 PaddleSeg 仓库²⁵。
3. 在终端中执行 cd ~/PaddleSeg/contrib/MedicalSeg/。
4. 执行 pip install -r requirements.txt。
5. 执行 pip install medpy。

3. 数据预处理

1. 修改“yml”文件中的参数，包括 data_root, batch_size, model, train_dataset, val_dataset, optimizer, lr_scheduler 和 loss 等。
2. 指定执行所需参数，包括 "-config, --log_iters, --precision, --nnunet, --save_dir, --save_interval" 等。
   注意：原始的“raw_data”被转换为三种类型的数据文件：“Cropped”、“Decathlon”和“Preprocessed”，以满足不同的训练与评估需求，从而提升模型性能和训练效率。

4. 模型训练与比较

注意：作为图像分割领域广泛使用的基线方法，nnU-Net²³ 在本研究中被用作基准模型。具体模型比较过程如下。

1. 训练 nnU-Net
   1. 配置批量大小为 2，迭代次数为 80,000，每 20 次迭代记录一次日志，并使用 fp16 精度。
   2. 将 save_dir 设置为工作目录下的“output”文件夹，保存间隔设为每 1,000 次迭代一次。
   3. 添加参数 --use_vdl 以实现 VisualDL 参数可视化。
   4. 运行 train.py。
   5. 创建四个 yml 文件，指定折数为 1-4，并使用五折交叉验证训练模型。
2. 训练 Swin-UNet
   1. 保持与先前设定的实验条件一致。
   2. 移除参数 --nnunet，改用 --swinunet。
   3. 运行 train.py 文件。
   4. 分别训练第 1 至第 4 折模型，然后进行集成处理。
3. 训练 TransUNet
   1. 遵循步骤 4.1 中所述流程。
   2. 打开 yml 文件。
   3. 将 type: nnunet 修改为 type: transunet。
   4. 将参数 --nnunet 修改为 --transunet。
   5. 运行 train.py 文件。
   6. 分别训练第 1 至第 4 折模型，并进行集成处理。
4. 训练 UNETR
   1. 保持实验变量一致。
   2. 重复之前 Swin-UNet 和 TransUNet（步骤 4.2 和步骤 4.3）所述流程。
   3. 运行 train.py。
   4. 分别训练第 1 至第 4 折模型，并进行集成处理。
5. 训练 nnFormer
   1. 参照 Swin-UNet 和 TransUNet 的方式修改参数。
   2. 运行 single fold train.py。
   3. 执行集成处理。

5. 消融实验

1. 整个网络的 PSAA 模块
   1. 保持变量与步骤 4.1 一致。
   2. 设置 volume_shape = (64,192,192) 且 Axial_attention = True。
   3. 运行 single fold train.py。
   4. 对五个折叠进行集成处理。
2. 整个网络的 SPD 模块
   1. 保持变量与步骤 4.1 一致。
   2. 设置 SPD_enable = True 且 Axial_attention = False。
   3. 运行 single fold train.py。
   4. 对五个折叠进行集成处理。
3. 上采样部分使用 PSAA 模块
   注意：下采样部分不要使用 SPD 模块。
   1. 上采样部分使用 PSAA 模块。
   2. 保持变量与步骤 4.1 一致。
   3. 在下采样部分设置 Axial_attention = False，并设置 SPD_enable = False。
   4. 运行 single fold train.py。
   5. 对五个折叠进行集成处理。
4. 上采样部分使用 PSAA 模块
   注意：下采样部分使用 SPD 模块。
   1. 下采样部分使用 SPD 模块，上采样部分使用 PSAA 模块。
   2. 保持变量与步骤 4.1 一致。
   3. 在下采样部分设置 Axial_attention = False，并设置 SPD_enable = True。
   4. 运行 single fold train.py。
   5. 对五个折叠进行集成处理。
   6. 由于结果更优，将迭代次数设置为 200,000。

结果

本方案采用两个指标来评估模型：Dice 相似性分数（DSC）和 95% Hausdorff 距离（HD95）。DSC 用于衡量体素分割预测结果与真实标签之间的重叠程度，而 95% HD 则用于评估体素分割预测边界与真实标签边界的重叠情况，并过滤掉 5% 的离群值。DSC²⁶ 的定义如下：

      (1)

其中，T 和 P 分别表示所有体素的真值和预测值。95% HD²⁶ 的定义如下：

      (2)

其中，d_T^,_P^, 表示真实体素与预测体素之间最大 HD₉₅ 的第95百分位距离，而 d_T^,_P^, 表示预测体素与真实体素之间最大 HD₉₅ 的第95百分位距离。

Swin-PSAxialNet 与多种主流多器官分割算法在表1中进行了比较。结果表明，所提出的算法在腹部器官分割方面表现出更优的性能。该分割方法的 Dice 准确率比其他方法的平均 Dice 值高出 2-3 个百分点，甚至比 nnFormer 高出 1 个百分点。Swin-PSAxialNet 在边界分割方面也表现出高精度，其平均 HD95 系数为 5.63，是所有对比算法中最低的。

方法		Lkn	Spl	Liv	Pan	Aro	Bla	Sto	Gbd	Pos	Eso	Rkn	平均值
													DSC	HD95
nnU-Net		96.3	96.8	97.21	85.2	95.22	86.57	91.47	84.7	91.79	82.5	95.44	91.2	7.13
Swin-Unet		96.03	96.68	97.22	85.74	95.3	87.82	91.77	84.64	91.67	83.26	96.14	91.48	9.68
TransUNet		96.42	96.74	97.05	85.09	94.86	85.2	91.81	83.33	91.52	82.22	96.12	90.94	13.77
UNETR		95.94	96.48	96.85	84.14	94.62	85.01	90.18	81.16	90.79	80.34	95.15	90.06	12.92
nnFormer		96.59	97.22	97.33	86.79	95.31	89.74	92.82	84.88	92.43	84.24	96.32	92.15	6.59
Swin-PSAxialNet		96.85	97.67	97.64	88.39	95.97	92.03	94.4	87.78	93.15	86.24	96.59	93.34	5.63

表1：主流分割算法的比较。 缩写代表以下器官：Lkn 表示左肾，Spl 表示脾脏，Liv 表示肝脏，Pan 表示胰腺，Aro 表示主动脉，Bla 表示膀胱，Sto 表示胃，Gbd 表示胆囊，Pos 表示下腔静脉，Eso 表示食管，Rkn 表示右肾，DSC 表示平均 Dice 系数，HD95 表示平均 95% Hausdorff 距离。

此外，参数共享机制的引入以及对 Transformer 结构的改进，在模型复杂度和推理速度方面均带来了显著优势（表2, 图 4，以及 图5).

模型	参数量 (M)	浮点运算量 (G)	推理时间 (s)
nn-Unet	17.96	398.54	9.07
UNETR	89.58	39.67	12.51
nnFormer	134.78	152.43	11.24
Swin-PSAxialNet	37.64	43.15	10.31

表2：网络模型指标比较表。

图 4：模型训练网络的 Dice 变化曲线。 （A）当前网络。（B）nnFormer。（C）nnU-Net。（D）UNETR。 请点击此处查看此图的放大版本。

图5：模型训练结果对比。 该图展示了当前网络的结果、nnFormer的结果、nnU-Net的结果以及UNETR的结果。请点击此处查看此图的放大版本。

本研究对算法网络进行了全面的消融实验，以PSAA模块和SPD模块作为变量，比较了单折训练与五折训练的效果。实验结果的Dice系数如表3所示。结果表明，与未修改的nnU-Net相比，使用SPD模块进行下采样、PSAA模块进行上采样能有效提升分割性能（图6）。

模型	整个网络的 PSAA 模块	整个网络的 SPD 模块	仅上采样部分使用 PSAA 模块	上采样部分使用 PSAA 模块，并采用 SPD
单次折叠训练	92.17	90.51	92.24	92.39
五次折叠训练	92.82	91.26	92.79	93.34

表3：Swin-PSAxialNet消融实验的结果。 展示了Swin-PSAxialNet的消融实验结果。在消融实验中，比较了PSAA模块和SPD模块的使用情况，表中列出了DSC的训练结果。

图6：分割结果。 请点击此处查看此图的放大版本。

讨论

腹部器官的分割是一项复杂的工作。与其他人体内部结构（如大脑或心脏）相比，由于CT图像中对比度较低且器官形态变化较大，腹部器官的分割显得更具挑战性^27,28。本文提出Swin-PSAxialNet以解决这一难题。

在数据收集步骤中，本研究从AMOS2022官方网站下载了200张图像²⁴。在环境配置步骤中，实验环境包括Paddlepaddle框架及其对应的CUDA版本。在数据预处理步骤中，采用了随机裁剪和随机翻转以增强数据的鲁棒性。在模型训练与比较步骤中，基于基线模型nnU-Net，实验在下采样过程中引入SPD Block，并在上采样过程中引入PSAA特征提取模块，取得了优异的分割效果。Swin-PSAxialNet与多种主流分割模型进行了比较，在公开数据集AMOS2022上表现出良好的性能，其分割精度高于本研究中提到的其他分割方法^16,17,18,19。在消融实验部分，通过调整PSAA模块和SPD模块的使用方式，以寻找最优的分割网络结构。消融实验结果表明，仅在上采样阶段使用PSAA模块，其性能比在整个网络中使用该模块高出一个百分点。

本研究的创新之处如下：首先，将SPD模块引入网络，与Conv模块共同构成下采样单元，从而增强了网络的特征提取能力。其次，引入PSAA特征提取模块，在降低网络复杂度的同时提高了训练精度。最后，采用参数共享机制，显著降低了训练模型的复杂度。

本研究在以下方面存在局限性：仅应用于AMOS2022数据集，尚未在其他数据集上验证结果。此外，尽管nnU-Net作为自动化分割框架经过一系列优化后在腹部器官分割任务中已取得较高性能，但部分参数仍具有一定程度的主观性，例如初始学习率、迭代次数、网络结构等。未来的研究将探索这些参数的自动计算方法²⁹。

Swin-PSAxialNet 在腹部器官分割任务中取得了优异的分割效果，表现出较强的泛化能力和稳定性。在未来的研究中，将引入更多数据集以进一步探究该算法的泛化能力，并对模型进行相应调整。此外，该算法的结构还可进一步优化，以提升网络的整体精度和鲁棒性。目标是将该网络与多模态能力相结合，将其发展为一种医学多模态模型。

材料

本文使用的材料清单

姓名	公司	目录编号	评论
AMOS2022 数据集	无	无	用于网络训练和测试的数据集。网络链接为：https://pan.baidu.com/s/1x2ZW5FiZtVap0er55Wk4VQ?pwd=xhpb
ASUS 主机	ASUS		https://www.asusparts.eu/en/asus-13020-01910200
CUDA 版本 11.7	NVIDIA		https://developer.nvidia.com/cuda-11-7-0-download-archive
NVIDIA GeForce RTX 3090	NVIDIA		https://www.nvidia.com/en-in/geforce/graphics-cards/30-series/rtx-3090-3090ti/
Paddlepaddle 环境	Baidu	无	网络训练的环境准备。网络链接为：https://www.paddlepaddle.org.cn/
PaddleSeg	Baidu	无	我们使用的基线模型：https://github.com/PaddlePaddle/PaddleSeg