（附鏈接）三個優(yōu)秀的語義分割框架 PyTorch實現(xiàn)

點擊左上方藍字關(guān)注我們

一個專注于目標檢測與深度學習知識分享的公眾號

編者薦語

文章基于動手深度學習項目講解了FCN進行自然圖像語義分割的流程，并對U-Net和Deeplab網(wǎng)絡進行了實驗，在Github和谷歌網(wǎng)盤上開源了代碼和預訓練模型，訓練和預測的腳本已經(jīng)做好封裝，讀者可以自行下載使用。

轉(zhuǎn)載自 | 計算機視覺工坊

1 前言

使用的VOC數(shù)據(jù)集鏈接開放在文章中，預訓練模型已上傳Github，環(huán)境我使用Colab pro，大家下載模型做預測即可。

代碼鏈接： https://github.com/lixiang007666/segmentation-learning-experiment-pytorch

使用方法：

1. 下載VOC數(shù)據(jù)集，將JPEGImages SegmentationClass兩個文件夾放入到data文件夾下。
2. 終端切換到目標目錄，運行python train.py -h查看訓練

(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

3. 預測需要手動修改predict.py中的模型

如果對FCN非常了解的，可以直接跳過d2l(動手學深度學習)的講解到最后一部分。

2 數(shù)據(jù)集

VOC數(shù)據(jù)集一般是用來做目標檢測，在2012版本中，加入了語義分割任務。

基礎數(shù)據(jù)集中包括：含有1464張圖片的訓練集，1449的驗證集和1456的測試集。一共有21類物體。

PASCAL VOC分割任務中，共有20個類別的對象，其他內(nèi)容作為背景類，其中紅色代表飛機類，黑色是背景，飛機邊界部分用米黃色（看著像白色）線條描繪，表示分割模糊區(qū)。

其中，分割標簽都是png格式的圖像，該圖像其實是單通道的顏色索引圖像，該圖像除了有一個單通道和圖像大小一樣的索引圖像外，還存儲了256個顏色值列表（調(diào)色板），每一個索引值對應調(diào)色板里一個RGB顏色值，因此，一個單通道的索引圖+調(diào)色板就能表示彩色圖。

原圖：

標簽：

挑選一張圖像可以發(fā)現(xiàn)，單張圖像分割類別不只兩類，且每張圖像類別不固定。

語義分割能對圖像中的每個像素分類。全卷積網(wǎng)絡 (fully convolutional network，F(xiàn)CN) 采用卷積神經(jīng)網(wǎng)絡實現(xiàn)了從圖像像素到像素類別的變換 。與我們之前在圖像分類或目標檢測部分介紹的卷積神經(jīng)網(wǎng)絡不同，全卷積網(wǎng)絡將中間層特征圖的高和寬變換回輸入圖像的尺寸：這是通過中引入的轉(zhuǎn)置卷積（transposed convolution）層實現(xiàn)的。因此，輸出的類別預測與輸入圖像在像素級別上具有一一對應關(guān)系：給定空間維上的位置，通道維的輸出即該位置對應像素的類別預測。

%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 網(wǎng)絡結(jié)構(gòu)

全卷積網(wǎng)絡先使用卷積神經(jīng)網(wǎng)絡抽取圖像特征，然后通過  卷積層將通道數(shù)變換為類別個數(shù)，最后再通過轉(zhuǎn)置卷積層將特征圖的高和寬變換為輸入圖像的尺寸。因此，模型輸出與輸入圖像的高和寬相同，且最終輸出的通道包含了該空間位置像素的類別預測。

下面，我們使用在ImageNet數(shù)據(jù)集上預訓練的ResNet-18模型來提取圖像特征，并將該網(wǎng)絡實例記為pretrained_net。該模型的最后幾層包括全局平均匯聚層和全連接層，然而全卷積網(wǎng)絡中不需要它們。

pretrained_net = torchvision.models.resnet18(pretrained=True)
list(pretrained_net.children())[-3:]

創(chuàng)建一個全卷積網(wǎng)絡實例net。它復制了Resnet-18中大部分的預訓練層，但除去最終的全局平均匯聚層和最接近輸出的全連接層。

net = nn.Sequential(*list(pretrained_net.children())[:-2])

給定高度和寬度分別為320和480的輸入，net的前向計算將輸入的高和寬減小至原來的，即10和15。

X = torch.rand(size=(1, 3, 320, 480))
net(X).shape

使用卷積層將輸出通道數(shù)轉(zhuǎn)換為Pascal VOC2012數(shù)據(jù)集的類數(shù)（21類)。最后，我們需要將要素地圖的高度和寬度增加32倍，從而將其變回輸入圖像的高和寬。

回想一下卷積層輸出形狀的計算方法：

