这是重新构建了的Unet语义分割网络,主要是文件框架上的构建,还有代码的实现,和之前的语义分割网络相比,更加完整也更清晰一些。建议还是学习这个版本的Unet。
还是快乐的tensorflow人。
Unet是一个优秀的语义分割模型,其主要执行过程与其它语义分割模型类似。
Unet可以分为三个部分,如下图所示:
第一部分是主干特征提取部分,我们可以利用主干部分获得一个又一个的特征层,Unet的主干特征提取部分与VGG相似,为卷积和最大池化的堆叠。利用主干特征提取部分我们可以获得五个初步有效特征层,在第二步中,我们会利用这五个有效特征层可以进行特征融合。
第二部分是加强特征提取部分,我们可以利用主干部分获取到的五个初步有效特征层进行上采样,并且进行特征融合,获得一个最终的,融合了所有特征的有效特征层。
第三部分是预测部分,我们会利用最终获得的最后一个有效特征层对每一个特征点进行分类,相当于对每一个像素点进行分类。
Github源码下载地址为:
https://github.com/bubbliiiing/unet-tf2
Unet的主干特征提取部分由卷积+最大池化组成,整体结构与VGG类似。
本文所采用的主干特征提取网络为VGG16,这样也方便使用imagnet上的预训练权重。
VGG是由Simonyan 和Zisserman在文献《Very Deep Convolutional Networks for Large Scale Image Recognition》中提出卷积神经网络模型,其名称来源于作者所在的牛津大学视觉几何组(Visual Geometry Group)的缩写。
该模型参加2014年的 ImageNet图像分类与定位挑战赛,取得了优异成绩:在分类任务上排名第二,在定位任务上排名第一。
它的结构如下图所示:
这是一个VGG16被用到烂的图,但确实很好的反应了VGG16的结构。
当我们使用VGG16作为主干特征提取网络的时候,我们只会用到两种类型的层,分别是卷积层和最大池化层。
当输入的图像大小为512x512x3的时候,具体执行方式如下:
1、conv1:进行两次[3,3]的64通道的卷积,获得一个[512,512,64]的初步有效特征层,再进行2X2最大池化,获得一个[256,256,64]的特征层。
2、conv2:进行两次[3,3]的128通道的卷积,获得一个[256,256,128]的初步有效特征层,再进行2X2最大池化,获得一个[128,128,128]的特征层。
3、conv3:进行三次[3,3]的256通道的卷积,获得一个[128,128,256]的初步有效特征层,再进行2X2最大池化,获得一个[64,64,256]的特征层。
4、conv4:进行三次[3,3]的512通道的卷积,获得一个[64,64,512]的初步有效特征层,再进行2X2最大池化,获得一个[32,32,512]的特征层。
5、conv5:进行三次[3,3]的512通道的卷积,获得一个[32,32,512]的初步有效特征层。
from tensorflow.keras import layers
from tensorflow.keras.initializers import RandomNormal
def VGG16(img_input):
# Block 1
# 512,512,3 -> 512,512,64
x = layers.Conv2D(64, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block1_conv1')(img_input)
x = layers.Conv2D(64, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block1_conv2')(x)
feat1 = x
# 512,512,64 -> 256,256,64
x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block1_pool')(x)
# Block 2
# 256,256,64 -> 256,256,128
x = layers.Conv2D(128, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block2_conv1')(x)
x = layers.Conv2D(128, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block2_conv2')(x)
feat2 = x
# 256,256,128 -> 128,128,128
x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block2_pool')(x)
# Block 3
# 128,128,128 -> 128,128,256
x = layers.Conv2D(256, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block3_conv1')(x)
x = layers.Conv2D(256, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block3_conv2')(x)
x = layers.Conv2D(256, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block3_conv3')(x)
feat3 = x
# 128,128,256 -> 64,64,256
x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block3_pool')(x)
# Block 4
# 64,64,256 -> 64,64,512
x = layers.Conv2D(512, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block4_conv1')(x)
x = layers.Conv2D(512, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block4_conv2')(x)
x = layers.Conv2D(512, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block4_conv3')(x)
feat4 = x
# 64,64,512 -> 32,32,512
x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block4_pool')(x)
# Block 5
# 32,32,512 -> 32,32,512
x = layers.Conv2D(512, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block5_conv1')(x)
x = layers.Conv2D(512, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block5_conv2')(x)
x = layers.Conv2D(512, (3, 3),
activation='relu',
padding='same',
kernel_initializer = RandomNormal(stddev=0.02),
name='block5_conv3')(x)
feat5 = x
return feat1, feat2, feat3, feat4, feat5
Unet所使用的加强特征提取网络是一个U的形状。
利用第一步我们可以获得五个初步的有效特征层,在加强特征提取网络这里,我们会利用这五个初步的有效特征层进行特征融合,特征融合的方式就是对特征层进行上采样并且进行堆叠。
为了方便网络的构建与更好的通用性,我们的Unet和上图的Unet结构有些许不同,在上采样时直接进行两倍上采样再进行特征融合,最终获得的特征层和输入图片的高宽相同。
具体示意图如下:
from tensorflow.keras.initializers import RandomNormal
from tensorflow.keras.layers import *
from tensorflow.keras.models import *
from nets.vgg16 import VGG16
def Unet(input_shape=(256,256,3), num_classes=21):
inputs = Input(input_shape)
#-------------------------------#
# 获得五个有效特征层
# feat1 512,512,64
# feat2 256,256,128
# feat3 128,128,256
# feat4 64,64,512
# feat5 32,32,512
#-------------------------------#
feat1, feat2, feat3, feat4, feat5 = VGG16(inputs)
channels = [64, 128, 256, 512]
# 32, 32, 512 -> 64, 64, 512
P5_up = UpSampling2D(size=(2, 2))(feat5)
# 64, 64, 512 + 64, 64, 512 -> 64, 64, 1024
P4 = Concatenate(axis=3)([feat4, P5_up])
# 64, 64, 1024 -> 64, 64, 512
P4 = Conv2D(channels[3], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P4)
P4 = Conv2D(channels[3], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P4)
# 64, 64, 512 -> 128, 128, 512
P4_up = UpSampling2D(size=(2, 2))(P4)
# 128, 128, 256 + 128, 128, 512 -> 128, 128, 768
P3 = Concatenate(axis=3)([feat3, P4_up])
# 128, 128, 768 -> 128, 128, 256
P3 = Conv2D(channels[2], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P3)
P3 = Conv2D(channels[2], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P3)
# 128, 128, 256 -> 256, 256, 256
P3_up = UpSampling2D(size=(2, 2))(P3)
# 256, 256, 256 + 256, 256, 128 -> 256, 256, 384
P2 = Concatenate(axis=3)([feat2, P3_up])
# 256, 256, 384 -> 256, 256, 128
P2 = Conv2D(channels[1], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P2)
P2 = Conv2D(channels[1], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P2)
# 256, 256, 128 -> 512, 512, 128
P2_up = UpSampling2D(size=(2, 2))(P2)
# 512, 512, 128 + 512, 512, 64 -> 512, 512, 192
P1 = Concatenate(axis=3)([feat1, P2_up])
# 512, 512, 192 -> 512, 512, 64
P1 = Conv2D(channels[0], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P1)
P1 = Conv2D(channels[0], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P1)
# 512, 512, 64 -> 512, 512, num_classes
P1 = Conv2D(num_classes, 1, activation="softmax")(P1)
model = Model(inputs=inputs, outputs=P1)
return model
利用1、2步,我们可以获取输入进来的图片的特征,此时,我们需要利用特征获得预测结果。
利用特征获得预测结果的过程为:
利用一个1x1卷积进行通道调整,将最终特征层的通道数调整成num_classes。
from tensorflow.keras.initializers import RandomNormal
from tensorflow.keras.layers import *
from tensorflow.keras.models import *
from nets.vgg16 import VGG16
def Unet(input_shape=(256,256,3), num_classes=21):
inputs = Input(input_shape)
#-------------------------------#
# 获得五个有效特征层
# feat1 512,512,64
# feat2 256,256,128
# feat3 128,128,256
# feat4 64,64,512
# feat5 32,32,512
#-------------------------------#
feat1, feat2, feat3, feat4, feat5 = VGG16(inputs)
channels = [64, 128, 256, 512]
# 32, 32, 512 -> 64, 64, 512
P5_up = UpSampling2D(size=(2, 2))(feat5)
# 64, 64, 512 + 64, 64, 512 -> 64, 64, 1024
P4 = Concatenate(axis=3)([feat4, P5_up])
# 64, 64, 1024 -> 64, 64, 512
P4 = Conv2D(channels[3], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P4)
P4 = Conv2D(channels[3], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P4)
# 64, 64, 512 -> 128, 128, 512
P4_up = UpSampling2D(size=(2, 2))(P4)
# 128, 128, 256 + 128, 128, 512 -> 128, 128, 768
P3 = Concatenate(axis=3)([feat3, P4_up])
# 128, 128, 768 -> 128, 128, 256
P3 = Conv2D(channels[2], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P3)
P3 = Conv2D(channels[2], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P3)
# 128, 128, 256 -> 256, 256, 256
P3_up = UpSampling2D(size=(2, 2))(P3)
# 256, 256, 256 + 256, 256, 128 -> 256, 256, 384
P2 = Concatenate(axis=3)([feat2, P3_up])
# 256, 256, 384 -> 256, 256, 128
P2 = Conv2D(channels[1], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P2)
P2 = Conv2D(channels[1], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P2)
# 256, 256, 128 -> 512, 512, 128
P2_up = UpSampling2D(size=(2, 2))(P2)
# 512, 512, 128 + 512, 512, 64 -> 512, 512, 192
P1 = Concatenate(axis=3)([feat1, P2_up])
# 512, 512, 192 -> 512, 512, 64
P1 = Conv2D(channels[0], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P1)
P1 = Conv2D(channels[0], 3, activation='relu', padding='same', kernel_initializer = RandomNormal(stddev=0.02))(P1)
# 512, 512, 64 -> 512, 512, num_classes
P1 = Conv2D(num_classes, 1, activation="softmax")(P1)
model = Model(inputs=inputs, outputs=P1)
return model
我们使用的训练文件采用VOC的格式。
语义分割模型训练的文件分为两部分。
第一部分是原图,像这样:
第二部分标签,像这样:
原图就是普通的RGB图像,标签就是灰度图或者8位彩色图。
原图的shape为[height, width, 3],标签的shape就是[height, width],对于标签而言,每个像素点的内容是一个数字,比如0、1、2、3、4、5……,代表这个像素点所属的类别。
语义分割的工作就是对原始的图片的每一个像素点进行分类,所以通过预测结果中每个像素点属于每个类别的概率与标签对比,可以对网络进行训练。
本文所使用的LOSS由两部分组成:
1、Cross Entropy Loss。
2、Dice Loss。
Cross Entropy Loss就是普通的交叉熵损失,当语义分割平台利用Softmax对像素点进行分类的时候,进行使用。
Dice loss将语义分割的评价指标作为Loss,Dice系数是一种集合相似度度量函数,通常用于计算两个样本的相似度,取值范围在[0,1]。
计算公式如下:
就是预测结果和真实结果的交乘上2,除上预测结果加上真实结果。其值在0-1之间。越大表示预测结果和真实结果重合度越大。所以Dice系数是越大越好。
如果作为LOSS的话是越小越好,所以使得Dice loss = 1 - Dice,就可以将Loss作为语义分割的损失了。
实现代码如下:
def dice_loss_with_CE(beta=1, smooth = 1e-5):
def _dice_loss_with_CE(y_true, y_pred):
y_pred = K.clip(y_pred, K.epsilon(), 1.0 - K.epsilon())
CE_loss = - y_true[...,:-1] * K.log(y_pred)
CE_loss = K.mean(K.sum(CE_loss, axis = -1))
tp = K.sum(y_true[...,:-1] * y_pred, axis=[0,1,2])
fp = K.sum(y_pred , axis=[0,1,2]) - tp
fn = K.sum(y_true[...,:-1], axis=[0,1,2]) - tp
score = ((1 + beta ** 2) * tp + smooth) / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + smooth)
score = tf.reduce_mean(score)
dice_loss = 1 - score
# dice_loss = tf.Print(dice_loss, [dice_loss, CE_loss])
return CE_loss + dice_loss
return _dice_loss_with_CE
def CE():
def _CE(y_true, y_pred):
y_pred = K.clip(y_pred, K.epsilon(), 1.0 - K.epsilon())
CE_loss = - y_true[...,:-1] * K.log(y_pred)
CE_loss = K.mean(K.sum(CE_loss, axis = -1))
# dice_loss = tf.Print(CE_loss, [CE_loss])
return CE_loss
return _CE
首先前往Github下载对应的仓库,下载完后利用解压软件解压,之后用编程软件打开文件夹。
注意打开的根目录必须正确,否则相对目录不正确的情况下,代码将无法运行。
一定要注意打开后的根目录是文件存放的目录。
本文使用VOC格式进行训练,训练前需要自己制作好数据集,如果没有自己的数据集,可以通过Github连接下载VOC12+07的数据集尝试下。
训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。
训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的SegmentationClass中。
在完成数据集的摆放之后,我们需要对数据集进行下一步的处理,目的是获得训练用的train.txt以及val.txt,需要用到根目录下的voc_annotation.py。
如果下载的是我上传的voc数据集,那么就不需要运行根目录下的voc_annotation.py。
如果是自己制作的数据集,那么需要运行根目录下的voc_annotation.py,从而生成train.txt和val.txt。
通过voc_annotation.py我们已经生成了train.txt以及val.txt,此时我们可以开始训练了。训练的参数较多,大家可以在下载库后仔细看注释,其中最重要的部分依然是train.py里的num_classes。
num_classes用于指向检测类别的个数+1!训练自己的数据集必须要修改!
训练结果预测需要用到两个文件,分别是unet.py和predict.py。
我们首先需要去unet.py里面修改model_path以及num_classes,这两个参数必须要修改。
model_path指向训练好的权值文件,在logs文件夹里。
num_classes指向检测类别的个数+1。
完成修改后就可以运行predict.py进行检测了。运行后输入图片路径即可检测。