语义分割计算机视觉领域中的一个重要模块，它与之前的图像分类、目标检测任务不同，它是一个精细到每个像素块的一个图像任务，当然，它包括分类和定位。更多关于语义分割的方法和原理请读者自行搜索相关论文和博文，本文不做过多阐述。
1、数据集
在计算机视觉领域，Pascal VOC 2012 数据集是比较经典的，该数据集包含了目标检测、对象分割的数据及相关标签，本文使用的语义分割数据集就是来自于VOC2012。首先，需要下载该数据集：官网下载链接，注意，该数据集大概有2个G的大小，建议提前下载。
2、读取数据集
写一段程序来读取并显示一下相关数据图片及其标签：
def show_images(imgs, num_rows, num_cols, scale=2):
"""Plot a list of images."""
figsize = (num_cols * scale, num_rows * scale)
_, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
for i in range(num_rows):
for j in range(num_cols):
axes[i][j].imshow(imgs[i * num_cols + j].asnumpy())
axes[i][j].axes.get_xaxis().set_visible(False)
axes[i][j].axes.get_yaxis().set_visible(False)
plt.show()
root='F:/test/VOC2012_dataset/VOCtrainval_11-May-2012/VOCdevkit/VOC2012'
# 1、将训练图片和标注图标读进内存
def read_images(root,train=True):
if train:
txt_fname=root+"/ImageSets/Segmentation/train.txt"
else:
txt_fname = root + "/ImageSets/Segmentation/val.txt"
with open(txt_fname,'r') as f:
images=f.read().split()
features, labels = [None] * len(images), [None] * len(images)
for i, fname in enumerate(images):
features[i] = image_deal.imread('%s/JPEGImages/%s.jpg' % (root, fname))
labels[i] = image_deal.imread(
'%s/SegmentationClass/%s.png' % (root, fname))
return features, labels
n=5 #显示前5张图片
train_features,train_labels=read_images(root)
show_images(train_features[0:n] + train_labels[0:n],2,5)
for im in train_features[0:n]:
print(im.shape)
显示结果：
我们可以看到，每个图片均对应一个分割的图像，这个分割的图像就是标签(实际上是对一张图片的所以像素块分别赋予了一个label)。在打印的图像大小信息中我们也可以看到，每张图片的大小可能并不是一样的，在CNN训练中，为了批量训练，我们都会把它resize成同样的大小，但是这里不行，为什么？因为它是对每个pixel做了标签，如果resize(比如插值resize)之后，这样会出来一个新的pixel，那么新出来的pixel可能是介于两边的中间值，这就导致标签信息无法匹配上，结果不准。
2、裁剪
无法resize，那怎么办？我们可以注意到，每张图片的宽度好像都是500，为了能批量训练，我们可以使用剪切的方法来使得它们成为一样的大小。
def rand_crop(image, label, height, width):
data, rect = image_deal.random_crop(image, (width, height))
label = image_deal.fixed_crop(label, *rect)
return data, label
imgs=[]
for _ in range(3):
imgs += rand_crop(train_features[0], train_labels[0], 200, 300)
show_images(imgs,3,2)
结果：
3、每个物体和背景对应的RGB值(官方网站给的)
colormap = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
[0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
[64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
[64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
[0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
[0, 64, 128]]
colorclass = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
'diningtable', 'dog', 'horse', 'motorbike', 'person',
'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
4、查找pixel中每个类别的像素索引
colormap2label = nd.zeros(256 ** 3)
for i, color_map in enumerate(colormap):
colormap2label[(color_map[0] * 256 + color_map[1]) * 256 + color_map[2]] = i
def label_indices(colormap):
colormap = colormap.astype('int32')
idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256
+ colormap[:, :, 2])
return colormap2label[idx]
y = label_indices(train_labels[0])
print(y[105:115, 130:140])
结果：
可以看到是飞机的那块地方均被标记成了1，而背景则被标记成了0。为什么飞机是1？因为飞机排在第2个，那么索引就是1：
5、图像预处理
rgb_mean=nd.array([0.485,0.456,0.406]) #官方给的
rgb_std=nd.array([0.229,0.224,0.225])
def normlize(image):
return ((image.astype("float32")/255)-rgb_mean)/rgb_std
6、自定义语义分割数据集类
在mxnet中，可以通过继承Gluon提供的Dataset类自定义了语义分割数据集类VOCSegDataset。通过实现__getitem__函数，可以任意访问数据集中索引为idx的输入图像及其每个像素的类别索引。由于数据集中有些图像的尺寸可能小于随机裁剪所指定的输出尺寸，这些样本需要通过自定义的filter函数所移除。
class VOCSegDataset(gn.data.Dataset):
def __init__(self, is_train, crop_size, voc_dir):
self.crop_size = crop_size
data, labels = read_images(root=voc_dir, train=is_train)
data=self.filter(data)
self.data = [normlize(im) for im in data]
self.labels = self.filter(labels)
print('read ' + str(len(self.data)) + ' examples')
#只保留大于裁剪尺寸的图片
def filter(self, images):
return [img for img in images if (
img.shape[0] >= self.crop_size[0] and
img.shape[1] >= self.crop_size[1])]
# 每读一张图片，就把它裁剪，并把label的图片转成标签
def __getitem__(self, idx):
data, label = rand_crop(self.data[idx], self.labels[idx],
*self.crop_size)
data=data.transpose((2,0,1))
label=label_indices(label)
return data,label
def __len__(self):
return len(self.data)
看看这样处理之后有多少张图片，这里设裁剪的图片为(h,w)=(320,480)：
input_size=(320,480)
voc_train=VOCSegDataset(True,input_size,root)
voc_test=VOCSegDataset(False,input_size,root)
结果：
可以看到，训练集有1114张，测试集有1078张，对于这样的样本集，使用深度学习方法从头训练显然是不可能的，到了这里，就要提前想好后面要用到什么方法了——微调(Fine Turning)
7、定义批量读取
batch_size=5
train_data=gn.data.DataLoader(voc_train,batch_size,shuffle=True)
test_data=gn.data.DataLoader(voc_test,batch_size,shuffle=False)
# 查看一下数据维度
for data,label in train_data:
print(data.shape," ",label.shape)
break
结果：
可以看到，图片的维度与我们常见的一样，满足NCHW的要求，而标签则变成了3维的形状。
全卷积神经网络(Fully Convolutional Networks,FCN)
从上面标签是3维的我们可以知道，分割的任务与分类、检测的任务完全不一样，预测的标号不再是一个数字，而是每个pixel。那么在预测的时候，输出的结果也应该是一个3维的结果，因为要与输入的标签一一对应，但是，我们知道，CNN都是把一个3维的数据变成一个一维的标量，这显然与分割预测的结果不同，对于这种情况，转置卷积（transposed convolution）出现了，全卷积网络通过转置卷积层将中间层特征图的高和宽变换回输入图像的尺寸，从而令预测结果与输入图像在空间维（高和宽）上一一对应：给定空间维上的位置，通道维的输出即该位置对应像素的类别预测。（FCN就是在Forward时把维度变小，在Backward时把维度变大，卷积本身就是一个对偶函数，卷积的导数的导数还是卷积自己）
举个例子:一个(3,320,480)的图片，通过步长为2，padding为1的卷积之后，它的形状变成(3,160,240)，那么，通过转置卷积之后，它的形状变成(3,320,480)，这样就进行了还原。
1、转置卷积
gluon中已经实现了转置卷积的函数：
# 除了替换输出的通道数以外，其余的参数都不变，可以将输出还原为输入的大小
conv=gn.nn.Conv2D(channels=10,kernel_size=4,strides=2,padding=1)
conv_trans=gn.nn.Conv2DTranspose(channels=3,kernel_size=4,strides=2,padding=1) #图片最开始输入通道数为3
conv.initialize()
conv_trans.initialize()
x=nd.random_normal(shape=(1,3,16,16))
print(conv(x).shape," ",conv_trans(conv(x)).shape)
结果：
可以看到，我们定义了输入是(1,3,16,16)，通过卷积层之后，大小变成了(1,10,8,8)，再将结果通过转置卷积之后，结果变成了(1,3,16,16)，与输入的维度一样。
另外需要注意的是，在最后的卷积层我们同样使用Flatten或GAP来使得数据偏平化，使其能输入到FC中，而这样操作会损害空间信息，这对语义分割非常重要，其中一个解决办法是去掉不需要的池化层，并利用1X1的卷积层来替代FC。所以给定一个FCN，它需要做以下工作：
1）利用1X1的卷积替代FC；
2）去掉损失空间信息的池化层；
3）最后接上卷积转置层来得到需要输出的大小；
4）为了训练更快，可使用微调的办法（数据量大可以忽略这一条）。
2、下载预训练模型(ResNet18，可自行选择预训练模型)
pretrained_model=models.resnet18_v2(pretrained=True)
print(pretrained_model.features[-4:],pretrained_model.output) #打印看看最后几层
结果：
从上图的结果来看，我们是不需要GAP和FC的，所以要把它替换掉。
3、添加训练好的权重信息，并修改GAP和FC
net=gn.nn.Sequential()
for layer in pretrained_model.features[:-2]:
net.add(layer)
m=nd.random_normal(shape=(1,3,320,480))
print("input shape:",m.shape)
print("out shape:",net(m).shape)
# 添加1X1卷积和转置卷积
num_class=len(colorclass)# 几个类别
with net.name_scope():
net.add(gn.nn.Conv2D(num_class,1,1),
gn.nn.Conv2DTranspose(num_class,kernel_size=64,padding=16,strides=32)) # kernel_size最好大于32，padding=(kernel_size-strides)/2
结果：
可以看到，一张(3,320,480)通过预训练模型之后，输出的维度是(512,10,15)，其中，宽高均缩小了32倍，为什么要明确到32倍？以为最后转置卷积还原的时候会用到这个系数。
4、训练
因为上面我们把转置卷积的核大小设成了64，所以很难训练，而转置卷积我们可以把它看成是插值的操作，在实际操作中发现，把转置卷积初始化从双线性插值函数可以使得训练更加容易。
# 双线性插值
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 = np.ogrid[:kernel_size, :kernel_size]
filt = (1 - abs(og[0] - center) / factor) * (1 - abs(og[1] - center) / factor)
weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size),
dtype='float32')
weight[range(in_channels), range(out_channels), :, :] = filt
return nd.array(weight)
# 初始化
net[-2].initialize(init=init.Xavier()) # 1X1 conv
net[-1].initialize(init.Constant(bilinear_kernel(in_channels=num_class,out_channels=num_class,kernel_size=64)),ctx=ctx)
之前做分类的时候，使用了gn.loss.SoftmaxCrossEntropyLoss()这个函数，它默认会把结果Flatten化，所以，这里要加入axis=1的命令，其它的训练都是一样的。
cross_loss=gn.loss.SoftmaxCrossEntropyLoss(axis=1)
net.collect_params().reset_ctx(ctx)
trainer=gn.Trainer(net.collect_params(),'sgd',{"learning_rate":0.1,"wd":1e-3})
# 定义准确率
def accuracy(output,label):
return nd.mean(output.argmax(axis=1)==label).asscalar()
def evaluate_accuracy(data_iter, net,ctx):
acc_sum, n = 0.0, 0
for features,label in data_iter:
features = features.as_in_context(ctx)
label = label.as_in_context(ctx)
output=net(features)
acc_sum+=accuracy(output,label)
n += 1
return acc_sum/ n
def train(train_iter, test_iter, net, cross_loss, trainer, ctx, num_epochs):
print('training on', ctx)
for epoch in range(num_epochs):
train_l_sum, train_acc_sum, n,start = 0.0, 0.0, 0, time.time()
for Xs, ys in train_iter:
Xs=Xs.as_in_context(ctx)
ys=ys.as_in_context(ctx)
with ag.record():
output = net(Xs)
loss = nd.mean(cross_loss(output, ys))
loss.backward()
trainer.step(batch_size)
train_l_sum += loss.asscalar()
train_acc_sum+=accuracy(output,ys)
n += 1
test_acc = evaluate_accuracy(test_iter, net,ctx)
print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f ''time %.1f sec'
% (epoch + 1, train_l_sum / n, train_acc_sum / n,test_acc,
time.time() - start))
train(train_data,test_data,net,cross_loss,trainer,ctx,10)
训练结果：
训练完了别忘了保存一下模型，好不容易训练的别让它丢了：
net.save_parameters("FCN.params")# 保存模型
5、预测
预测的时候与图像分类类似，主要的区别在于我们需要在axis=1上做argmax，同时我们定义img2label的反函数，它将预测值转成图片。
net.load_parameters("FCN.params") #读取参数
def predict(image):
data=normlize(image)
data=data.transpose((2,0,1)).expand_dims(axis=0)
y_hat=net(data.as_in_context(ctx))
pred=nd.argmax(y_hat,axis=1)
return pred.reshape((pred.shape[1],pred.shape[2]))
def label2img(pred,colormap):
cm = nd.array(colormap, ctx=ctx, dtype='uint8')
X = pred.astype('int32')
return cm[X, :]
# 读取测试集前几张图片做预测
test_image,test_label=read_images(root,False)
image_6=[]
for i in range(6):
x=test_image[i]
pred=label2img(predict(x),colormap)
image_6+=[x,pred,test_label[i]]
show_images(image_6,6,3)
结果：
上图中，中间的是预测结果，最右边的是真实结果，预测结果离真实结果还有一段距离，可以多加几个epoch跑跑看。
下面附上所有源码：
import mxnet.ndarray as nd
import mxnet.gluon as gn
import mxnet.autograd as ag
import mxnet.initializer as init
import mxnet.image as image_deal
import matplotlib.pyplot as plt
import numpy as np
from mxnet.gluon.model_zoo import vision as models
import time
import mxnet as mx
# 显示图片
def show_images(imgs, num_rows, num_cols, scale=2):
"""Plot a list of images."""
figsize = (num_cols * scale, num_rows * scale)
_, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
for i in range(num_rows):
for j in range(num_cols):
axes[i][j].imshow(imgs[i * num_cols + j].asnumpy())
axes[i][j].axes.get_xaxis().set_visible(False)
axes[i][j].axes.get_yaxis().set_visible(False)
plt.show()
root = 'F:/test/VOC2012_dataset/VOCtrainval_11-May-2012/VOCdevkit/VOC2012'
# 1、将训练图片和标注图标读进内存
def read_images(root, train=True):
if train:
txt_fname = root + "/ImageSets/Segmentation/train.txt"
else:
txt_fname = root + "/ImageSets/Segmentation/val.txt"
with open(txt_fname, 'r') as f:
images = f.read().split()
features, labels = [None] * len(images), [None] * len(images)
for i, fname in enumerate(images):
features[i] = image_deal.imread('%s/JPEGImages/%s.jpg' % (root, fname))
labels[i] = image_deal.imread(
'%s/SegmentationClass/%s.png' % (root, fname))
return features, labels
n = 5 # 显示前5张图片
train_features, train_labels = read_images(root)
show_images(train_features[0:n] + train_labels[0:n], 2, 5)
# for im in train_features:
# print(im.shape)
# 2、裁剪
def rand_crop(image, label, height, width):
data, rect = image_deal.random_crop(image, (width, height))
label = image_deal.fixed_crop(label, *rect)
return data, label
imgs = []
for _ in range(3):
imgs += rand_crop(train_features[0], train_labels[0], 200, 300)
# print(imgs)
show_images(imgs, 3, 2)
# 3、每个物体和背景对应的RGB值
colormap = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
[0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
[64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
[64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
[0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
[0, 64, 128]]
colorclass = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
'diningtable', 'dog', 'horse', 'motorbike', 'person',
'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
print(len(colorclass)) # 共21类
# 4、查找pixel中每个类别的像素索引
colormap2label = nd.zeros(256 ** 3)
for i, color_map in enumerate(colormap):
colormap2label[(color_map[0] * 256 + color_map[1]) * 256 + color_map[2]] = i
def label_indices(colormap):
colormap = colormap.astype('int32')
idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256
+ colormap[:, :, 2])
return colormap2label[idx]
y = label_indices(train_labels[0])
print(y[105:115, 130:140])
#5、图像预处理
rgb_mean=nd.array([0.485,0.456,0.406]) #官方给的
rgb_std=nd.array([0.229,0.224,0.225])
def normlize(image):
return ((image.astype("float32")/255)-rgb_mean)/rgb_std
# for im in train_features:
# print(normlize(im))
#6、自定义语义分割数据集类
class VOCSegDataset(gn.data.Dataset):
def __init__(self, is_train, crop_size, voc_dir):
self.crop_size = crop_size
data, labels = read_images(root=voc_dir, train=is_train)
data=self.filter(data)
self.data = [normlize(im) for im in data]
self.labels = self.filter(labels)
print('read ' + str(len(self.data)) + ' examples')
#只保留大于裁剪尺寸的图片
def filter(self, images):
return [img for img in images if (
img.shape[0] >= self.crop_size[0] and
img.shape[1] >= self.crop_size[1])]
# 每读一张图片，就把它裁剪，并把label的图片转成标签
def __getitem__(self, idx):
data, label = rand_crop(self.data[idx], self.labels[idx],
*self.crop_size)
data=data.transpose((2,0,1))
label=label_indices(label)
return data,label
def __len__(self):
return len(self.data)
input_size=(320,480)
voc_train=VOCSegDataset(True,input_size,root)
voc_test=VOCSegDataset(False,input_size,root)
#7、定义批量读取
batch_size=5
train_data=gn.data.DataLoader(voc_train,batch_size,shuffle=True)
test_data=gn.data.DataLoader(voc_test,batch_size,shuffle=False)
# 查看一下数据维度
for data,label in train_data:
print(data.shape," ",label.shape)
break
'''---------FCN---------'''
#1、转置卷积实例
# 除了替换输出的通道数以外，其余的参数都不变，可以将输出还原为输入的大小
conv=gn.nn.Conv2D(channels=10,kernel_size=4,strides=2,padding=1)
conv_trans=gn.nn.Conv2DTranspose(channels=3,kernel_size=4,strides=2,padding=1) #图片最开始输入通道数为3
conv.initialize()
conv_trans.initialize()
x=nd.random_normal(shape=(1,3,16,16))
print(conv(x).shape," ",conv_trans(conv(x)).shape)
# 2、下载预训练模型resnet18
ctx=mx.gpu(0)
pretrained_model=models.resnet18_v2(pretrained=True,ctx=ctx)
print(pretrained_model.features[-4:],pretrained_model.output)
#3、添加训练好的权重信息，并修改GAP和FC
net=gn.nn.Sequential()
for layer in pretrained_model.features[:-2]:
net.add(layer)
m=nd.random_normal(shape=(1,3,320,480),ctx=ctx)
print("input shape:",m.shape)
print("out shape:",net(m).shape)
# 添加1X1卷积和转置卷积
num_class=len(colorclass)# 几个类别
with net.name_scope():
net.add(gn.nn.Conv2D(num_class,1,1),
gn.nn.Conv2DTranspose(num_class,kernel_size=64,padding=16,strides=32)) # kernel_size最好大于32，padding=(kernel_size-strides)/2
# 4、训练
# 双线性插值
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 = np.ogrid[:kernel_size, :kernel_size]
filt = (1 - abs(og[0] - center) / factor) * (1 - abs(og[1] - center) / factor)
weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size),
dtype='float32')
weight[range(in_channels), range(out_channels), :, :] = filt
return nd.array(weight)
# 初始化
net[-2].initialize(init=init.Xavier(),ctx=ctx) # 1X1 conv
net[-1].initialize(init.Constant(bilinear_kernel(in_channels=num_class,out_channels=num_class,
kernel_size=64)),ctx=ctx)
cross_loss=gn.loss.SoftmaxCrossEntropyLoss(axis=1)
net.collect_params().reset_ctx(ctx)
trainer=gn.Trainer(net.collect_params(),'sgd',{"learning_rate":0.1,"wd":1e-3})
# 定义准确率
def accuracy(output,label):
return nd.mean(output.argmax(axis=1)==label).asscalar()
def evaluate_accuracy(data_iter, net,ctx):
acc_sum, n = 0.0, 0
for features,label in data_iter:
features = features.as_in_context(ctx)
label = label.as_in_context(ctx)
output=net(features)
acc_sum+=accuracy(output,label)
n += 1
return acc_sum/ n
def train(train_iter, test_iter, net, cross_loss, trainer, ctx, num_epochs):
print('training on', ctx)
for epoch in range(num_epochs):
train_l_sum, train_acc_sum, n,start = 0.0, 0.0, 0, time.time()
for Xs, ys in train_iter:
Xs=Xs.as_in_context(ctx)
ys=ys.as_in_context(ctx)
with ag.record():
output = net(Xs)
loss = nd.mean(cross_loss(output, ys))
loss.backward()
trainer.step(batch_size)
train_l_sum += loss.asscalar()
train_acc_sum+=accuracy(output,ys)
n += 1
test_acc = evaluate_accuracy(test_iter, net,ctx)
print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f ''time %.1f sec'
% (epoch + 1, train_l_sum / n, train_acc_sum / n,test_acc,
time.time() - start))
# train(train_data,test_data,net,cross_loss,trainer,ctx,10)
# net.save_parameters("FCN.params")# 保存模型
# 5、预测
net.load_parameters("FCN.params") #读取参数
def predict(image):
data=normlize(image)
data=data.transpose((2,0,1)).expand_dims(axis=0)
y_hat=net(data.as_in_context(ctx))
pred=nd.argmax(y_hat,axis=1)
return pred.reshape((pred.shape[1],pred.shape[2]))
def label2img(pred,colormap):
cm = nd.array(colormap, ctx=ctx, dtype='uint8')
X = pred.astype('int32')
return cm[X, :]
# 读取测试集前几张图片做预测
test_image,test_label=read_images(root,False)
image_6=[]
for i in range(6):
x=test_image[i]
pred=label2img(predict(x),colormap)
image_6+=[x,pred,test_label[i]]
show_images(image_6,6,3)