医疗AI：轻量级图像分割新突破
在医学图像领域，精准的图像分割技术一直是提高诊断效率和准确性的关键，然而传统的图像分割方法常常受到计算资源和处理速度的限制，尤其在资源紧张的医疗环境中更是如此。随着人工智能技术的飞速发展，我们迎来了一个激动人心的新时代——轻量级医疗AI图像分割的突破正悄然改变着医疗行业的面貌。
目录
概述
核心逻辑
复现过程
写在最后
概述
本文将解读并且复现一篇来自2022年医学图像顶会MICCAI的文章：UNeXt: MLP-based Rapid Medical Image Segmentation Network 。这篇文章主要解决了Unet网络及其扩展网络近年来虽然取得了极大的进展，但是在网络参数、计算复杂度和速度上非常慢，不能够适应实时的医学图像分割的问题。文章主要从以下方面解决了上述问题：
1）提出UNeXt网络
UNeXt是一种基于卷积多层感知机（MLP）的图像分割网络
设计上结合了早期卷积阶段和潜在层中的MLP阶段，以实现高效分割
2）核心创新点
Tokenized MLP块：有效地对卷积特征进行标记（Tokenization）和投影，使用MLP建模，减少了参数量和计算复杂度
通道移位：在输入MLP时移动输入通道，以更专注于学习局部依赖，提升性能
3）网络结构特点
编码器与解码器之间的不同层级包含跳跃连接（Skip Connections），有助于信息流动和分割精度
尽管UNet及其多种扩展（如UNet++、UNet3+、3D UNet、V-Net、Y-Net和KiUNet）以及基于Transformer的网络（如TransUNet、MedT、TransBTS和UNETR）在医学图像分割方面取得了显著进展，但这些方法大多侧重于提升网络性能，而忽略了计算复杂度、推理时间和参数数量，这在许多实际应用中至关重要。
紧接着又说随着医学成像解决方案从实验室向床边（即“即时成像”）的转移，对高效、低参数、快速推理的分割网络的需求日益增加。即时成像有助于临床医生提供更快捷的服务和更好的患者护理，减少患者前往放射科中心的时间和程序。
然后为了解决这些问题提出了UNeXt网络，该网络结合了卷积网络和多层感知机（MLP），旨在减少计算开销、降低参数数量、提高推理速度，同时保持良好的分割性能。UNeXt网络采用UNet的编码器-解码器架构，但每个块的设计有所不同，包括一个卷积阶段和一个MLP阶段。在瓶颈部分，引入了新颖的Tokenized MLP（Tok-MLP）块，该块通过将卷积特征投影为抽象标记并使用MLP学习有意义的分割信息，有效降低了计算量。此外，还在MLP中引入了轴向移位操作，以提取不同轴向的局部信息。
最后实验结果表明，UNeXt在ISIC 皮肤病变数据集和乳腺超声图像（BUSI）数据集上取得了优于最新通用分割架构的性能。更重要的是，与TransUNet相比，UNeXt将参数数量减少了72倍，计算复杂度降低了68倍，推理速度提高了10倍，使其非常适合于即时医疗成像应用。
核心逻辑
UNext网络结构同样是encoder-decoder 结构，有两个阶段：卷积阶段；tokenized MLP阶段
UNeXt 是一种编码器-解码器架构，分为两个阶段：卷积阶段和tokenized MLP阶段。编码器包括前3个卷积块和后2个 MLP 块，逐步减小特征分辨率；解码器反向操作，包括2个 MLP 块和3个卷积块，逐步增大分辨率，并通过跳跃连接与编码器相连。通道数（C1 到 C5）设置为32、64、128、160和256，比 UNet 及其变体的滤波器数量更少，从而减少了参数和计算量。整体架构如下图所示：
卷积阶段 ：卷积阶段中的每个卷积块包含一个卷积层、批量归一化层和 **ReLU **激活。编码器块使用 2 × 2 最大池化层，解码器块使用双线性插值进行上采样，以减少可学习参数数量，这里使用双线性插值不需要学习任何参数，它通过加权邻近像素值来计算新的像素值，从而在不引入额外的可学习参数的情况下实现图像的上采样。
tokenized MLP阶段 ：这一部分内容是本文的核心创新点，首先先介绍Shifted MLP在 Shifted MLP 中，先对卷积特征进行移动，以便 MLP 更好地专注于特定区域。可以把它想象成在分析图片时，先把图片分成小块，每次只看其中的一部分。通过移动这些小块，MLP 能够更好地理解局部细节，而不是只看整体，这种方法类似于Swin Transformer的做法，在原本需要全局处理的模型中，通过分区来更好地关注局部信息。具体来说，Tokenized MLP 有两个部分，一个部分负责沿宽度方向移动特征，另一个部分则沿高度方向移动特征。这样一来，模型可以更有针对性地在不同方向上关注细节，从而提高效果。
举个简单的例子：假设你有一张4×4的小图片，每个像素代表一个特征。原始图片可能是这样的：
1 2 3 4
5 6 7 8
9 10 11 12
13 14 15 16如果你向右平移一格，图片可能会变成这样：
4 1 2 3
8 5 6 7
12 9 10 11
16 13 14 15这种平移操作让每个像素的位置发生了变化，但图片的内容还是相同的。通过这种方式，MLP 在处理这些移动后的特征时，会看到特征图中不同的部分，从而增加模型对局部信息的敏感度 。从而不需要看到全局的信息，减少了参数量。如图所示：
灰色是特征块的位置，白色是移动之后的padding，然后介绍Tokenized MLP 在 Tokenized MLP 块中，首先对特征进行移动并将其投射为标记。然后，将这些标记输入一个沿宽度方向平移的 MLP，接着通过深度卷积层（DWConv）来编码位置信息并提高效率。使用 GELU 激活函数代替 ReLU，因为 GELU 提供了更平滑的激活效果。随后，特征通过另一个沿高度方向平移的 MLP，将维度从 H 转换为 O，并使用残差连接和层归一化（LN），以增强性能和稳定性。 如图所示：
为什么使用DWConv？
1）编码位置信息：DWConv 能够有效地编码位置信息。研究表明，卷积层在 MLP 块中能够有效捕捉位置信息，比传统的位置编码技术（如 ViT 中使用的技术）更优。传统的位置编码技术在测试和训练阶段分辨率不一致时需要插值处理，这可能导致性能下降，而卷积层则可以避免这种问题。
2）参数效率：DWConv 通过对每个通道单独进行卷积操作，减少了需要学习的参数数量。这提高了计算效率，并使模型在保持性能的同时更加高效。与标准卷积层相比，DWConv 的参数更少，因此可以降低计算开销和内存占用。
复现过程
实验数据集和代码已经下载好在附件和readme中。本复现过程只使用了部分的数据，且数据存放格式保持一致：
安装依赖包两种方式：基于conda和基于pip
# 基于conda
conda env create -f environment.yml
conda activate unext
# 基于pip
timm==0.3.2
mmcv-full==1.2.7
torch==1.7.1
torchvision==0.8.2
opencv-python==4.5.1.48
请确保你的数据集文件夹存放的结构如下：
inputs
└── <dataset name>
├── images
| ├── 001.png
│ ├── 002.png
│ ├── 003.png
│ ├── ...
||
└── masks
├── 0
| ├── 001.png
| ├── 002.png
| ├── 003.png
| ├── ...
||
└── 1
├── 001.png
├── 002.png
├── 003.png
├── ...
模型训练和验证，进入到相应的文件夹后分别进行以下操作:
训练模型 ：其中–dataset
是数据集路径；–name 可以为模型取个名字；–img_ext .png --mask_ext .png要根据具体的图片格式改成png或者jpg
模型代码及训练过程截图如下：
python train.py --dataset <dataset name> --arch UNext --name <exp name> --img_ext .png --mask_ext .png --lr 0.0001 --epochs 500 --input_w 512 --input_h 512 --b 8
模型评估 ：–name
这个就是上面自己取的那个名字
模型代码及验证过程截图如下：
对于这样一个能够兼顾到训练速度和精度的模型，可以在此基础上进行改进后为自己所用，是一个非常不错的选择。改进建议如下：
多尺度特征提取 ：将原UNext网络中用于提取图像特征的单一卷积模块替换为多尺度特征提取模块
注意力模块 ：在模型关键的位置分别增加自注意力机制和残差注意力机制，使得模型能够在关键位置建立权重连接，从而更好的捕获和理解图像中上下文的信息以及图像不同区域之间的关系，以及捕获图像中的目标细节和边界
写在最后
在医疗领域，轻量级图像分割技术的突破不仅仅是一个技术进步，它预示着未来医疗服务的全新模式。通过减少计算资源的需求并提高处理速度，这些创新将极大地推动医疗图像分析的普及，让更广泛的医疗机构能够享受到先进的AI技术带来的好处。随着这些技术不断成熟和应用，未来的医疗工作者将能够以更高的精确度和效率应对各种挑战，从而为患者提供更优质的医疗服务。无论是在偏远地区的诊断，还是在复杂的手术规划中，轻量级图像分割的进步都将成为改变医疗格局的关键力量。我们正站在医疗AI新纪元的门槛上，未来的每一步都值得我们共同期待和探索。
详细复现过程的项目源码、数据和预训练好的模型可从地址获取。