ai-station-code/dimaoshibie

SegFormer语义分割模型在Pytorch当中的实现

---

目录

1. 仓库更新 Top News
2. 相关仓库 Related code
3. 性能情况 Performance
4. 所需环境 Environment
5. 文件下载 Download
6. 训练步骤 How2train
7. 预测步骤 How2predict
8. 评估步骤 miou
9. 参考资料 Reference

Top News

2022-06:创建仓库、支持训练时评估、支持多backbone、支持step、cos学习率下降法、支持adam、sgd优化器选择、支持学习率根据batch_size自适应调整。

相关仓库

模型	路径
Unet	https://github.com/bubbliiiing/unet-pytorch
PSPnet	https://github.com/bubbliiiing/pspnet-pytorch
deeplabv3+	https://github.com/bubbliiiing/deeplabv3-plus-pytorch
hrnet	https://github.com/bubbliiiing/hrnet-pytorch
segformer	https://github.com/bubbliiiing/segformer-pytorch

性能情况

训练数据集	权值文件名称	测试数据集	输入图片大小	mIOU
VOC12+SBD	segformer_b0_weights_voc.pth	VOC-Val12	512x512	73.34
VOC12+SBD	segformer_b1_weights_voc.pth	VOC-Val12	512x512	76.80
VOC12+SBD	segformer_b2_weights_voc.pth	VOC-Val12	512x512	80.38

所需环境

torch==1.2.0

文件下载

训练所需的权值可在百度网盘中下载。
链接: https://pan.baidu.com/s/1tH4wdGnACtIuGOoXb0_rAw
提取码: tyjr

VOC拓展数据集的百度网盘如下：
链接: https://pan.baidu.com/s/1vkk3lMheUm6IjTXznlg7Ng
提取码: 44mk

训练步骤

a、训练voc数据集

1、将我提供的voc数据集放入VOCdevkit中（无需运行voc_annotation.py）。
2、在train.py中设置对应参数，默认参数已经对应voc数据集所需要的参数了，所以只要修改backbone和model_path即可。
3、运行train.py进行训练。

b、训练自己的数据集

1、本文使用VOC格式进行训练。
2、训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的SegmentationClass中。
3、训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。
4、在训练前利用voc_annotation.py文件生成对应的txt。
5、在train.py文件夹下面，选择自己要使用的主干模型。 6、注意修改train.py的num_classes为分类个数+1。
7、运行train.py即可开始训练。

预测步骤

a、使用预训练权重

1、下载完库后解压，在百度网盘下载权值，放入model_data，修改segformer.py的backbone和model_path之后再运行predict.py，输入。

img/street.jpg

可完成预测。
2、在predict.py里面进行设置可以进行fps测试、整个文件夹的测试和video视频检测。

b、使用自己训练的权重

1、按照训练步骤训练。
2、在segformer.py文件里面，在如下部分修改model_path、num_classes、backbone使其对应训练好的文件；model_path对应logs文件夹下面的权值文件，num_classes代表要预测的类的数量加1，backbone是所使用的主干特征提取网络。

_defaults = {
    #-------------------------------------------------------------------#
    #   model_path指向logs文件夹下的权值文件
    #   训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
    #   验证集损失较低不代表miou较高，仅代表该权值在验证集上泛化性能较好。
    #-------------------------------------------------------------------#
    "model_path"        : "model_data/segformer_b0_weights_voc.pth",
    #----------------------------------------#
    #   所需要区分的类的个数+1
    #----------------------------------------#
    "num_classes"       : 21,
    #----------------------------------------#
    #   所使用的的主干网络：
    #   b0、b1、b2、b3、b4、b5
    #----------------------------------------#
    "phi"               : "b0",
    #----------------------------------------#
    #   输入图片的大小
    #----------------------------------------#
    "input_shape"       : [512, 512],
    #-------------------------------------------------#
    #   mix_type参数用于控制检测结果的可视化方式
    #
    #   mix_type = 0的时候代表原图与生成的图进行混合
    #   mix_type = 1的时候代表仅保留生成的图
    #   mix_type = 2的时候代表仅扣去背景，仅保留原图中的目标
    #-------------------------------------------------#
    "mix_type"          : 0,
    #-------------------------------#
    #   是否使用Cuda
    #   没有GPU可以设置成False
    #-------------------------------#
    "cuda"              : True,
}

3、运行predict.py，输入

img/street.jpg

可完成预测。
4、在predict.py里面进行设置可以进行fps测试、整个文件夹的测试和video视频检测。

评估步骤

1、设置get_miou.py里面的num_classes为预测的类的数量加1。
2、设置get_miou.py里面的name_classes为需要去区分的类别。
3、运行get_miou.py即可获得miou大小。

Reference

https://github.com/ggyyzm/pytorch_segmentation
https://github.com/bonlime/keras-deeplab-v3-plus https://github.com/NVlabs/SegFormer