keras-unet-collection

tensorflow.keras 实现的 U-net、V-net、U-net++、UNET 3+、Attention U-net、R2U-net、ResUnet-a、U^2-Net、TransUNET 和 Swin-UNET，可选使用 ImageNet 预训练的骨干网络。

---

keras_unet_collection.models 包含用于配置带有超参数选项的 Keras 模型的函数。

• U-net、U-net++、UNET 3+、Attention U-net 和 TransUNET 支持预训练的 ImageNet 骨干网络。
• U-net++、UNET 3+ 和 U^2-Net 支持深度监督。
• 其他选项和用例请参阅 用户指南。

keras_unet_collection.models	名称	参考文献
unet_2d	U-net	Ronneberger 等 (2015)
vnet_2d	V-net（修改为适用于 2D 输入）	Milletari 等 (2016)
unet_plus_2d	U-net++	Zhou 等 (2018)
r2_unet_2d	R2U-Net	Alom 等 (2018)
att_unet_2d	Attention U-net	Oktay 等 (2018)
resunet_a_2d	ResUnet-a	Diakogiannis 等 (2020)
u2net_2d	U^2-Net	Qin 等 (2020)
unet_3plus_2d	UNET 3+	Huang 等 (2020)
transunet_2d	TransUNET	Chen 等 (2021)
swin_unet_2d	Swin-UNET	Hu 等 (2021)

注意：这两个 Transformer 模型与 NumPy 1.20 不兼容；建议使用 NumPy 1.19.5。

---

keras_unet_collection.base 包含用于构建 Unet 变体基础架构（即不含模型头部）的函数，以便进行模型自定义和调试。

keras_unet_collection.base	备注
unet_2d_base、vnet_2d_base、unet_plus_2d_base、unet_3plus_2d_base、att_unet_2d_base、r2_unet_2d_base、resunet_a_2d_base、u2net_2d_base、transunet_2d_base、swin_unet_2d_base	这些函数接受输入张量和相应模型的超参数，并生成基础架构的输出张量。

---

keras_unet_collection.activations 和 keras_unet_collection.losses 提供额外的激活层和损失函数。

keras_unet_collection.activations	名称	参考文献
GELU	高斯误差线性单元（GELU）	Hendrycks 等 (2016)
Snake	Snake 激活	Liu 等 (2020)

keras_unet_collection.losses	名称	参考文献
dice	Dice 损失	Sudre 等 (2017)
tversky	Tversky 损失	Hashemi 等 (2018)
focal_tversky	Focal Tversky 损失	Abraham 等 (2019)
ms_ssim	多尺度结构相似度指数损失	Wang 等 (2003)
iou_seg	分割任务中的交并比（IoU）损失	Rahman 和 Wang (2016)
iou_box	（广义）IoU 损失，用于目标检测	Rezatofighi 等 (2019)
triplet_1d	半硬三元组损失（实验性）
crps2d_tf	CRPS 损失（实验性）

安装

该项目托管在 PyPI 上，因此可以通过以下命令安装：

pip install keras-unet-collection

使用

from keras_unet_collection import models
# 例如 models.unet_2d(...)

• 注意：目前支持的骨干网络包括：VGG[16,19]、ResNet[50,101,152]、ResNet[50,101,152]V2、DenseNet[121,169,201] 和 EfficientNetB[0-7]。详细信息请参阅 Keras Applications。

• 注意：本包生成的神经网络可能包含不属于 TensorFlow 的自定义层。建议保存和加载模型权重。

• 变更日志

示例

• Jupyter 笔记本作为示例提供，位于 examples 目录下：

  • 带有 VGG16 骨干的 Attention U-net [链接]。

  • 带有深度监督、分类引导模块和混合损失的 UNET 3+ [链接]。

  • 基于 Vision-Transformer 的示例正在开发中，可在 keras-vision-transformer 中找到。

依赖项

• TensorFlow 2.5.0、Keras 2.5.0、NumPy 1.19.5。

• （示例可选）Pillow、matplotlib 等。

概述

U-net 是一种具有编码器-解码器架构和跳跃连接的卷积神经网络，其概念大致属于“全卷积网络”。U-net 最初是为医学图像的语义分割而提出的，后来被改进以解决更广泛的网格状学习问题。

U-net 及其许多变体以三或四维张量作为输入，并产生相同形状的输出。这些模型的一个技术亮点是从下采样层到上采样层的跳跃连接，这有助于重建高分辨率的网格状输出。

联系方式

