基于飞桨复现语义分割网络U²-Net,保持高分辨率特征同时降低计算成本

点击左上方蓝字关注我们

【飞桨开发者说】路星奎,沈阳化工大学信息工程学院研究生在读,计算机视觉技术爱好者,研究方向为图像分类、目标检测、图像分割等

本项目尝试使用PaddlePaddle复现U2-Net,将其改写为通用的图像分割模型,并整合到PaddleSeg开发套件中。

PaddleSeg GitHub:

https://github.com/PaddlePaddle/PaddleSeg

内容简介

本文包含以下4部分内容:

  • U²-Net介绍

  • U²-Net结构说明

  • 改写U²-Net适配语义分割任务

  • 使用U²-Net完成活性污泥显微图像分割任务

U²-Net介绍

  • 发表时间:2020年5月发表于模式识别(Pattern Recognition)

  • 地址:https://arxiv.org/abs/2005.09007

  • 论文研究领域:显著性目标检测

  • 论文出发点:能否为SOD设计一个新的网络,允许从头开始训练,并实现与现有预训练的主干网络相当或更好的性能?能否深入研究在保持高分辨率特征的同时,使用较低的内存和计算成本?

  • U²-Net在Reddit、Twitter火爆一时,号称2020年最强的静态背景分割算法,下图展示了U²-Net与其他方法的性能对比:

显著性目标检测、目标检测的区别:

  • 显著性目标检测:源于视觉显著性,即找到一张图像中人眼最关注的目标区域;

  • 目标检测:找到一张图像中感兴趣的目标位置与类别;

  • 目标数量:显著性目标检测一般只检一个目标,一般目标检测不会限制数量;

  • 目标类别:显著性目标检测不关心目标类别,一般目标检测会得到目标位置和类别。

U²-Net结构说明

1. U²-Net的优势

  • 提出的ReSidual U-blocks (RSU)中混合了不同大小的感受场,能够从不同的尺度捕捉更多的上下文信息;

  • 在不显著增加计算代价的情况下,增加了整个体系结构的深度;

  • 从头开始训练深度网络,无需使用图像分类任务中的Backbone。

2. U²-Net的设计思想

  • U2-Net是一种为SOD设计的两级嵌套的U型结构网络;

  • 在底层设计了一种新的ReSidual U-blocks (RSU),它能够在不降低特征图分辨率的情况下提取多尺度特征;

  • 在顶层设计了一种类似U-Net的结构,其中每一级都由RSU块填充。

3. U²-Net的整体结构

  • 6个stages的编码器

  • 5个stages的解码器

  • 解码器阶段的所有输出和编码器最后阶段输出的显著图融合模块

  • 以EN_6为轴,形成对称的U型结构,左侧为编码器,右侧为解码器、显著性图融合模块

