U-KAN：为医学图像分割和生成提供强大的骨干网络

:pushpin: 这是U-KAN：为医学图像分割和生成提供强大的骨干网络的官方PyTorch实现

[项目主页] [arXiv论文] [引用]

U-KAN：为医学图像分割和生成提供强大的骨干网络<br> 李晨鑫*, 刘欣宇*, 李武阳*, 王成*, 刘恒宇, 袁逸轩<sup>✉</sup><br>香港中文大学

我们探索了Kolmogorov-Anold网络（简称KAN）在改进视觉任务骨干网络方面的潜力。我们通过在分词后的中间表示上集成专用KAN层，对成熟的U-Net架构进行了研究、修改和重新设计，称之为U-KAN。严格的医学图像分割基准验证了U-KAN的优越性，即使在计算成本更低的情况下也能达到更高的准确率。我们进一步探讨了U-KAN作为扩散模型中U-Net噪声预测器的替代方案的潜力，展示了其在生成任务导向模型架构方面的适用性。这些努力揭示了宝贵的见解，并阐明了U-KAN在医学图像分割和生成方面可以成为强大骨干网络的前景。

📰新闻

[2024年6月] 为了更好地复现性能，对Seg_UKAN进行了一些修改。可以通过用新的train.py和archs.py内容替换原有文件来快速更新之前的代码。

[2024年6月] 模型检查点和训练日志已发布！

[2024年6月] U-KAN的代码和论文已发布！

💡主要特点

首次尝试将新兴KAN的优势融入成熟的U-Net架构，使其更加准确、高效和可解释。
分割U-KAN采用分词KAN块有效引导KAN算子，使其与现有的基于卷积的设计兼容。
扩散U-KAN作为改进的噪声预测器展示了其在生成任务和更广泛视觉设置中作为骨干网络的潜力。

🛠环境配置

git clone https://github.com/CUHK-AIM-Group/U-KAN.git
cd U-KAN
conda create -n ukan python=3.10
conda activate ukan
cd Seg_UKAN && pip install -r requirements.txt

提示A：我们使用pytorch=1.13.0和CUDA编译版本=11.6测试了该框架。其他版本应该也可以，但不能完全保证。

📚数据准备

BUSI：数据集可以在这里找到。

GLAS：数据集可以在这里找到。

CVC-ClinicDB：数据集可以在这里找到。

我们还提供了所有预处理过的数据集，无需进行任何进一步的数据处理。您可以直接下载并将它们放入数据目录。

最终的文件结构如下：

Seg_UKAN
├── inputs
│   ├── busi
│     ├── images
│           ├── malignant (1).png
|           ├── ...
|     ├── masks
│        ├── 0
│           ├── malignant (1)_mask.png
|           ├── ...
│   ├── GLAS
│     ├── images
│           ├── 0.png
|           ├── ...
|     ├── masks
│        ├── 0
│           ├── 0.png
|           ├── ...
│   ├── CVC-ClinicDB
│     ├── images
│           ├── 0.png
|           ├── ...
|     ├── masks
│        ├── 0
│           ├── 0.png
|           ├── ...

🔖评估分割U-KAN

您可以直接从检查点模型评估U-KAN。以下是使用我们在分割模型库中的预训练模型的快速使用示例：

下载预训练权重并将它们放入{args.output_dir}/{args.name}/model.pth
运行以下脚本：

cd Seg_UKAN
python val.py --name ${dataset}_UKAN --output_dir [YOUR_OUTPUT_DIR]

⏳训练分割U-KAN

你可以通过指定数据集名称 --dataset 和输入大小 --input_size 在单个GPU上简单地训练U-KAN。

cd Seg_UKAN
python train.py --arch UKAN --dataset ${dataset} --input_w ${input_size} --input_h ${input_size} --name ${dataset}_UKAN  --data_dir [YOUR_DATA_DIR]

例如，在 inputs 目录中的BUSI数据集上使用单个GPU训练分辨率为256x256的U-KAN：

cd Seg_UKAN
python train.py --arch UKAN --dataset busi --input_w 256 --input_h 256 --name busi_UKAN  --data_dir ./inputs

更多详细信息请参见Seg_UKAN/scripts.sh。请注意，glas的分辨率为512x512，与其他数据集（256x256）不同。

🎪分割模型库

我们提供了所有预训练模型检查点以下是已发布性能和检查点的概览。注意，单次运行的结果可能与论文中报告的平均结果有所不同。

方法	数据集	IoU	F1	检查点
Seg U-KAN	BUSI	65.26	78.75	链接
Seg U-KAN	GLAS	87.51	93.33	链接
Seg U-KAN	CVC-ClinicDB	85.61	92.19	链接

参数 ``--no_kan'' 表示将KAN层替换为MLP层的基线模型。请注意：性能偶尔出现不一致是合理的，多次运行的平均结果可以显示出更明显的趋势。

方法	层类型	IoU	F1	检查点
Seg U-KAN (--no_kan)	MLP层	63.49	77.07	链接
Seg U-KAN	KAN层	65.26	78.75	链接

🎇使用扩散U-KAN生成医学图像

请参考 Diffusion_UKAN

🛒待办事项

发布Seg U-KAN的代码。
发布Diffusion U-KAN的代码。
上传预训练检查点。

🎈致谢

非常感谢以下项目的巨大努力！

CKAN

📜引用

如果您发现这项工作对您的项目有帮助，请考虑引用以下论文：

@article{li2024u,
  title={U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation},
  author={Li, Chenxin and Liu, Xinyu and Li, Wuyang and Wang, Cheng and Liu, Hengyu and Yuan, Yixuan},
  journal={arXiv preprint arXiv:2406.02918},
  year={2024}
}