手把手教你实现语义分割

廖政飞
2022-05-28

上一期给大家介绍了语义分割的基础知识,下面小编就带大家一起实际操作一下。话不多说,让我们开始吧!

1、项目结构

本文介绍的语义分割项目为在PASCAL VOC 2012语义分割数据集上训练FCN-8s网络,并对训练好的网络进行测试。

PASCAL VOC 2012数据集的官方下载地址为:

http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar

项目代码的下载地址为:

https://github.com/AI-RESEARCH-GROUP/segmentation_tutorial

在下载好数据集和项目代码后,将两者放置于同一文件夹下。其中项目代码的目录结构如下:

项目代码的目录结构

  • checkpoints文件夹用于存放训练好的模型;
  • logs文件夹用于存放训练过程中记录损失值、mIoU等指标的日志文件;
  • pic文件夹中存放的是用于演示的图片;
  • result文件夹中存放的是对演示图片的分割结果;
  • arg文件夹中的args.py文件用于设置学习率、训练epoch数、batch size等;
  • dataset文件夹中的seg.py文件用于从PASCAL VOC 2012数据集中读入训练集和验证集数据;
  • func文件夹中的metric.py文件定义了计算mIoU的函数,early_stopping.py文件实现了当损失函数值下降小于某一值时,使训练停止;
  • model文件夹中的my_model.py文件定义了FCN-8s网络;
  • utils文件夹中的util.py文件定义了打印和记录日志等函数;
  • src文件夹下的train.py文件可执行FCN-8s模型的训练,demo.py文件可加载训练好的模型对pic文件夹中的图片进行分割,并将分割结果存放在result文件夹中。

2、基础设置

arg文件夹中的args.py文件是对项目的基础设置,如下图所示:

args.py文件

该文件中设定了训练时使用的GPU默认个数为1;设定训练epoch数为100;设定初始学习率为0.0001;设定训练及验证时的batch size为16,即训练和验证时每次将16张图像输入网络进行处理;最后一行代码设定了数据集的位置,需要说明的是,该代码中红色部分应更改为自己的数据集存放位置。

3、数据处理

dataset文件夹中的seg.py文件可实现读入数据集中的图片及标签,并经过填充和缩放处理后将其转变为tensor格式。

seg.py文件

此处代码实现了从PASCAL VOC 2012数据集的存放位置分别读入训练集和验证集中的图片和标签。

此处代码实现了依次对训练集和验证集中的图片和标签进行预处理和格式转换。首先通过填充操作将每张图片及标签变为正方形,然后再通过缩放操作将所有的图片和标签的尺寸统一,最后将图片和标签的格式转换为tensor。处理好的图片和标签即可输入网络进行训练。

4、FCN-8s 网络定义

model文件夹中的my_model.py文件定义了FCN-8s的网络结构和运算流程。

https://www.aimagician.cn/aimagician/api/img/download/20220628shou6-4a18.png?no_download_just_preview=1

此处代码定义了FCN-8s的前两个卷积块,卷积块中包含若干个卷积运算和一个池化运算,后三个卷积块也按相同的方式进行了定义。

此处代码定义了FCN-8s的三个转置卷积运算,用于扩大特征图的尺寸。

此处代码中的forward函数定义了FCN-8s的运算流程,即输入图像依次经过5个卷积块后,再分别由三个转置卷积分支融合网络深层和浅层的信息,最后得到网络的最终输出。

5、网络训练

src文件夹下的train.py实现了FCN-8s网络在PASCAL VOC数据集上的训练。在项目根目录下运行以下命令即可开始训练:

python train.py

此处代码设置了训练时采用cross entropy损失函数,设置采用Adam优化器,并采用StepLR学习率更新策略,即每经过20个epoch,将当前学习率乘以0.1。训练时也采用了early stopping策略,即在连续若干个训练周期中验证集损失都没有得到降低时,就停止模型的训练,此处将策略中的训练周期设置为10。

训练和验证阶段的流程为:首先训练一个epoch,训练过程中将每个batch的平均损失值和平均mIoU记录到log日志文件中;然后进行一次验证,并将验证过程中的每个batch的平均损失值和平均mIoU记录到log日志文件中;进一步判断此次验证集上的平均mIoU是否优于当前验证集上最优的mIoU,并据此确定是否保存此轮的模型参数。重复以上步骤,直到达到设定的训练epoch次数或满足early stopping策略的停止条件时,即停止模型的训练。

训练和验证阶段的流程在以下代码处实现:

6、模型演示

src文件夹下的demo.py实现了项目的演示功能,即采用训练好的FCN-8s网络对PASCAL VOC 2012测试集中的图片进行分割。在项目根目录下运行以下命令即可开始演示:

python demo.py

此处代码实现了从checkpoints文件夹下加载训练好的FCN-8s模型。

代码中的gen_img函数可实现对pic文件夹下的图片依次采用FCN-8s模型分割,生成对应的分割结果,并将分割结果保存在result文件夹中。

因此想要对测试集中某些图片进行分割,将对应的图片放入pic文件夹下,运行演示命令,即可在result文件夹下得到对应图片的分割结果。

部分图片的分割结果如下:

奇数行和偶数行分别为测试集图片及其语义分割结果图。在分割结果图中,不同类别的物体区域以不同的颜色标记出来。

-The End-

文案:廖政飞

指导老师:曹菁菁 赵强伟

排版:王雪 黄雅文

最新热门文章
关注AI魔术师
学习最新AI智能算法

联系我们

AI魔术师
bettycao@whut.edu.cn
湖北省武汉市