一.简介
PSPNet模型采用了PSP模块,因此取名为PSPNet。该模型提出的金字塔池化模块(Pyramid Pooling Module)能够聚合不同区域上下文信息,从而提高获取全局信息的能力。
PSP结构的功能是将获取到的特征层按照不同的尺寸划分成不同的网络,每个网络内各自进行平均池化。
在PSPNet中,网络会将输入进来的特征层分别划分为6x6,3x3,2x2,1x1的网格,对应图片中绿色,蓝色,橙色,红色的输出。比如一个特征图的边长为72,划分为6x6的网格后,每个网格边长就是12;划分为3x3的网格后,每个网格边长就是24;划分为2x2的网格后,每个网格边长就是36;划分为1x1的网格后,每个网格边长就是72。
二.编码器Encoder
from tensorflow.keras.layers import *
from tensorflow.keras import Model
# 1.基于 vgg16 的 encoder编码器
def encoder_vgg16(height=416, width=416):
img_input = Input(shape=(height, width, 3))
# block1
# 416,416,3 -- 208,208,64
x = Conv2D(64, 3, padding='same', activation='relu', name='b1_c1')(img_input)
x = Conv2D(64, 3, padding='same', activation='relu', name='b1_c2')(x)
x = MaxPooling2D((2, 2), strides=2, name='b1_pool')(x)
out1 = x
# block2
# 208,208,64 -- 104,104,128
x = Conv2D(128, 3, padding='same', activation='relu', name='b2_c1')(x)
x = Conv2D(128, 3, padding='same', activation='relu', name='b2_c2')(x)
x = MaxPooling2D((2, 2), strides=2, name='b2_pool')(x)
out2 = x
# block3
# 104,104,128 -- 52,52,256
x = Conv2D(256, 3, padding='same', activation='relu', name='b3_c1')(x)
x = Conv2D(256, 3, padding='same', activation='relu', name='b3_c2')(x)
x = Conv2D(256, 3, padding='same', activation='relu', name='b3_c3')(x)
x = MaxPooling2D((2, 2), strides=2, name='b3_pool')(x)
out3 = x
# block4
# 52,52,256 -- 26,26,512
x = Conv2D(512, 3, padding='same', activation='relu', name='b4_c1')(x)
x = Conv2D(512, 3, padding='same', activation='relu', name='b4_c2')(x)
x = Conv2D(512, 3, padding='same', activation='relu', name='b4_c3')(x)
x = MaxPooling2D((2, 2), strides=2, name='b4_pool')(x)
out4 = x
# block5
# 26,26,512 -- 13,13,512
x = Conv2D(512, 3, padding='same', activation='relu', name='b5_c1')(x)
x = Conv2D(512, 3, padding='same', activation='relu', name='b5_c2')(x)
x = Conv2D(512, 3, padding='same', activation='relu', name='b5_c3')(x)
x = MaxPooling2D((2, 2), strides=2, name='b5_pool')(x)
out5 = x
return img_input, [out1,out2,out3,out4,out5]
# 2.基于 MobilenetV1 的 encoder编码器(DepthwiseConv2D + Conv1x1 实现)
def conv_block(inputs, filters, kernel, strides):
'''
:param inputs: 输入的 tensor
:param filters: 卷积核数量
:param kernel: 卷积核大小
:param strides: 卷积步长
:return:
'''
x = ZeroPadding2D(1)(inputs)
x = Conv2D(filters, kernel, strides, padding='valid', use_bias=False)(x)
x = BatchNormalization()(x)
x = ReLU(max_value=6)(x)
return x
def dw_pw_block(inputs, dw_strides, pw_filters, name):
'''
:param inputs: 输入的tensor
:param dw_strides: 深度卷积的步长
:param pw_filters: 逐点卷积的卷积核数量
:param name:
:return:
'''
x = ZeroPadding2D(1)(inputs)
# dw
x = DepthwiseConv2D((3, 3), dw_strides, padding='valid', use_bias=False, name=name)(x)
x = BatchNormalization()(x)
x = ReLU(max_value=6)(x)
# pw
x = Conv2D(pw_filters, (1, 1), 1, padding='valid', use_bias=False)(x)
x = BatchNormalization()(x)
x = ReLU(max_value=6)(x)
return x
def encoder_MobilenetV1_1(height=416, width=416):
img_input = Input(shape=(height, width, 3))
# block1:con1 + dw_pw_1
# 416,416,3 -- 208,208,32 -- 208,208,64
x = conv_block(img_input, 32, (3, 3), (2, 2))
x = dw_pw_block(x, 1, 64, 'dw_pw_1')
out1 = x
# block2:dw_pw_2
# 208,208,64 -- 104,104,128
x = dw_pw_block(x, 2, 128, 'dw_pw_2_1')
x = dw_pw_block(x, 1, 128, 'dw_pw_2_2')
out2 = x
# block3:dw_pw_3
# 104,104,128 -- 52,52,256
x = dw_pw_block(x, 2, 256, 'dw_pw_3_1')
x = dw_pw_block(x, 1, 256, 'dw_pw_3_2')
out3 = x
# block4:dw_pw_4
# 52,52,256 -- 26,26,512
x = dw_pw_block(x, 2, 512, 'dw_pw_4_1')
for i in range(5):
x = dw_pw_block(x, 1, 512, 'dw_pw_4_' + str(i + 2))
out4 = x
# block5:dw_pw_5
# 26,26,512 -- 13,13,1024
x = dw_pw_block(x, 2, 1024, 'dw_pw_5_1')
x = dw_pw_block(x, 1, 1024, 'dw_pw_5_2')
out5 = x
return img_input, [out1,out2,out3,out4,out5]
# 3.基于 MobilenetV1 的 encoder编码器(SeparableConv2D实现)
def sp_block(x, dw_strides, pw_filters, name):
'''
:param x: 输入的 tensor
:param dw_strides: 深度卷积的步长
:param pw_filters: 逐点卷积核的数量
:param name:
:return:
'''
x = ZeroPadding2D(1)(x)
x = SeparableConv2D(pw_filters, (3, 3), dw_strides, use_bias=False, name=name)(x)
x = BatchNormalization()(x)
x = ReLU(max_value=6)(x)
return x
def encoder_MobilenetV1_2(height=416, width=416):
img_input = Input(shape=(height, width, 3))
# block1:con1 + dw_pw_1
# 416,416,3 -- 208,208,32 -- 208,208,64
x = conv_block(img_input, 32, (3, 3), (2, 2))
x = sp_block(x, 1, 64, 'dw_pw_1')
out1 = x
# block2:dw_pw_2
# 208,208,64 -- 104,104,128
x = sp_block(x, 2, 128, 'dw_pw_2_1')
x = sp_block(x, 1, 128, 'dw_pw_2_2')
out2 = x
# block3:dw_pw_3
# 104,104,128 -- 52,52,256
x = sp_block(x, 2, 256, 'dw_pw_3_1')
x = sp_block(x, 1, 256, 'dw_pw_3_2')
out3 = x
# block4:dw_pw_4
# 52,52,256 -- 26,26,512
x = sp_block(x, 2, 512, 'dw_pw_4_1')
for i in range(5):
x = sp_block(x, 1, 512, 'dw_pw_4_' + str(i + 2))
out4 = x
# block5:dw_pw_5
# 26,26,512 -- 13,13,1024
x = sp_block(x, 2, 1024, 'dw_pw_5_1')
x = sp_block(x, 1, 1024, 'dw_pw_5_2')
out5 = x
return img_input, [out1,out2,out3,out4,out5]
# 4.基于 MobilenetV2 的 encoder编码器
# 倒残差结构
def inv_res_block(inputs,filters,strides,expansion,is_add,block_id=1,rate=1):
'''
:param inputs: 输入的 tensor
:param filters: 深度可分离卷积卷积核数量
:param strides: 深度可分离卷积步长
:param expansion: 倒残差通道扩张的倍数
:param is_add: 是否进行残差相加
:param rate: 空洞卷积扩张率
:param block_id:
:return:
'''
in_channels = inputs.shape[-1]
x = inputs
# 如果是第0个倒残差块,不进行通道扩张
if block_id:
x = Conv2D(in_channels*expansion,kernel_size=1,padding='same',use_bias=False)(x)
x = BatchNormalization()(x)
x = ReLU(max_value=6)(x)
# 深度可分离卷积提取特征
x = DepthwiseConv2D(kernel_size=3,strides=strides,padding='same',use_bias=False,dilation_rate=(rate,rate))(x)
x = BatchNormalization()(x)
x = ReLU(max_value=6)(x)
# 使用 1x1 卷积进行通道缩小
x = Conv2D(filters,kernel_size=1,padding='same',use_bias=False)(x)
x = BatchNormalization()(x)
if is_add:
return Add()([inputs,x])
return x
# img_input=size,residual_1 = size/4,x = size/8
def encoder_MobilenetV2(height=416, width=416):
img_input = Input(shape=(height, width, 3))
# 416,416,3 -- 208,208,32
x = Conv2D(32,3,2,padding='same',use_bias=False)(img_input)
x = BatchNormalization()(x)
x = ReLU(max_value=6)(x)
# 208,208,32 -- 208,208,16;首个倒残差块内部不进行通道先扩张后缩小
x = inv_res_block(x,filters=16,strides=1,expansion=1,is_add=False,block_id=0)
# 208,208,16 -- 104,104,24
x = inv_res_block(x,filters=24,strides=2,expansion=6,is_add=False)
x = inv_res_block(x,filters=24,strides=1,expansion=6,is_add=True)
residual_1 = x
# 104,104,24 -- 52,52,32
x = inv_res_block(x, filters=32, strides=2, expansion=6, is_add=False)
x = inv_res_block(x, filters=32, strides=1, expansion=6, is_add=True)
x = inv_res_block(x, filters=32, strides=1, expansion=6, is_add=True)
# 52,52,32 -- 52,52,64
x = inv_res_block(x, filters=64, strides=1, expansion=6, is_add=False)
x = inv_res_block(x, filters=64, strides=1, expansion=6, is_add=True,rate=2)
x = inv_res_block(x, filters=64, strides=1, expansion=6, is_add=True,rate=2)
x = inv_res_block(x, filters=64, strides=1, expansion=6, is_add=True,rate=2)
# 52,52,64 -- 52,52,96
x = inv_res_block(x, filters=96, strides=1, expansion=6, is_add=False, rate=2)
x = inv_res_block(x, filters=96, strides=1, expansion=6, is_add=True, rate=2)
x = inv_res_block(x, filters=96, strides=1, expansion=6, is_add=True, rate=2)
# 52,52,96 -- 52,52,160
x = inv_res_block(x, filters=160, strides=1, expansion=6, is_add=False, rate=2)
x = inv_res_block(x, filters=160, strides=1, expansion=6, is_add=True, rate=4)
x = inv_res_block(x, filters=160, strides=1, expansion=6, is_add=True, rate=4)
# 52,52,160 -- 52,52,320
x = inv_res_block(x, filters=320, strides=1, expansion=6, is_add=False, rate=4)
# 以下为MoilenetV2其余部分,这里用不到
o = x
# 52,52,320 -- 52,52,1280
o = Conv2D(1280,kernel_size=1,use_bias=False)(o)
o = BatchNormalization()(o)
o = ReLU(max_value=6)(o)
# 52,52,1280 -- 1280
o = GlobalAveragePooling2D()(o)
n_classes = 20
# 1280 -- n_classes
o = Dense(n_classes,activation='softmax')(o)
return img_input,residual_1,x
三.解码并创建模型
from tensorflow.keras.layers import *
from encoders import encoder_MobilenetV1_1,encoder_MobilenetV1_2
from tensorflow.keras.models import Model
import tensorflow
def pool_block(inputs,pool_num):
# 设 inputs shape 为 (18,18,1024)
h = inputs.shape[1]
w = inputs.shape[2]
# 设置池化的边长和步长 (18/1,18/2,18/3,18/6) - (18,9,6,3)
pool_size = strides = (int(h/pool_num),int(w/pool_num))
# 网格化池化,得到的特征图为 (1,1,1024),(2,2,1024),(3,3,1024),(6,6,1024)
x = AveragePooling2D(pool_size,strides,padding='same')(inputs)
# 调整通道数
x = Conv2D(512,1,padding='same')(x)
x = BatchNormalization()(x)
x = ReLU()(x)
# 将x大小变为inputs同样大小 (18*1 9*2 6*3 3*6)
x = UpSampling2D(strides)(x)
return x
def build_pspnet(n_classes,input_height=576,input_width=576,encoder_type='MobilenetV1_1'):
# 边长不能整除192时报错
if input_height%(32*6) != 0 or input_width%(32*6) != 0:
raise RuntimeError('Picture size cannot be divided by 32 * 6')
# 1.获取encoder的输出 (576,576,3--18,18,1024)
if encoder_type == 'MobilenetV1_1':
img_input, [out1,out2,out3,out4,out5] = encoder_MobilenetV1_1(input_height, input_width)
elif encoder_type == 'MobilenetV1_2':
img_input, [out1,out2,out3,out4,out5] = encoder_MobilenetV1_2(input_height, input_width)
else:
raise RuntimeError('pspnet encoder name is error')
# out5 shape 18,18,1024
# 2.PSP获取最终特征
# 对 out5 进行不同边长的网格池化
pool_nums = [1,2,3,6]
pool_outs = [out5]
# 获取池化后并resize后的特征
for pool_num in pool_nums:
p = pool_block(out5,pool_num)
pool_outs.append(p)
# 将 pool_outs 中的特征堆叠合并 (一个(18,18,1024),四个(18,18,512))
# 18, 18, 1024 -- 18,18,3072
x = Concatenate()(pool_outs)
# 18, 18, 1024 -- 18,18,512
x = Conv2D(512, 1, use_bias=False)(x)
x = BatchNormalization()(x)
x = ReLU()(x)
# 18, 18, 512 -- 18*8,18*8,n_classes
x = Conv2D(n_classes, 1, use_bias=False)(x)
x = UpSampling2D((8, 8))(x)
# x = tensorflow.image.resize(x,(x.shape[1]*8,x.shape[2]*8))
# 18*8,18*8,n_classes -- 144*144,n_classes
x = Reshape((-1,n_classes))(x)
x = Softmax()(x)
return Model(img_input,x)
四.Pspnet训练斑马线语义分割
from pspnet import build_pspnet
from tensorflow.keras.callbacks import ModelCheckpoint,ReduceLROnPlateau,EarlyStopping
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import CategoricalCrossentropy
import numpy as np
from PIL import Image
import os
import argparse
def parse_opt():
parse = argparse.ArgumentParser()
parse.add_argument('--datasets_path',type=str,default='../../datasets/banmaxian',help='数据集路径')
parse.add_argument('--n_classes',type=int,default=2,help='标签种类(含背景)')
parse.add_argument('--height',type=int,default=576,help='图片高度')
parse.add_argument('--width',type=int,default=576,help='图片宽度')
parse.add_argument('--batch_size',type=int,default=4)
parse.add_argument('--lr',type=float,default=0.0001)
parse.add_argument('--epochs',type=int,default=50)
parse.add_argument('--encoder_type',type=str,default='MobilenetV1_2',help='pspnet模型编码器的类型[MobilenetV1_1,MobilenetV1_2]')
opt = parse.parse_args()
return opt
def get_data_from_file(opt):
datasets_path,height,width,n_classes = opt.datasets_path,opt.height,opt.width,opt.n_classes
with open(os.path.join(datasets_path,'train.txt')) as f:
lines = f.readlines()
lines = [line.replace('\n','') for line in lines]
X = []
Y = []
for i in range(len(lines)):
names = lines[i].split(';')
real_name = names[0] # xx.jpg
label_name = names[1] # xx.png
# 读取真实图像
real_img = Image.open(os.path.join(datasets_path,'jpg',real_name))
real_img = real_img.resize((height,width))
real_img = np.array(real_img)/255 # (576,576,3) [0,1]
X.append(real_img)
# 读取标签图像,3通道,每个通道的数据都一样,每个像素点就是对应的类别,0表示背景
label_img = Image.open(os.path.join(datasets_path, 'png', label_name))
label_img = label_img.resize((int(height/4), int(width/4)))
label_img = np.array(label_img) # (144,144,3) [0,1]
# 根据标签图像来创建训练标签数据,n类对应的 seg_labels 就有n个通道
# 此时 seg_labels 每个通道的都值为 0
seg_labels = np.zeros((int(height/4), int(width/4),n_classes)) # (144,144,2)
# 第0通道表示第0类
# 第1通道表示第1类
# .....
# 第n_classes通道表示第n_classes类
for c in range(n_classes):
seg_labels[:,:,c] = (label_img[:,:,0]==c).astype(int)
# 此时 seg_labels 每个通道的值为0或1, 1 表示该像素点是该类,0 则不是
seg_labels = np.reshape(seg_labels,(-1,n_classes)) # (144*144,2)
Y.append(seg_labels)
return np.array(X),np.array(Y)
if __name__ == '__main__':
# 1.参数初始化
opt = parse_opt()
# 2.获取数据集
X,Y = get_data_from_file(opt)
# 3.创建模型
# 每5个epoch保存一次
weight_path = 'weights/pspnet_' + opt.encoder_type+'_weight/'
model = build_pspnet(opt.n_classes,opt.height,opt.width,opt.encoder_type,)
os.makedirs(weight_path,exist_ok=True)
checkpoint = ModelCheckpoint(
filepath=weight_path+'acc{accuracy:.4f}-ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
monitor='val_loss',
verbose=1,save_best_only=True,save_weights_only=True,period=5
)
lr_sh = ReduceLROnPlateau(monitor='val_loss',factor=0.5,patience=5,verbose=1)
es = EarlyStopping(monitor='val_loss',patience=10,verbose=1)
model.compile(loss=CategoricalCrossentropy(),optimizer=Adam(opt.lr),metrics='accuracy')
# 4.模型训练
model.fit(
x=X,y=Y,
batch_size=opt.batch_size,
epochs=opt.epochs,
callbacks=[checkpoint,lr_sh],
verbose=1,
validation_split=0.3,
shuffle=True,
)
# 5.模型保存
model.save_weights(weight_path+'/last.h5')
五.测试
from pspnet import build_pspnet
from PIL import Image
import numpy as np
import copy
import os
import argparse
def parse_opt():
parse = argparse.ArgumentParser()
parse.add_argument('--test_imgs', type=str, default='test_imgs', help='测试数据集')
parse.add_argument('--test_out', type=str, default='test_res', help='测试数据集')
parse.add_argument('--n_classes', type=int, default=2, help='标签种类(含背景)')
parse.add_argument('--height', type=int, default=576, help='输入模型的图片高度')
parse.add_argument('--width', type=int, default=576, help='输入模型的图片宽度')
parse.add_argument('--encoder_type', type=str, default='MobilenetV1_2', help='pspnet模型编码器的类型[MobilenetV1_1,MobilenetV1_2]')
opt = parse.parse_args()
return opt
def resize_img(path,real_width,real_height):
img_names = os.listdir(path)
for img_name in img_names:
img = Image.open(os.path.join(path, img_name))
img = img.resize((real_width,real_height))
img.save(os.path.join(path, img_name))
if __name__ == '__main__':
# 1.参数初始化
opt = parse_opt()
# class_colors 要根据图像的语义标签来设定;n_classes 行 3 列;
# 3列为RGB的值
class_colors = [[0, 0, 0],
[0, 255, 0]]
imgs_path = os.listdir(opt.test_imgs)
imgs_test = []
imgs_init = []
jpg_names = []
real_width,real_height = 1280,720
resize_img(opt.test_imgs, real_width,real_height)
# 2.获取测试图片
for i,jpg_name in enumerate(imgs_path):
img_init = Image.open(os.path.join(opt.test_imgs, jpg_name))
img = copy.deepcopy(img_init)
img = img.resize((opt.width,opt.height))
img = np.array(img) / 255 # (576,576,3) [0,1]
imgs_test.append(img)
imgs_init.append(img_init)
jpg_names.append(jpg_name)
imgs_test = np.array(imgs_test) # (-1,576,576,3)
# 3.模型创建
weight_path = 'weights/pspnet_' + opt.encoder_type + '_weight/'
model = build_pspnet(opt.n_classes,opt.height,opt.width, opt.encoder_type)
model.load_weights(os.path.join(weight_path, 'last.h5'))
# 4.模型预测语义分类结果
prs = model.predict(imgs_test) # (-1, 20736, 2)
# 结果 reshape
prs = prs.reshape(-1, int(opt.height / 4), int(opt.width / 4), opt.n_classes) # (-1, 144, 144, 2)
# 找到结果每个像素点所属类别的索引 两类就是 0 或 1
prs = prs.argmax(axis=-1) # (-1, 144, 144)
# 此时 prs 就是预测出来的类别,argmax 求得是最大值所在的索引,这个索引和类别值相同
# 所以 prs 每个像素点就是对应的类别
# 5.创建语义图像
# 和训练集中的语义标签图像不同,这里要显示图像,所以固定3通道
imgs_seg = np.zeros((len(prs), int(opt.height / 4), int(opt.width / 4), 3)) # (-1,144,144,3)
for c in range(opt.n_classes):
# 每个通道都要判断是否属于第0,1,2... n-1 类,是的话就乘以对应的颜色,每个类别都要判断一次
# 因为是RGB三个通道,所以3个通道分别乘以class_colors的每个通道颜色值
imgs_seg[:,:,:,0] += ((prs[:,:,:]==c)*(class_colors[c][0])).astype(int)
imgs_seg[:,:,:,1] += ((prs[:,:,:]==c)*(class_colors[c][1])).astype(int)
imgs_seg[:,:,:,2] += ((prs[:,:,:]==c)*(class_colors[c][2])).astype(int)
# 6.保存结果
save_path = opt.test_out+'/'+opt.encoder_type
os.makedirs(save_path,exist_ok=True)
for img_init,img_seg,img_name in zip(imgs_init,imgs_seg,jpg_names):
img_seg = Image.fromarray(np.uint8(img_seg)).resize((real_width,real_height))
images = Image.blend(img_init,img_seg,0.3)
images.save(os.path.join(opt.test_out+'/'+opt.encoder_type,img_name))