由于且，我們構(gòu)造一個步幅為的轉(zhuǎn)置卷積層，并將卷積核的高和寬設為，填充為。

我們可以看到如果步幅為，填充為（假設是整數(shù))且卷積核的高和寬為，轉(zhuǎn)置卷積核會將輸入的高和寬分別放大倍。

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 初始化轉(zhuǎn)置卷積層

將圖像放大通常使用上采樣(upsampling)方法。雙線性插值（bilinear interpolation） 是常用的上采樣方法之一，它也經(jīng)常用于初始化轉(zhuǎn)置卷積層。

為了解釋雙線性插值，假設給定輸入圖像，我們想要計算上采樣輸出圖像上的每個像素。

首先，將輸出圖像的坐標 (??,??) 映射到輸入圖像的坐標 (??′,??′) 上。例如，根據(jù)輸入與輸出的尺寸之比來映射。請注意，映射后的 ??′ 和 ??′ 是實數(shù)。

然后，在輸入圖像上找到離坐標 (??′,??′) 最近的4個像素。

最后，輸出圖像在坐標 (??,??) 上的像素依據(jù)輸入圖像上這4個像素及其與 (??′,??′) 的相對距離來計算。

雙線性插值的上采樣可以通過轉(zhuǎn)置卷積層實現(xiàn)，內(nèi)核由以下bilinear_kernel函數(shù)構(gòu)造。限于篇幅，我們只給出bilinear_kernel函數(shù)的實現(xiàn)，不討論算法的原理。

def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = (torch.arange(kernel_size).reshape(-1, 1),
          torch.arange(kernel_size).reshape(1, -1))
    filt = (1 - torch.abs(og[0] - center) / factor) * \
           (1 - torch.abs(og[1] - center) / factor)
    weight = torch.zeros((in_channels, out_channels,
                          kernel_size, kernel_size))
    weight[range(in_channels), range(out_channels), :, :] = filt
    return weight

用雙線性插值的上采樣實驗它由轉(zhuǎn)置卷積層實現(xiàn)。我們構(gòu)造一個將輸入的高和寬放大2倍的轉(zhuǎn)置卷積層，并將其卷積核用bilinear_kernel函數(shù)初始化。

conv_trans = nn.ConvTranspose2d(3, 3, kernel_size=4, padding=1, stride=2,
                                bias=False)
conv_trans.weight.data.copy_(bilinear_kernel(3, 3, 4));

在全卷積網(wǎng)絡中，我們用雙線性插值的上采樣初始化轉(zhuǎn)置卷積層。對于 1×1卷積層，我們使用Xavier初始化參數(shù)。

W = bilinear_kernel(num_classes, num_classes, 64)
net.transpose_conv.weight.data.copy_(W);

3.3 訓練

損失函數(shù)和準確率計算與圖像分類中的并沒有本質(zhì)上的不同，因為我們使用轉(zhuǎn)置卷積層的通道來預測像素的類別，所以在損失計算中通道維是指定的。此外，模型基于每個像素的預測類別是否正確來計算準確率。

def loss(inputs, targets):
    return F.cross_entropy(inputs, targets, reduction='none').mean(1).mean(1)

num_epochs, lr, wd, devices = 5, 0.001, 1e-3, d2l.try_all_gpus()
trainer = torch.optim.SGD(net.parameters(), lr=lr, weight_decay=wd)
d2l.train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs, devices)

4 開源代碼和Dataset

數(shù)據(jù)集下載地址：http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar

輸入樣本：

輸出樣本：

運行Segmentat_pytorch.ipynp：

訓練：

!python3 train.py -m Unet -g 0

預測：

模型代碼包括FCN、U-Net和Deeplab的實現(xiàn)，大家可以更方便的更換模型訓練和預測。

DeeplabV3分割結(jié)果：

FCN分割結(jié)果：

U-Net分割結(jié)果：

記得點個Star哦！

5 總結(jié)

通過與分割標準圖像的對比，可以發(fā)現(xiàn)該模型的輸出分割圖像與分割標準圖像幾乎一致，同時模型的輸出分割圖像與原圖也較好的融合，說明該模型具有較好的準確性。

此外，從輸入圖像大小來看，該模型可以輸入任意大小的圖像，并輸出相同大小的已經(jīng)標簽好的分割圖像。由于是針對PASCAL VOC數(shù)據(jù)集圖像進行的分割，PASCAL VOC數(shù)據(jù)集中只支持20個類別（背景為第21個類別），所以在分割時，遇到不在20個類別中的事物都將其標為背景。

但總體來說，該模型對PASCAL VOC數(shù)據(jù)集的圖像分割達到了較高準確率。

6 參考

[1].https://zh-v2.d2l.ai/index.html

END

雙一流大學研究生團隊創(chuàng)建，專注于目標檢測與深度學習，希望可以將分享變成一種習慣！