UNeXt

用于 UNeXt：基于 MLP 的快速医学图像分割网络 的官方 PyTorch 代码库，MICCAI 2022

论文 | 项目主页

简介

近年来，UNet 及其最新扩展如 TransUNet 一直是医学图像分割领域的主流方法。然而，这些网络由于参数量大、计算复杂且推理速度较慢，难以在即时医疗应用中实现快速图像分割。为此，我们提出了 UNeXt，这是一种基于卷积多层感知机（MLP）的图像分割网络。我们设计了一种高效的架构，包含早期的卷积阶段和潜在空间中的 MLP 阶段。我们提出了一种标记化 MLP 块，能够高效地对卷积特征进行标记化和投影，并利用 MLP 对特征表示进行建模。为进一步提升性能，我们在将输入送入 MLP 时对其通道进行移位，以更好地捕捉局部依赖关系。在潜在空间中使用标记化的 MLP 不仅减少了参数量和计算复杂度，还能生成更优的特征表示，从而提高分割效果。该网络还包含编码器和解码器各层级之间的跳跃连接。我们在多个医学图像分割数据集上测试了 UNeXt，结果表明，与当前最先进的医学图像分割架构相比，UNeXt 的参数量减少了 72 倍，计算复杂度降低了 68 倍，推理速度提升了 10 倍，同时取得了更好的分割性能。

使用代码：

代码在 Python 3.6.13 和 CUDA >=10.1 环境下运行稳定。

• 克隆本仓库：

git clone https://github.com/jeya-maria-jose/UNeXt-pytorch
cd UNeXt-pytorch

使用 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

数据集

1. ISIC 2018 - 链接
2. BUSI - 链接

数据格式

请确保按照以下结构组织文件（例如，类别数为 2）：

inputs
└── <数据集名称>
    ├── images
    |   ├── 001.png
    │   ├── 002.png
    │   ├── 003.png
    │   ├── ...
    |
    └── masks
        ├── 0
        |   ├── 001.png
        |   ├── 002.png
        |   ├── 003.png
        |   ├── ...
        |
        └── 1
            ├── 001.png
            ├── 002.png
            ├── 003.png
            ├── ...

对于二分类问题，只需使用文件夹 0 即可。

训练与验证

1. 训练模型。

python train.py --dataset <数据集名称> --arch UNext --name <实验名称> --img_ext .png --mask_ext .png --lr 0.0001 --epochs 500 --input_w 512 --input_h 512 --b 8

2. 评估模型。

python val.py --name <实验名称>

致谢：

本代码库使用了来自 UNet++、Segformer 和 AS-MLP 的部分代码块及辅助函数。特别感谢 Poojan 提供的命名贡献。

引用：

@article{valanarasu2022unext,
  title={UNeXt: MLP-based Rapid Medical Image Segmentation Network},
  author={Valanarasu, Jeya Maria Jose and Patel, Vishal M},
  journal={arXiv preprint arXiv:2203.04967},
  year={2022}
}

UNeXt-pytorch 快速上手指南

UNeXt 是一种基于卷积多层感知机（MLP）的快速医学图像分割网络。相比传统的 UNet 和 TransUNet，它在大幅减少参数量（72 倍）和计算复杂度（68 倍）的同时，将推理速度提升了 10 倍，并保持了更优的分割性能。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• Python: 3.6.13 (代码在此版本下最稳定)
• CUDA: >= 10.1
• GPU: 支持 CUDA 的 NVIDIA 显卡

安装步骤

1. 克隆项目

首先从 GitHub 克隆代码库并进入目录：

git clone https://github.com/jeya-maria-jose/UNeXt-pytorch
cd UNeXt-pytorch

2. 安装依赖

方式一：使用 Conda（推荐） 项目提供了完整的环境配置文件，可一键安装所有依赖：

conda env create -f environment.yml
conda activate unext

方式二：使用 Pip 如果您偏好使用 pip，请手动安装以下指定版本的依赖包：

pip install timm==0.3.2 mmcv-full==1.2.7 torch==1.7.1 torchvision==0.8.2 opencv-python==4.5.1.48

提示：国内用户安装 mmcv-full 或 torch 较慢时，可添加清华源或阿里源加速，例如： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple ...

基本使用

1. 数据准备

请按照以下目录结构组织数据集（以二分类为例，类别数为 2）：

inputs
└── <dataset name>
    ├── images
    |   ├── 001.png
    │   ├── 002.png
    │   └── ...
    |
    └── masks
        ├── 0
        |   ├── 001.png
        |   ├── 002.png
        |   └── ...
        |
        └── 1
            ├── 001.png
            ├── 002.png
            └── ...

注：对于二值分割任务，只需使用文件夹 0 存放掩码即可。

支持的数据集示例：

• ISIC 2018
• BUSI (Breast Ultrasound Images)

2. 训练模型

使用 train.py 启动训练。以下是一个典型的训练命令示例：

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

参数说明：

• --dataset: 数据集名称（需与 inputs 下的文件夹名一致）
• --name: 实验名称（用于保存模型和日志）
• --input_w / --input_h: 输入图像的宽和高
• --b: Batch Size

3. 验证与评估

训练完成后，使用 val.py 对模型进行评估：

python val.py --name <exp name>

将 <exp name> 替换为您训练时指定的实验名称即可。