4. U²-Net各部分说明

  • residual U-block/RSU

  1. 受U-Net的启发,提出了一种新的residual U-block/RSU,用于捕获stage内多尺度特征。

  2. RSU-L(Cin,M,Cout)的结构如图所示,其中L是编码器中的层数;Cin、Cout表示输入和输出通道,M表示RSU内部层中的通道数。

  3. 具体实现过程如下。

  4. 首先,经过第一个卷积层,将输入特征图X(H×W×Cin)变换为通道数为Cout的中间特征映射F1(X)。这是一个用于局部特征提取的普通卷积层;

  5. 其次,经过一种高度为L的类似U型网络的对称编解码器结构,它以中间特征映射F1(X)为输入,提取特征的多尺度上下文信息U(F1(X));

  6. 最后,从逐步下采样的特征图中提取多尺度特征,并通过逐步上采样、Concatenation、卷积将多尺度特征解码成高分辨率的特征图。该过程减轻了由小尺度特征图直接上采样造成的细节损失。

  7. 下图展示了RSU-L网络结构,以及RSU-L与ResNet的残差结构的类比:

  • 编码器阶段

  1. encoder stages分别由EN_1、EN_2、EN_3、EN_4、EN_5、EN_6组成:

  2. EN_1、EN_2、EN_3、EN_4,分别使用RSU-7、RSU-6、RSU-5、RSU-4,其中“7”、“6”、“5”、“4”表示RSU的深度(L)

  3. EN_5、EN_6,使用RSU-4F

  • 解码器阶段

  1. decoder stages分别由DE_1、DE_2、DE_3、DE_4、DE_5组成,其结构分别与对应encoder stages的EN_1、EN_2、EN_3、EN_4、EN_5相同。

  2. 每个decoder stage将来自其前一个stages的上采样特征映射与来自其对称encoder stage的上采样特征映射的级联作为输入

  • 显著性图融合模块(the saliency map fusion module)

  1. 用于生成显著性概率图

  2. 首先,通过一个3 * 3的卷积和sigmoid函数分别从DE_1、DE_2、DE_3、DE_4、DE_5、EN_6中得到6个显著性概率图S(1)side、S(2)side、S(3)side、S(4)side、S(5)side、S(6)side

  3. 最后,将这些显著性图上采样到输入图像大小,利用1×1卷积层和Sigmoid函数将它们融合,以生成最终的显著性概率图Sfuse

  • 解释说明

  1. L通常根据输入图像的空间分辨率进行配置。对于尺寸较大的图像,我们使用较大的L来捕捉更多的大尺度的特征信息

  2. EN5、EN6中特征映射的分辨率相对较低,进一步对这些特征映射进行下采样会导致丢失有用的上下文

  3. RSU-4F表示带有空洞卷积的RSU结构,其中使用空洞卷积代替池化与上采样的操作

改写U²-Net适配语义分割任务

  • U²-Net首要解决的是显著性目标检测,其主要特点是,将所有的目标看作一类,背景看作一类,可以将其看做是一个只有背景与目标的2类语义分割任务。

  • U²-Net在网络的最后使用sigmoid函数将特征值限定在了[0,1]以便适配2类语义分割任务,因此需要改写这个位置以便适配范围更广的语义分割任务。

  • 本文对U²-Net的后处理做了2个改进,一种是图1所示的直接使用1 * 1卷积替换原操作中的3 * 3 卷积,同时移除sigmoid函数(u2net.py);另一种是使用HRNet的后处理方式先使用1 * 1 卷积跨通道整合信息再使用1 * 1卷积完成最后的分类,同时移除sigmoid函数(u2net2.py)。经过试验发现两种方法的差距并不大,指标基本持平,但是第二种的参数量会适当的增加一些。本文选择第一种修改方式进行实验。

  • U²-Net已被整合到PaddleSeg组件中,这样可以更加方便的使用PaddleSeg提供的各种增强、调试、训练等功能

使用U²-Net实现活性污泥

显微图像分割任务

  • 活性污泥显微图像数据集介绍

  1. 来源:该数据集由国外某研究机构开源

  2. 数据集规模:由4X、10X、20X、40X倍数的83张活性污泥显微图像组成,本文按照4:1的比例划分训练集与测试集

  3. 类别:背景、絮体

下图展示了活性污泥显微图像与标签图像:

动手实践

1. 环境要求

PaddlePaddle >= 1.7.0

Python >= 3.5+

由于图像分割模型计算开销大,推荐在GPU版本的PaddlePaddle下使用PaddleSeg.

pip install -U paddlepaddle-gpu

安装PaddleSeg套件:安装过程

git clone https://github.com/PaddlePaddle/PaddleSeg.git

安装PaddleSeg依赖:

cd PaddleSegpip install -r requirements.txt

将活性污泥显微图像数据集按照4:1的比例划分成训练集与测试集,并生成对应的txt文件。

2. 模型选择参数配置

数据处理之后,需要根据模型配置相关的参数,关于各参数的具体含义,可访问链接:

https://aistudio.baidu.com/aistudio/projectdetail/894141 学习相关内容。

在PaddleSeg/configs/u2net_optic.yaml中的修改必要参数

  • DATASET:关于数据集的相关配置,如类别数、训练数据列表、测试数据列表

  • MODEL:模型配置

  • MULTI_LOSS_WEIGHT:模型输出权重配置

  • TRAIN_CROP_SIZE:训练时输入数据大小  

  • EVAL_CROP_SIZE:测试时输入数据大小

  • BATCH_SIZE:输入网络中的BATCH_SIZE,需要适配显存

  • SNAPSHOT_EPOCH: 阶段性保存EPOCH

  • NUM_EPOCHS:总的训练轮数

  • LOSS:损失函数类别

  • LR:学习率

