三个优秀的语义分割框架 PyTorch实现

2021-09-14 16:08:52 浏览数 (1)

【GiantPandaCV导语】

本文基于动手深度学习项目讲解了FCN进行自然图像语义分割的流程，并对U-Net和Deeplab网络进行了实验，在Github和谷歌网盘上开源了代码和预训练模型，训练和预测的脚本已经做好封装，读者可以自行下载使用。

1 前言

使用的VOC数据集链接开放在文章中，预训练模型已上传Github，环境我使用Colab pro，大家下载模型做预测即可。

代码链接： https://github.com/lixiang007666/segmentation-learning-experiment-pytorch

使用方法：

  1. 下载VOC数据集，将JPEGImages SegmentationClass两个文件夹放入到data文件夹下。
  2. 终端切换到目标目录，运行python train.py -h查看训练
代码语言：javascript复制
(torch) qust116-jq@qustx-X299-WU8:~/语义分割$ python train.py -h
usage: train.py [-h] [-m {Unet,FCN,Deeplab}] [-g GPU]

choose the model

optional arguments:
  -h, --help            show this help message and exit
  -m {Unet,FCN,Deeplab}, --model {Unet,FCN,Deeplab}
                        输入模型名字
  -g GPU, --gpu GPU     输入所需GPU

选择模型和GPU编号进行训练，例如运行python train.py -m Unet -g 0

  1. 预测需要手动修改predict.py中的模型

如果对FCN非常了解的，可以直接跳过d2l(动手学深度学习)的讲解到最后一部分。

2 数据集

VOC数据集一般是用来做目标检测，在2012版本中，加入了语义分割任务。

基础数据集中包括：含有1464张图片的训练集，1449的验证集和1456的测试集。一共有21类物体。

PASCAL VOC分割任务中，共有20个类别的对象，其他内容作为背景类，其中红色代表飞机类，黑色是背景，飞机边界部分用米黄色（看着像白色）线条描绘，表示分割模糊区。

其中，分割标签都是png格式的图像，该图像其实是单通道的颜色索引图像，该图像除了有一个单通道和图像大小一样的索引图像外，还存储了256个颜色值列表（调色板），每一个索引值对应调色板里一个RGB颜色值，因此，一个单通道的索引图 调色板就能表示彩色图。

原图：

在这里插入图片描述

标签：

在这里插入图片描述

挑选一张图像可以发现，单张图像分割类别不只两类，且每张图像类别不固定。

3 全卷积神经网络

语义分割能对图像中的每个像素分类。全卷积网络 (fully convolutional network，FCN) 采用卷积神经网络实现了从图像像素到像素类别的变换 。与我们之前在图像分类或目标检测部分介绍的卷积神经网络不同，全卷积网络将中间层特征图的高和宽变换回输入图像的尺寸：这是通过中引入的转置卷积（transposed convolution）层实现的。因此，输出的类别预测与输入图像在像素级别上具有一一对应关系：给定空间维上的位置，通道维的输出即该位置对应像素的类别预测。

代码语言：javascript复制
%matplotlib inline
import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l

3.1 网络结构

全卷积网络先使用卷积神经网络抽取图像特征，然后通过

1times 1

卷积层将通道数变换为类别个数，最后再通过转置卷积层将特征图的高和宽变换为输入图像的尺寸。因此，模型输出与输入图像的高和宽相同，且最终输出的通道包含了该空间位置像素的类别预测。

下面，我们使用在ImageNet数据集上预训练的ResNet-18模型来提取图像特征，并将该网络实例记为pretrained_net。该模型的最后几层包括全局平均汇聚层和全连接层，然而全卷积网络中不需要它们。

代码语言：javascript复制
pretrained_net = torchvision.models.resnet18(pretrained=True)
list(pretrained_net.children())[-3:]

创建一个全卷积网络实例net。它复制了Resnet-18中大部分的预训练层，但除去最终的全局平均汇聚层和最接近输出的全连接层。

代码语言：javascript复制
net = nn.Sequential(*list(pretrained_net.children())[:-2])

给定高度和宽度分别为320和480的输入，net的前向计算将输入的高和宽减小至原来的

1/32

，即10和15。

代码语言：javascript复制
X = torch.rand(size=(1, 3, 320, 480))
net(X).shape

使用

1times1

卷积层将输出通道数转换为Pascal VOC2012数据集的类数（21类)。最后，我们需要将要素地图的高度和宽度增加32倍，从而将其变回输入图像的高和宽。

回想一下卷积层输出形状的计算方法：

由于

(320-64 16times2 32)/32=10

且

(480-64 16times2 32)/32=15

，我们构造一个步幅为

32

的转置卷积层，并将卷积核的高和宽设为

64

，填充为

16

。

我们可以看到如果步幅为

s

，填充为

s/2

（假设

s/2

是整数)且卷积核的高和宽为

2s

，转置卷积核会将输入的高和宽分别放大

s

倍。

代码语言：javascript复制
num_classes = 21
net.add_module('final_conv', nn.Conv2d(512, num_classes, kernel_size=1))
net.add_module('transpose_conv', nn.ConvTranspose2d(num_classes, num_classes,
                                    kernel_size=64, padding=16, stride=32))

3.2 初始化转置卷积层

将图像放大通常使用上采样(upsampling)方法。双线性插值（bilinear interpolation） 是常用的上采样方法之一，它也经常用于初始化转置卷积层。

为了解释双线性插值，假设给定输入图像，我们想要计算上采样输出图像上的每个像素。

首先，将输出图像的坐标 (

卷积神经网络 数据分析 深度学习 kernel

0 人点赞