基于UNet和camvid数据集的道路分割
基于UNet和camvid數(shù)據(jù)集的道路分割h(1.3.0+)：
背景
語義分割是深度學(xué)習(xí)中的一個非常重要的研究方向，并且UNet是語義分割中一個非常經(jīng)典的模型。在本次博客中，我嘗試用UNet對camvid dataset數(shù)據(jù)集進行道路分割，大致期望的效果如下：
原圖
道路分割效果
本博客的代碼參考了以下鏈接：
https://github.com/milesial/Pytorch-UNet
https://github.com/qubvel/segmentation_models.pytorch
1
2
數(shù)據(jù)集介紹及處理
之前的博客里，我?guī)缀醪辉趺唇榻B數(shù)據(jù)集，因為用到的數(shù)據(jù)集比較簡單；但是在使用camvid dataset的時候，我腦袋都大了，用了兩三個小時才搞清楚這個數(shù)據(jù)集到底是啥情況。
數(shù)據(jù)集下載鏈接
雖然數(shù)據(jù)集的主頁還可以訪問，但是下載鏈接好像都失效了，所以最后還是用了aws上存儲鏈接。
https://s3.amazonaws.com/fast-ai-imagelocal/camvid.tgz
1
數(shù)據(jù)說明
camvid數(shù)據(jù)集里包括三種重要信息，分別是RGB影像、語義分割圖和標(biāo)簽說明。
RGB影像就不用多少了，為三通道RGB。
語義分割圖為單通道，其中像素值代表了當(dāng)前像素的類別，其對應(yīng)關(guān)系存儲在標(biāo)簽說明里。
標(biāo)簽說明對應(yīng)了語義分割圖像素值和類別的關(guān)系，如下：
0?? ? Animal
1?? ? Archway
2?? ? Bicyclist
3?? ? Bridge
4 ?? ?Building
5?? ? Car
6?? ? CartLuggagePram
7?? ? Child
8?? ? Column_Pole
9 ?? ? Fence
10 LaneMkgsDriv
11 LaneMkgsNonDriv
12 Misc_Text
13 MotorcycleScooter
14 OtherMoving
15 ParkingBlock
16 Pedestrian
17 Road
18 RoadShoulder
19 Sidewalk
20 SignSymbol
21 Sky
22 SUVPickupTruck
23 TrafficCone
24 TrafficLight
25 Train
26 Tree
27 Truck_Bus
28 Tunnel
29 VegetationMisc
30 Void
31 Wall
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
數(shù)據(jù)處理
下載后數(shù)據(jù)后會有一個壓縮包，包括images和labels，分別對應(yīng)的是RGB的影像和像素的標(biāo)簽。
首先要做以下的一些處理，包括：
重命名labels的名稱，去掉名稱里的_P，保證labels和images的名稱一致
將原始數(shù)據(jù)集按照7:2:1的規(guī)則，分割成train:valid:test
rename.py
import os,sys
cur_path = 'D:/camvid/camvid/labels' # 你的數(shù)據(jù)集路徑
labels = os.listdir(cur_path)
for label in labels:
? ? old_label = str(label)
? ? new_label = label.replace('_P.png','.png')
? ? print(old_label, new_label)
? ? os.rename(os.path.join(cur_path,old_label),os.path.join(cur_path,new_label))
? ??
1
2
3
4
5
6
7
8
9
10
11
12
split_dataset.py
import os
import random
import shutil
# 數(shù)據(jù)集路徑
dataset_path = 'D:/camvid/camvid'
images_path = 'D:/camvid/camvid/images'
labels_path ? = 'D:/camvid/camvid/labels'
images_name = os.listdir(images_path)
images_num ?= len(images_name)
alpha ?= int( images_num ?* 0.7 )
beta ? = int( images_num ?* 0.9 )
print(images_num)
random.shuffle(images_name)
train_list = images_name[0:alpha]
valid_list = images_name[alpha:beta]
test_list ?= images_name[beta:images_num]
# 確認(rèn)分割正確
print('train list: ',len(train_list))
print('valid list: ',len(valid_list))
print('test list: ',len(test_list))
print('total num: ',len(test_list)+len(valid_list)+len(train_list))
# 創(chuàng)建train,valid和test的文件夾
train_images_path = os.path.join(dataset_path,'train_images')
train_labels_path ?= os.path.join(dataset_path,'train_labels')
if os.path.exists(train_images_path)==False:
? ? os.mkdir(train_images_path )
if os.path.exists(train_labels_path)==False:
? ? os.mkdir(train_labels_path)
valid_images_path = os.path.join(dataset_path,'valid_images')
valid_labels_path ?= os.path.join(dataset_path,'valid_labels')
if os.path.exists(valid_images_path)==False:
? ? os.mkdir(valid_images_path )
if os.path.exists(valid_labels_path)==False:
? ? os.mkdir(valid_labels_path)
test_images_path = os.path.join(dataset_path,'test_images')
test_labels_path ?= os.path.join(dataset_path,'test_labels')
if os.path.exists(test_images_path)==False:
? ? os.mkdir(test_images_path )
if os.path.exists(test_labels_path)==False:
? ? os.mkdir(test_labels_path)
# 拷貝影像到指定目錄
for image in train_list:
? ? shutil.copy(os.path.join(images_path,image), os.path.join(train_images_path,image))
? ? shutil.copy(os.path.join(labels_path,image), os.path.join(train_labels_path,image))
for image in valid_list:
? ? shutil.copy(os.path.join(images_path,image), os.path.join(valid_images_path,image))
? ? shutil.copy(os.path.join(labels_path,image), os.path.join(valid_labels_path,image))
for image in test_list:
? ? shutil.copy(os.path.join(images_path,image), os.path.join(test_images_path,image))
? ? shutil.copy(os.path.join(labels_path,image), os.path.join(test_labels_path,image))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
代碼
代碼鏈接：https://github.com/Yannnnnnnnnnnn/learnPyTorch/blob/master/road%20segmentation%20(camvid).ipynb
# 導(dǎo)入庫
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import numpy as np
import cv2
import matplotlib.pyplot as plt
# 設(shè)置數(shù)據(jù)集路徑
DATA_DIR = 'D:/camvid/camvid' # 根據(jù)自己的路徑來設(shè)置
x_train_dir = os.path.join(DATA_DIR, 'train_images')
y_train_dir = os.path.join(DATA_DIR, 'train_labels')
x_valid_dir = os.path.join(DATA_DIR, 'valid_images')
y_valid_dir = os.path.join(DATA_DIR, 'valid_labels')
x_test_dir = os.path.join(DATA_DIR, 'test_images')
y_test_dir = os.path.join(DATA_DIR, 'test_labels')
# 導(dǎo)入pytorch
import torch
from torch.utils.data import DataLoader
from torch.utils.data import Dataset as BaseDataset
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
# 自定義Dataloader
class Dataset(BaseDataset):
? ? """CamVid Dataset. Read images, apply augmentation and preprocessing transformations.
? ??
? ? Args:
? ? ? ? images_dir (str): path to images folder
? ? ? ? masks_dir (str): path to segmentation masks folder
? ? ? ? class_values (list): values of classes to extract from segmentation mask
? ? ? ? augmentation (albumentations.Compose): data transfromation pipeline?
? ? ? ? ? ? (e.g. flip, scale, etc.)
? ? ? ? preprocessing (albumentations.Compose): data preprocessing?
? ? ? ? ? ? (e.g. noralization, shape manipulation, etc.)
? ??
? ? """
? ??
? ? def __init__(
? ? ? ? ? ? self,?
? ? ? ? ? ? images_dir,?
? ? ? ? ? ? masks_dir,?
? ? ? ? ? ? augmentation=None,
? ? ):
? ? ? ? self.ids = os.listdir(images_dir)
? ? ? ? self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]
? ? ? ? self.masks_fps = [os.path.join(masks_dir, image_id) for image_id in self.ids]
? ? ? ??
? ? ? ? self.augmentation = augmentation
? ??
? ? def __getitem__(self, i):
? ? ? ? ? ? ? ??
? ? ? ? # read data
? ? ? ? image = cv2.imread(self.images_fps[i])
? ? ? ? image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
? ? ? ? mask = cv2.imread(self.masks_fps[i], 0)
? ? ? ??
? ? ? ? # 抱歉代碼寫的這么粗暴，意思就是講mask里的道路設(shè)置為前景，而其他設(shè)置為背景
? ? ? ? # road
? ? ? ? mask = (mask==17)
? ? ? ? mask = mask.astype('float') ??
? ? ? ??
? ? ? ? # apply augmentations
? ? ? ? if self.augmentation:
? ? ? ? ? ? sample = self.augmentation(image=image, mask=mask)
? ? ? ? ? ? image, mask = sample['image'], sample['mask']
? ? ? ?
? ? ? ? # 這里必須設(shè)置一個mask的shape，因為前邊的形狀是（320,320）
? ? ? ? return image, mask.reshape(1,320,320)
? ? ? ??
? ? def __len__(self):
? ? ? ? return len(self.ids)
# 數(shù)據(jù)增強
# 關(guān)于albumentations 怎么用我就不廢話了
# 需要說明的是，我本身是打算用pytorch自帶的transform
# 然而我實在沒有搞明白，怎么同時對image和mask進行增強
# 如果連續(xù)調(diào)用兩次transform，那么image和mask的增強方式都不一致，肯定不行
# 如果將[image;mask]堆砌在一起，放到transform里，image和mask的增強方式倒是一樣了，但是transform最后一步的toTensor會把mask歸一化，這肯定也是不行的
import albumentations as albu
def get_training_augmentation():
? ? train_transform = [
? ? ? ? albu.HorizontalFlip(p=0.5),
? ? ? ? albu.Resize(height=320, width=320, always_apply=True),
? ? ? ? albu.ShiftScaleRotate(scale_limit=0.1, rotate_limit=20, shift_limit=0.1, p=1, border_mode=0),
? ? ]
? ? return albu.Compose(train_transform)
def get_test_augmentation():
? ? train_transform = [
? ? ? ? albu.Resize(height=320, width=320, always_apply=True),
? ? ]
? ? return albu.Compose(train_transform) ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ??
augmented_dataset = Dataset(
? ? x_train_dir,?
? ? y_train_dir,?
? ? augmentation=get_training_augmentation(),?
)
# 定義UNet的基本模塊
# 代碼來自https://github.com/milesial/Pytorch-UNet
class DoubleConv(nn.Module):
? ? """(convolution => [BN] => ReLU) * 2"""
? ? def __init__(self, in_channels, out_channels):
? ? ? ? super().__init__()
? ? ? ? self.double_conv = nn.Sequential(
? ? ? ? ? ? nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
? ? ? ? ? ? nn.BatchNorm2d(out_channels),
? ? ? ? ? ? nn.ReLU(inplace=True),
? ? ? ? ? ? nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
? ? ? ? ? ? nn.BatchNorm2d(out_channels),
? ? ? ? ? ? nn.ReLU(inplace=True)
? ? ? ? )
? ? def forward(self, x):
? ? ? ? return self.double_conv(x)
class Down(nn.Module):
? ? """Downscaling with maxpool then double conv"""
? ? def __init__(self, in_channels, out_channels):
? ? ? ? super().__init__()
? ? ? ? self.maxpool_conv = nn.Sequential(
? ? ? ? ? ? nn.MaxPool2d(2),
? ? ? ? ? ? DoubleConv(in_channels, out_channels)
? ? ? ? )
? ? def forward(self, x):
? ? ? ? return self.maxpool_conv(x)
class Up(nn.Module):
? ? """Upscaling then double conv"""
? ? def __init__(self, in_channels, out_channels, bilinear=True):
? ? ? ? super().__init__()
? ? ? ? # if bilinear, use the normal convolutions to reduce the number of channels
? ? ? ? if bilinear:
? ? ? ? ? ? self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
? ? ? ? else:
? ? ? ? ? ? self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)
? ? ? ? self.conv = DoubleConv(in_channels, out_channels)
? ? def forward(self, x1, x2):
? ? ? ? x1 = self.up(x1)
? ? ? ? # input is CHW
? ? ? ? diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
? ? ? ? diffX = torch.tensor([x2.size()[3] - x1.size()[3]])
? ? ? ? x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
? ? ? ? ? ? ? ? ? ? ? ? diffY // 2, diffY - diffY // 2])
? ? ? ? # if you have padding issues, see
? ? ? ? # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
? ? ? ? # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
? ? ? ? x = torch.cat([x2, x1], dim=1)
? ? ? ? return self.conv(x)
? ? ? ??
class OutConv(nn.Module):
? ? def __init__(self, in_channels, out_channels):
? ? ? ? super(OutConv, self).__init__()
? ? ? ? self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
? ? def forward(self, x):
? ? ? ? return self.conv(x)
# UNet
class UNet(nn.Module):
? ? def __init__(self, n_channels, n_classes, bilinear=True):
? ? ? ? super(UNet, self).__init__()
? ? ? ? self.n_channels = n_channels
? ? ? ? self.n_classes = n_classes
? ? ? ? self.bilinear = bilinear
?? ??? ?# 考慮到我電腦的顯卡大小，我降低了參數(shù)~~，無奈之舉
? ? ? ? self.inc = DoubleConv(n_channels, 32)
? ? ? ? self.down1 = Down(32, 64)
? ? ? ? self.down2 = Down(64, 128)
? ? ? ? self.down3 = Down(128, 256)
? ? ? ? self.down4 = Down(256, 256)
? ? ? ? self.up1 = Up(512, 128, bilinear)
? ? ? ? self.up2 = Up(256, 64, bilinear)
? ? ? ? self.up3 = Up(128, 32, bilinear)
? ? ? ? self.up4 = Up(64, 32, bilinear)
? ? ? ? self.outc = OutConv(32, n_classes)
? ? ? ? self.out ?= torch.sigmoid #此處記得有sigmoid
? ? def forward(self, x):
? ? ? ? x1 = self.inc(x)
? ? ? ? x2 = self.down1(x1)
? ? ? ? x3 = self.down2(x2)
? ? ? ? x4 = self.down3(x3)
? ? ? ? x5 = self.down4(x4)
? ? ? ? x = self.up1(x5, x4)
? ? ? ? x = self.up2(x, x3)
? ? ? ? x = self.up3(x, x2)
? ? ? ? x = self.up4(x, x1)
? ? ? ? logits = self.outc(x)
? ? ? ? logits = self.out(logits)
? ? ? ? return logits
# 設(shè)置train數(shù)據(jù)集
# 原諒我偷懶，并沒有valid，因為我并沒有train多少epoch
train_dataset = Dataset(
? ? x_train_dir,?
? ? y_train_dir,?
? ? augmentation=get_training_augmentation(),?
)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
# 準(zhǔn)備訓(xùn)練，定義模型，我只做了兩分類（偷懶）
# 另外，由于我修改了UNet模型，所以encoder部分，肯定不能用預(yù)訓(xùn)練模型
# 并且，我真的很反感每次都用預(yù)訓(xùn)練模型，沒啥成就感。。。
net = UNet(n_channels=3, n_classes=1)
# 訓(xùn)練
from torch.autograd import Variable
net.cuda()
# 這里我說一下我是怎么train的
# 先lr=0.01,train大概40個epoch
# 然后lr=0.005,train大概40個epoch
# 最后在lr=0.0001,train大概20個epoch
optimizer = optim.RMSprop(net.parameters(), lr=0.4, weight_decay=1e-8)
# 這個loss是專門用于二分類的，吳恩達的課程我記得前幾節(jié)課就講了
criterion = nn.BCELoss()
device = 'cuda'
for epoch in range(10):
? ??
? ? net.train()
? ? epoch_loss = 0
? ??
? ? for data in train_loader:
? ? ? ??
? ? ? ? # 修改一下數(shù)據(jù)格式
? ? ? ? images,labels = data
? ? ? ? images = images.permute(0,3,1,2) # 交換通道順序
? ? ? ? images = images/255. # 把image的值歸一化到[0,1]
? ? ? ? images = Variable(images.to(device=device, dtype=torch.float32))
? ? ? ? labels = Variable(labels.to(device=device, dtype=torch.float32))
? ? ? ??
? ? ? ? pred = net(images)
? ? ? ??
? ? ? ? # 這里我不知道是看了哪里的代碼
? ? ? ? # 最開始犯傻寫成了 loss = criterion(pred.view(-1), labels.view(-1))
? ? ? ? # 結(jié)果loss很久都不下降
? ? ? ? # 還不知道為啥
? ? ? ? loss = criterion(pred, labels)
? ? ? ? epoch_loss += loss.item()
? ? ? ??
? ? ? ? optimizer.zero_grad()
? ? ? ? loss.backward()
? ? ? ? optimizer.step()
? ? ? ? print('loss: ', loss.item())
? ? ? ?
?# 測試
?test_dataset_noaug = Dataset(
? ? x_train_dir,?
? ? y_train_dir,
? ? augmentation=get_test_augmentation(),
? ? )
image, mask = test_dataset_noaug[77]
show_image = image
with torch.no_grad():
? ? image = image/255.
? ? image = image.astype('float32')
? ? image = torch.from_numpy(image)
? ? image = image.permute(2,0,1)
? ? image = image.to()
? ? print(image.shape)
? ??
? ? pred = net(image.unsqueeze(0).cuda())
? ? pred = pred.cpu()
# 大于0.5我才認(rèn)為是對的
pred = pred>0.5
# 展示圖如下
visualize(image=show_image,GT=mask[0,:,:],Pred=pred[0,0,:,:])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
結(jié)果及分析
看一下最終結(jié)果，做一下分析討論，總結(jié)經(jīng)驗。
結(jié)果
關(guān)于結(jié)果，這里隨便展示幾個吧，感覺還行。
分析
這是我第一次train分割的網(wǎng)絡(luò)，有一些經(jīng)驗，寫一寫。
最開始train的時候，我比較心貪，用的原始分辨率的影像，720*960；結(jié)果網(wǎng)絡(luò)參數(shù)太多，根本train不了，而且訓(xùn)練效果也不好；最后降采樣才正常了，且效果變好了。
在訓(xùn)練之前，務(wù)必搞清楚數(shù)據(jù)集的格式，不然都不知道在train啥。
我在選擇分割對象的時候，其實最開始也是用car,但是明顯這個類別在影像里特別少，效果一直不好；最后選取了sky,road和wall這種樣本較多的，效果才比較好；這說明樣本數(shù)量還是很重要的。
————————————————
版權(quán)聲明：本文為CSDN博主「Stone_Yannn」的原創(chuàng)文章，遵循CC 4.0 BY-SA版權(quán)協(xié)議，轉(zhuǎn)載請附上原文出處鏈接及本聲明。
原文鏈接：https://blog.csdn.net/u012348774/article/details/104300366
總結(jié)
以上是生活随笔為你收集整理的基于UNet和camvid数据集的道路分割的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： LABLEME UPDATE DAMO
- 下一篇： cJONS序列化工具解读二(数据解析)