3. 参数校验

切换根目录,所有代码都是基于PaddleSeg这一路径运行,所以需要切换路径至PaddleSeg。

cd PaddleSeg/

将已经修改好的配置文件移动至指定位置。

cp ../u2net_optic.yaml ./configs/cp ../config.py ./pdseg/utils/cp ../model_builder.py ./pdseg/models/cp ../u2net.py ./pdseg/models/modeling/

验证参数配置是否正确,验证指令如下:

python pdseg/check.py --cfg ./configs/u2net_optic.yaml

4. 模型训练

使用已经配置好的u2net_optic.yaml文件进行模型训练,训练指令如下:

python pdseg/train.py --use_gpu --cfg ./configs/u2net_optic.yaml --do_eval

本次项目中设置的是阶段性模型评估,同时保存评估结果最好的模型参数在下述目录:


PaddleSeg/saved_model/unet_optic/best_model。best_model文件夹下包含ppcls.pdmodel、ppcls.pdopt、ppcls.pdparams三个文件用来进行后续的评估推理使用。

5. 模型评估

使用已经配置好的u2net_optic.yaml文件进行模型评估,评估命令如下:

python pdseg/eval.py --use_gpu --cfg ./configs/u2net_optic.yaml

经过400个NUM_EPOCHS之后得到以下结果,可以看出 U²-Net效果可以与HRNet-W18持平,但是相比模型大小U²-Net要显得大一些。

6. 可视化展示

使用已经配置好的u2net_optic.yaml文件进行预测结果可视化,可视化命令如下:

python pdseg/vis.py  --cfg ./configs/u2net_optic.yaml

得到可视化结果之后,可以使用如下代码展示可视化结果:

import matplotlib.pyplot as plt
import os
import cv2 
def display(img_name):
    image_dir = os.path.join("./dataset/sludge/images", img_name)
    label_dir = os.path.join("./dataset/sludge/annotations_visualization",img_name)
    mask_dir = os.path.join("./visual", img_name)
    img_dir = [image_dir, label_dir, mask_dir]
    plt.figure(figsize=(15, 15))
    title = ['Image', 'label', 'Predict']
    for i in range(len(title)):
        plt.subplot(1, len(title), i+1)
        plt.title(title[i])
        img = cv2.imread(img_dir[i])
        b,g,r = cv2.split(img)
        img_rgb = cv2.merge([r,g,b])     
        plt.imshow(img_rgb)
        plt.axis('off')
    plt.show()
img_list=os.listdir("./visual")
for img_name in img_list:    
    display(img_name)

以下展示部分U2-Net的预测结果,可以看出预测结果与正确标注图的差距还是比较小的。

原始图像

标注图像

预测结果

心得体会

本项目使用了飞桨开源深度学习框架以及PaddleSeg套件,在AI Studio上完成了数据处理、模型训练、模型评估等工作。PaddleSeg套件让图像分割技术变得更为简单便捷,降低了开发者的上手难度。

在此强烈安利AI Studio。AI Studio是基于百度深度学习平台飞桨的人工智能学习与实训社区,提供在线编程环境、免费GPU算力、海量开源算法和开放数据,帮助开发者快速创建和部署模型。对于像笔者一样没有硬件条件的学习者是一个很大的助力。

整个项目包括数据集与相关代码已公开在AI Studio上,欢迎小伙伴们Fork。

https://aistudio.baidu.com/aistudio/projectdetail/897158

如在使用过程中有问题,可加入飞桨官方QQ群进行交流:1108045677。

如果您想详细了解更多飞桨的相关内容,请参阅以下文档。


官网地址:https://www.paddlepaddle.org.cn

·飞桨 PaddleSeg 项目地址·

GitHub: https://github.com/PaddlePaddle/PaddleSeg 

Gitee: https://gitee.com/paddlepaddle/PaddleSeg

·飞桨官网地址·

https://www.paddlepaddle.org.cn/


扫描二维码 | 关注我们

微信号 : PaddleOpenSource