Yingkai (Kyle) Sha <yingkai@eoas.ubc.ca> <yingkaisha@gmail.com>

许可证

MIT 许可证

引用

• Sha, Y., 2021: Keras-unet-collection. GitHub 仓库，于 2021年9月4日访问，https://doi.org/10.5281/zenodo.5449801

@misc{keras-unet-collection,
  author = {Sha, Yingkai},
  title = {Keras-unet-collection},
  year = {2021},
  publisher = {GitHub},
  journal = {GitHub 仓库},
  howpublished = {\url{https://github.com/yingkaisha/keras-unet-collection}},
  doi = {10.5281/zenodo.5449801}
}

keras-unet-collection 快速上手指南

keras-unet-collection 是一个基于 tensorflow.keras 的开源库，集成了多种先进的 U-Net 变体（如 U-net++, Attention U-net, TransUNET, Swin-UNET 等），支持预训练的 ImageNet 骨干网络，适用于医学图像分割及其他网格化学习问题。

环境准备

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

• 操作系统：Linux, macOS 或 Windows
• Python 版本：建议 Python 3.7 - 3.9
• 核心依赖：
  • TensorFlow 2.5.0+
  • Keras 2.5.0+
  • NumPy 1.19.5 (⚠️ 重要提示：由于 Transformer 模型兼容性问题，请勿使用 NumPy 1.20 及以上版本，推荐锁定为 1.19.5)

安装步骤

您可以直接通过 PyPI 进行安装。国内开发者建议使用清华或阿里镜像源以加速下载。

标准安装命令：

pip install keras-unet-collection

使用国内镜像源加速安装（推荐）：

pip install keras-unet-collection -i https://pypi.tuna.tsinghua.edu.cn/simple

安装完成后，建议检查 NumPy 版本是否符合要求：

pip show numpy
# 如果版本高于 1.19.5，请执行降级：
pip install numpy==1.19.5

基本使用

该库提供了丰富的模型构建函数，位于 keras_unet_collection.models 模块中。以下是最简单的使用示例，展示如何构建一个标准的 2D U-Net 模型。

1. 导入模块并构建模型

from keras_unet_collection import models

# 构建一个基础的 2D U-Net 模型
# input_shape: 输入图像尺寸 (高，宽，通道数)
# n_labels: 分割类别的数量
model = models.unet_2d(
    input_shape=(256, 256, 3),
    n_labels=2,
    filter_num_down=[64, 128, 256, 512],
    filter_num_up=[512, 256, 128, 64],
    activation='ReLU',
    output_activation='Softmax',
    batch_norm=True,
    pool=True,
    unpool=True,
    name='unet'
)

model.summary()

2. 使用预训练骨干网络（可选）

部分模型（如 U-net, Attention U-net, TransUNET 等）支持加载 ImageNet 预训练的骨干网络以提升性能。

from keras_unet_collection import models

# 构建带有 VGG16 骨干网络的 Attention U-Net
model = models.att_unet_2d(
    input_shape=(256, 256, 3),
    n_labels=2,
    filter_num=[64, 128, 256, 512],
    backbone='VGG16',          # 指定骨干网络
    weights='imagenet',        # 加载 ImageNet 预训练权重
    freeze_backbone=True,      # 可选：冻结骨干网络参数
    name='att_unet_vgg16'
)

3. 编译与训练提示

由于该库可能包含 TensorFlow 原生不支持的自定义层，强烈建议保存和加载模型权重（weights）而非整个模型结构。

# 编译模型（需配合库中提供的损失函数或自定义损失）
from keras_unet_collection import losses

model.compile(
    optimizer='adam',
    loss=losses.dice,  # 使用库中提供的 Dice Loss
    metrics=['accuracy']
)

# 训练...
# model.fit(x_train, y_train, ...)

# 推荐保存权重的方法
model.save_weights('my_model_weights.h5')

# 加载权重时，需先重建相同结构的模型，再加载权重
# new_model = models.unet_2d(...) 
# new_model.load_weights('my_model_weights.h5')

常用模型列表

函数名	模型名称	特点
unet_2d	U-net	经典架构
unet_plus_2d	U-net++	支持深度监督
att_unet_2d	Attention U-net	引入注意力机制
transunet_2d	TransUNET	基于 Vision Transformer
swin_unet_2d	Swin-UNET	基于 Swin Transformer
u2net_2d	U^2-Net	嵌套 U 结构，支持深度监督

更多高级用法（如深度监督、混合损失函数等）请参考官方 GitHub 仓库中的 examples 目录。