医疗Transformer

用于论文《Medical Transformer: Gated Axial-Attention for Medical Image Segmentation》（MICCAI 2021）的PyTorch代码
论文 | 海报

新闻：

:rocket: : 欢迎查看我们的最新工作 UNeXt，这是一种更快、更高效的分割架构，同时易于训练和实现！代码可在这里获取。

关于此仓库：

本仓库包含以下网络的代码实现：

1. 门控轴向注意力U-Net
2. MedT

简介

大多数为视觉任务设计的基于Transformer的网络架构都需要大规模数据集才能进行有效训练。然而，与视觉领域的数据集相比，医学影像的数据样本数量相对较少，这使得在医学领域高效训练Transformer模型变得困难。为此，我们提出了一种门控轴向注意力模型，通过在自注意力模块中引入额外的控制机制来扩展现有架构。此外，为了在医学图像上有效地训练该模型，我们还提出了一种局部-全局训练策略（LoGo），以进一步提升性能。具体而言，我们分别对整张图像和图像块进行操作，以学习全局和局部特征。所提出的Medical Transformer（MedT）在门控轴向注意力U-Net的基础上使用了LoGo训练策略。

使用代码：

• 克隆此仓库：

git clone https://github.com/jeya-maria-jose/Medical-Transformer
cd Medical-Transformer

代码在Python 3.6.10和PyTorch 1.4.0环境下运行稳定。

使用conda安装所有依赖项：

conda env create -f environment.yml
conda activate medt

使用pip安装所有依赖项：

pip install -r requirements.txt

公开数据集下载链接：

1. MoNuSeG数据集 - 链接（原始）
2. GLAS数据集 - 链接（原始）
3. 论文中使用的Brain Anatomy US数据集将很快公开！

在您的数据集上使用代码

数据集准备

请按照以下格式准备数据集，以便于代码的使用。训练和测试文件夹应各包含两个子文件夹：img和label。确保图像与其对应的分割掩码位于这些文件夹中，并且文件名相同，以便于对应。如果您不希望按此格式准备数据集，可以相应地修改数据加载器。

训练文件夹-----
      img----
          0001.png
          0002.png
          .......
      label----
          0001.png
          0002.png
          .......
验证文件夹-----
      img----
          0001.png
          0002.png
          .......
      label----
          0001.png
          0002.png
          .......
测试文件夹-----
      img----
          0001.png
          0002.png
          .......
      label----
          0001.png
          0002.png
          .......

• 真值图像中的像素应与标签一一对应。例如，在二值分割中，真值图像的像素应为0或255。

训练命令：

python train.py --train_dataset "输入训练目录" --val_dataset "输入验证目录" --direc '保存结果的路径' --batch_size 4 --epoch 400 --save_freq 10 --modelname "gatedaxialunet" --learning_rate 0.001 --imgsize 128 --gray "no"

若要训练MedT或LoGo，请将modelname改为MedT或logo。

测试命令：

python test.py --loaddirec "./saved_model_path/model_name.pth" --val_dataset "测试数据集目录" --direc '保存结果的路径' --batch_size 1 --modelname "gatedaxialunet" --imgsize 128 --gray "no"

测试结果（包括预测的分割图）将与模型权重一同保存在results文件夹中。请使用MATLAB运行性能指标代码以计算F1分数和mIoU。

注意事项：

1. 请注意，这些实验是在配备48 GB显存的Nvidia Quadro 8000上进行的。
2. Google Colab代码是用于快速训练/测试的非官方实现。请遵循原始代码以确保正确训练。

致谢：

数据加载器代码灵感来源于 pytorch-UNet 。轴向注意力代码则基于axial-deeplab开发。

引用：

@InProceedings{jose2021medical,
author="Valanarasu, Jeya Maria Jose
and Oza, Poojan
and Hacihaliloglu, Ilker
and Patel, Vishal M.",
title="Medical Transformer: Gated Axial-Attention for Medical Image Segmentation",
booktitle="Medical Image Computing and Computer Assisted Intervention -- MICCAI 2021",
year="2021",
publisher="Springer International Publishing",
address="Cham",
pages="36--46",
isbn="978-3-030-87193-2"
}

如有任何疑问或建议，请随时提交Issue或直接联系我。

Medical-Transformer 快速上手指南

Medical-Transformer (MedT) 是 MICCAI 2021 发表的一种用于医学图像分割的 Transformer 架构，引入了门控轴向注意力机制（Gated Axial-Attention）和局部-全局训练策略（LoGo），旨在解决医疗数据量少导致 Transformer 难以训练的问题。

环境准备

• 系统要求：Linux/Windows/macOS
• 推荐配置：Nvidia GPU（原文实验基于 Nvidia Quadro 8000 48GB，显存较小需调整 batch_size）
• 核心依赖版本：
  • Python 3.6.10
  • PyTorch 1.4.0

安装步骤

1. 克隆代码库

git clone https://github.com/jeya-maria-jose/Medical-Transformer
cd Medical-Transformer

2. 安装依赖

你可以选择使用 Conda 或 Pip 进行安装。

方式一：使用 Conda（推荐）

conda env create -f environment.yml
conda activate medt

方式二：使用 Pip

pip install -r requirements.txt

提示：若在国内网络环境下安装缓慢，建议配置 pip 国内镜像源（如清华源）： pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

1. 数据集准备

请将数据集整理为以下目录结构，确保图像和标签文件名一致，且标签像素值对应类别（例如二分类为 0 和 255）。

Train Folder-----
      img----
          0001.png
          0002.png
          .......
      labelcol---
          0001.png
          0002.png
          .......
Validation Folder-----
      img----
          0001.png
          0002.png
          .......
      labelcol---
          0001.png
          0002.png
          .......
Test Folder-----
      img----
          0001.png
          0002.png
          .......
      labelcol---
          0001.png
          0002.png
          .......

支持的数据集示例：MoNuSeG, GLAS 等。

2. 模型训练

使用 train.py 进行训练。可以通过 --modelname 参数选择模型：gatedaxialunet、MedT 或 logo。

python train.py --train_dataset "enter train directory" --val_dataset "enter validation directory" --direc 'path for results to be saved' --batch_size 4 --epoch 400 --save_freq 10 --modelname "gatedaxialunet" --learning_rate 0.001 --imgsize 128 --gray "no"

关键参数说明：

• --train_dataset: 训练集路径
• --val_dataset: 验证集路径
• --direc: 结果保存路径
• --modelname: 模型名称 (gatedaxialunet, MedT, logo)
• --imgsize: 输入图像尺寸

3. 模型测试

使用 test.py 进行推理测试。

python test.py --loaddirec "./saved_model_path/model_name.pth" --val_dataset "test dataset directory" --direc 'path for results to be saved' --batch_size 1 --modelname "gatedaxialunet" --imgsize 128 --gray "no"

关键参数说明：

• --loaddirec: 预训练模型权重路径 (.pth)
• --val_dataset: 测试集路径
• --direc: 预测结果保存路径

测试完成后，预测的分割图将保存在结果文件夹中。如需计算 F1 Score 和 mIoU，请使用提供的 MATLAB 代码进行性能评估。