基于UNet和camvid数据集的道路分割

基于UNet和camvid数据集的道路分割h(1.3.0+)：

背景
语义分割是深度学习中的一个非常重要的研究方向，并且UNet是语义分割中一个非常经典的模型。在本次博客中，我尝试用UNet对camvid dataset数据集进行道路分割，大致期望的效果如下：

原图

道路分割效果

本博客的代码参考了以下链接：

https://github.com/milesial/Pytorch-UNet
https://github.com/qubvel/segmentation_models.pytorch
1
2
数据集介绍及处理
之前的博客里，我几乎不怎么介绍数据集，因为用到的数据集比较简单；但是在使用camvid dataset的时候，我脑袋都大了，用了两三个小时才搞清楚这个数据集到底是啥情况。

数据集下载链接
虽然数据集的主页还可以访问，但是下载链接好像都失效了，所以最后还是用了aws上存储链接。

https://s3.amazonaws.com/fast-ai-imagelocal/camvid.tgz
1
数据说明
camvid数据集里包括三种重要信息，分别是RGB影像、语义分割图和标签说明。
RGB影像就不用多少了，为三通道RGB。
语义分割图为单通道，其中像素值代表了当前像素的类别，其对应关系存储在标签说明里。
标签说明对应了语义分割图像素值和类别的关系，如下：

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
数据处理
下载后数据后会有一个压缩包，包括images和labels，分别对应的是RGB的影像和像素的标签。
首先要做以下的一些处理，包括：

重命名labels的名称，去掉名称里的_P，保证labels和images的名称一致
将原始数据集按照7:2:1的规则，分割成train:valid:test
rename.py
import os,sys

cur_path = 'D:/camvid/camvid/labels' # 你的数据集路径

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

# 数据集路径
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]

# 确认分割正确
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))

# 创建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

# 导入库
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import numpy as np
import cv2
import matplotlib.pyplot as plt

# 设置数据集路径
DATA_DIR = 'D:/camvid/camvid' # 根据自己的路径来设置

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')

# 导入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里的道路设置为前景，而其他设置为背景
# 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']

# 这里必须设置一个mask的shape，因为前边的形状是（320,320）
return image, mask.reshape(1,320,320)

def __len__(self):
return len(self.ids)

# 数据增强
# 关于albumentations 怎么用我就不废话了
# 需要说明的是，我本身是打算用pytorch自带的transform
# 然而我实在没有搞明白，怎么同时对image和mask进行增强
# 如果连续调用两次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

# 考虑到我电脑的显卡大小，我降低了参数~~，无奈之举
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

# 设置train数据集
# 原谅我偷懒，并没有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)

# 准备训练，定义模型，我只做了两分类（偷懒）
# 另外，由于我修改了UNet模型，所以encoder部分，肯定不能用预训练模型
# 并且，我真的很反感每次都用预训练模型，没啥成就感。。。
net = UNet(n_channels=3, n_classes=1)

# 训练
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是专门用于二分类的，吴恩达的课程我记得前几节课就讲了
criterion = nn.BCELoss()

device = 'cuda'
for epoch in range(10):

net.train()
epoch_loss = 0

for data in train_loader:

# 修改一下数据格式
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))
# 结果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我才认为是对的
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

结果及分析
看一下最终结果，做一下分析讨论，总结经验。

结果
关于结果，这里随便展示几个吧，感觉还行。

分析
这是我第一次train分割的网络，有一些经验，写一写。

最开始train的时候，我比较心贪，用的原始分辨率的影像，720*960；结果网络参数太多，根本train不了，而且训练效果也不好；最后降采样才正常了，且效果变好了。
在训练之前，务必搞清楚数据集的格式，不然都不知道在train啥。
我在选择分割对象的时候，其实最开始也是用car,但是明显这个类别在影像里特别少，效果一直不好；最后选取了sky,road和wall这种样本较多的，效果才比较好；这说明样本数量还是很重要的。
————————————————
版权声明：本文为CSDN博主「Stone_Yannn」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/u012348774/article/details/104300366

基于UNet和camvid数据集的道路分割相关推荐

使用Pytorch搭建U-Net网络并基于DRIVE数据集训练(语义分割)学习笔记
使用Pytorch搭建U-Net网络并基于DRIVE数据集训练(语义分割)学习笔记 https://www.bilibili.com/video/BV1rq4y1w7xM?spm_id_from=33 ...
基于FCN，U-Net的深度学习医学影像分割算法（细胞分割算法）以及传统算法分析
本博文主要分析了基于传统计算机视觉的细胞分割算法和基于深度学习的细胞分割算法.主要针对医学影像分割算法.包括了FCN全卷积网络的细胞分割效果,U-Net细胞分割效果.查阅了采用深度学习分割医学影像的几 ...
基于U-Net图像分割的划痕缺陷分割（课程设计）
基于U-Net的划痕缺陷分割(课程设计) 论文:**U-Net: Convolutional Networks for Biomedical Image Segmentation** https:// ...
基于 U-Net 的医学影像分割算法综述
来源:UNknown知识库转载自新机器视觉仅做学术分享,若侵权,请联系删文 ▍一.医学图像分割简介医学影像分割是医学影像分析中的重要领域,也是计算机辅助诊断.监视.干预和治疗所必需的一环,其关键 ...
camvid数据集介绍_语义分割的数据集
背景语义分割指的是把图像中的每个像素都划分到某一个类别上. 实现算法上,有传统时代的grab cut.ML时代的TextonForest.DL时代的FCN .SegNet .Dilated Conv ...
基于U-Net+残差网络的语义分割缺陷检测
一.介绍基于深度学习的缺陷检测主要集中在场景识别.object detection等方法,近年来,推出了一系列优秀的语义分割模型,比如SegNet.FCN.U-Net等.语义分割模型被广泛的应用到场 ...
Python基于改进Unet的新冠肺炎等级分割系统（源码＆教程）
1.研究背景新冠肺炎给人类带来极大威胁,自动精确分割新冠肺炎CT图像感染区域可以辅助医生进行诊断治疗,但新冠肺炎的弥漫性感染.感染区域形状多变.与其他肺部组织极易混淆等给CT图像分割带来挑战.为此, ...
【PytorchLearning】基于 UNet 的肺部影像语义分割案例保姆教程
基于 UNet 的肺部影像分割一般而言,计算机视觉领域包含三大主流任务:分类.检测.分割.其中,分类任务对模型的要求较为简单,在之前的Pytorch入门教程中已进行了较为详尽的介绍,有兴趣的小伙伴可 ...
计算机毕业设计（1）-- 基于UNet的肝脏肿瘤分割
基于UNet的肝脏肿瘤分割文章地址:arxiv.org/pdf/1505.04597.pdf 1.UNet概述 UNet是一种用于图像分割的深度学习网络架构,最初由Olaf Ronneberge ...

基于UNet和camvid数据集的道路分割

基于UNet和camvid数据集的道路分割相关推荐

最新文章

热门文章