TensorFlow2.x构建U^2Net网络
文章目录
- 前言
- 一、环境搭建
- 二、网络结构
-
- 1.网络结构图
- 2.导入头文件
- 3.RSU结构
- 4.扩展模块
- 5.整体结构
前言
U^2Net是我之前使用过的一个图像分割的网络，由于效果比较好，所以对其印象比较深刻，同样为了学习，当时用TensorFlow2.x重构了这个网络，虽说现在发有点晚，但为了防止自己忘记，还是把它记录了下来。相关文件链接如下：
论文地址
官方代码（pytorch）
详细解读
一、环境搭建
pip install tensorflow-gpu
二、网络结构
1.网络结构图
网络结构图如下：
2.导入头文件
import tensorflow.keras as k
import tensorflow as tf
from tensorflow.keras.models import Sequential,Model
from tensorflow.keras.layers import Conv2D,MaxPool2D,BatchNormalization,ReLU,UpSampling2D
3.RSU结构
RSU解耦股相当于一个小型的unet结构，前半部分通过一系列下采样来提取特征，后半部分通过上采样和concat的方式达到特征融合。需要注意的是：这里的连续4个RSU模块其实可以写的更简化，这里我就直接全部展开来写了，代码如下：
#基本卷积块
class REBNCONV(Model):
def __init__(self,out_ch=3,dirate=1):
super(REBNCONV, self).__init__()
self.conv=Sequential()
self.conv.add(Conv2D(out_ch,kernel_size=(3,3),strides=(1,1),padding="same",dilation_rate=dirate))
self.conv.add(BatchNormalization())
self.conv.add(ReLU())
def call(self, inputs, training=None, mask=None):
x=self.conv(inputs)
return x
#第一个RSU结构
class RSU7(Model):
def __init__(self,mid_ch=12,out_ch=3):
super(RSU7, self).__init__()
self.rebnconv0=REBNCONV(out_ch,dirate=1)
self.rebnconv1=REBNCONV(mid_ch,dirate=1)
self.pool1=MaxPool2D(pool_size=(2,2),strides=(2,2))
self.rebnconv2=REBNCONV(mid_ch,dirate=1)
self.pool2=MaxPool2D(pool_size=(2,2),strides=(2,2))
self.rebnconv3=REBNCONV(mid_ch,dirate=1)
self.pool3=MaxPool2D(pool_size=(2,2),strides=(2,2))
self.rebnconv4=REBNCONV(mid_ch,dirate=1)
self.pool4=MaxPool2D(pool_size=(2,2),strides=(2,2))
self.rebnconv5=REBNCONV(mid_ch,dirate=1)
self.pool5=MaxPool2D(pool_size=(2,2),strides=(2,2))
self.rebnconv6=REBNCONV(mid_ch,dirate=1)
self.rebnconv7=REBNCONV(mid_ch,dirate=2)
self.rebnconv6d=REBNCONV(mid_ch,dirate=1)
self.rebnconv5d=REBNCONV(mid_ch,dirate=1)
self.rebnconv4d=REBNCONV(mid_ch,dirate=1)
self.rebnconv3d=REBNCONV(mid_ch,dirate=1)
self.rebnconv2d=REBNCONV(mid_ch,dirate=1)
self.rebnconv1d=REBNCONV(out_ch,dirate=1)
def call(self, inputs, training=None, mask=None):
h0=self.rebnconv0(inputs)
hx1=self.rebnconv1(h0)
hx=self.pool1(hx1)
hx2=self.rebnconv2(hx)
hx=self.pool2(hx2)
hx3=self.rebnconv3(hx)
hx=self.pool3(hx3)
hx4=self.rebnconv4(hx)
hx=self.pool4(hx4)
hx5=self.rebnconv5(hx)
hx=self.pool5(hx5)
hx6=self.rebnconv6(hx)
hx7=self.rebnconv7(hx6)
hx6d=self.rebnconv6d(tf.concat((hx7,hx6),axis=3))
hx6d_up=UpSampling2D((2,2),interpolation="bilinear")(hx6d) #上采样
hx5d=self.rebnconv5d(tf.concat((hx6d_up,hx5),axis=3))
hx5d_up=UpSampling2D((2,2),interpolation="bilinear")(hx5d)
hx4d=self.rebnconv4d(tf.concat((hx5d_up,hx4),axis=3))
hx4d_up=UpSampling2D((2,2),interpolation="bilinear")(hx4d)
hx3d=self.rebnconv3d(tf.concat((hx4d_up,hx3),axis=3))
hx3d_up=UpSampling2D((2,2),interpolation="bilinear")(hx3d)
hx2d=self.rebnconv2d(tf.concat((hx3d_up,hx2),axis=3))
hx2d_up=UpSampling2D((2,2),interpolation="bilinear")(hx2d)
hx1d=self.rebnconv1d(tf.concat((hx2d_up,hx1),axis=3))
return hx1d+h0
#第二个RSU模块
class RSU6(Model):
def __init__(self,mid_ch=12,out_ch=3):
super(RSU6, self).__init__()
self.rebnconv0=REBNCONV(out_ch,dirate=1)
self.rebnconv1=REBNCONV(mid_ch,dirate=1)
self.pool1=MaxPool2D(pool_size=(2,2),strides=(2,2))
self.rebnconv2=REBNCONV(mid_ch,dirate=1)
self.pool2=MaxPool2D(pool_size=(2,2),strides=(2,2))
self.rebnconv3=REBNCONV(mid_ch,dirate=1)
self.pool3=MaxPool2D(pool_size=(2,2),strides=(2,2))
self.rebnconv4=REBNCONV(mid_ch,dirate=1)
self.pool4=MaxPool2D(pool_size=(2,2),strides=(2,2))
self.rebnconv5=REBNCONV(mid_ch,dirate=1)
self.rebnconv6=REBNCONV(mid_ch,dirate=2)
self.rebnconv5d=REBNCONV(mid_ch,dirate=1)
self.rebnconv4d=REBNCONV(mid_ch,dirate=1)
self.rebnconv3d=REBNCONV(mid_ch,dirate=1)
self.rebnconv2d=REBNCONV(mid_ch,dirate=1)
self.rebnconv1d=REBNCONV(out_ch,dirate=1)
def call(self, inputs, training=None, mask=None):
hx0=self.rebnconv0(inputs)
hx1=self.rebnconv1(hx0)
hx=self.pool1(hx1)
hx2=self.rebnconv2(hx)
hx=self.pool2(hx2)
hx3=self.rebnconv3(hx)
hx=self.pool3(hx3)
hx4=self.rebnconv4(hx)
hx=self.pool4(hx4)
hx5=self.rebnconv5(hx)
hx6=self.rebnconv6(hx5)
hx5d=self.rebnconv5d(tf.concat((hx6,hx5),axis=3))
hx5d_up= UpSampling2D((2,2),interpolation="bilinear")(hx5d)
# print(hx5d_up.shape)
# print(hx4.shape)
hx4d=self.rebnconv4d(tf.concat((hx5d_up,hx4),axis=3))
hx4d_up=UpSampling2D((2,2),interpolation="bilinear")(hx4d)
hx3d=self.rebnconv3d(tf.concat((hx4d_up,hx3),axis=3))
hx3d_up=UpSampling2D((2,2),interpolation="bilinear")(hx3d)
hx2d=self.rebnconv2d(tf.concat((hx3d_up,hx2),axis=3))
hx2d_up=UpSampling2D((2,2),interpolation="bilinear")(hx2d)
hx1d=self.rebnconv1d(tf.concat((hx2d_up,hx1),axis=3))
return hx1d+hx0
#第三个RSU模块
class RSU5(Model):
def __init__(self, mid_ch=12, out_ch=3):
super(RSU5, self).__init__()
self.rebnconv0 = REBNCONV(out_ch, dirate=1)
self.rebnconv1 = REBNCONV(mid_ch, dirate=1)
self.pool1 = MaxPool2D(pool_size=(2, 2), strides=(2, 2))
self.rebnconv2 = REBNCONV(mid_ch, dirate=1)
self.pool2 = MaxPool2D(pool_size=(2, 2), strides=(2, 2))
self.rebnconv3 = REBNCONV(mid_ch, dirate=1)
self.pool3 = MaxPool2D(pool_size=(2, 2), strides=(2, 2))
self.rebnconv4 = REBNCONV(mid_ch, dirate=1)
self.rebnconv5 = REBNCONV(mid_ch, dirate=2)
self.rebnconv4d = REBNCONV(mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(out_ch, dirate=1)
def call(self, inputs, training=None, mask=None):
hx0 = self.rebnconv0(inputs)
hx1 = self.rebnconv1(hx0)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(tf.concat((hx5, hx4), axis=3))
hx4d_up = UpSampling2D((2, 2), interpolation="bilinear")(hx4d)
hx3d = self.rebnconv3d(tf.concat((hx4d_up, hx3), axis=3))
hx3d_up = UpSampling2D((2, 2), interpolation="bilinear")(hx3d)
hx2d = self.rebnconv2d(tf.concat((hx3d_up, hx2), axis=3))
hx2d_up = UpSampling2D((2, 2), interpolation="bilinear")(hx2d)
hx1d = self.rebnconv1d(tf.concat((hx2d_up, hx1), axis=3))
return hx1d + hx0
#第四个RSU模块
class RSU4(Model):
def __init__(self, mid_ch=12, out_ch=3):
super(RSU4, self).__init__()
self.rebnconv0 = REBNCONV(out_ch, dirate=1)
self.rebnconv1 = REBNCONV(mid_ch, dirate=1)
self.pool1 = MaxPool2D(pool_size=(2, 2), strides=(2, 2))
self.rebnconv2 = REBNCONV(mid_ch, dirate=1)
self.pool2 = MaxPool2D(pool_size=(2, 2), strides=(2, 2))
self.rebnconv3 = REBNCONV(mid_ch, dirate=1)
self.rebnconv4 = REBNCONV(mid_ch, dirate=2)
self.rebnconv3d = REBNCONV(mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(out_ch, dirate=1)
def call(self, inputs, training=None, mask=None):
hx0 = self.rebnconv0(inputs)
hx1 = self.rebnconv1(hx0)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx)
hx3d = self.rebnconv3d(tf.concat((hx4, hx3), axis=3))
hx3d_up = UpSampling2D((2, 2), interpolation="bilinear")(hx3d)
hx2d = self.rebnconv2d(tf.concat((hx3d_up, hx2), axis=3))
hx2d_up = UpSampling2D((2, 2), interpolation="bilinear")(hx2d)
hx1d = self.rebnconv1d(tf.concat((hx2d_up, hx1), axis=3))
return hx1d + hx0
4.扩展模块
在经过4个RSU模块后，特征图的分辨率变得很小（输入分辨率为320320的情况下，这里的分辨率为1818），如果继续下采样的话，会导致更多信息的丢失，所以作者在这里用扩展卷积替换掉了concat和上采样的操作，也就是说扩展模块中的所有中间层的特征图和输入进来的特征图有着相同的分辨率。代码如下：
#扩展模块
class RSU4F(Model):
def __init__(self,mid_ch=12,out_ch=3):
super(RSU4F, self).__init__()
self.rebnconv0=REBNCONV(out_ch,dirate=1)
self.rebnconv1=REBNCONV(mid_ch,dirate=1)
self.rebnconv2=REBNCONV(mid_ch,dirate=2)
self.rebnconv3=REBNCONV(mid_ch,dirate=4)
self.rebnconv4=REBNCONV(mid_ch,dirate=8)
self.rebnconv3d=REBNCONV(mid_ch,dirate=4)
self.rebnconv2d=REBNCONV(mid_ch,dirate=2)
self.rebnconv1d=REBNCONV(out_ch,dirate=1)
def call(self, inputs, training=None, mask=None):
hx0=self.rebnconv0(inputs)
hx1=self.rebnconv1(hx0)
hx2=self.rebnconv2(hx1)
hx3=self.rebnconv3(hx2)
hx4=self.rebnconv4(hx3)
hx3d=self.rebnconv3d(tf.concat((hx4,hx3),axis=3))
hx2d=self.rebnconv2d(tf.concat((hx3d,hx2),axis=3))
hx1d=self.rebnconv1d(tf.concat((hx2d,hx1),axis=3))
return hx1d+hx0
5.整体结构
这里需要注意的是，在得到每层的概率图时，所使用的激活函数均为sigmoid，整体的网络实现代码如下：
#U^2Net
class U2NET(Model):
def __init__(self,out_ch=1):
super(U2NET, self).__init__()
#encode
self.stage1=RSU7(32,64)
self.pool1_1=MaxPool2D(pool_size=(2,2),strides=(2,2))#144*144
self.stage2=RSU6(32,128)
self.pool2_1=MaxPool2D(pool_size=(2,2),strides=(2,2))#72*72
self.stage3=RSU5(64,256)
self.pool3_1=MaxPool2D(pool_size=(2,2),strides=(2,2))#36*36
self.stage4=RSU4(128,512)
self.pool4_1=MaxPool2D(pool_size=(2,2),strides=(2,2))#18*18
self.stage5=RSU4F(256,512)
self.pool5_1=MaxPool2D(pool_size=(2,2),strides=(2,2))#9*9
self.stage6=RSU4F(256,512)
#decode
self.stage5d=RSU4F(256,512)
self.stage4d=RSU4(128,256)
self.stage3d=RSU5(64,128)
self.stage2d=RSU6(32,64)
self.stage1d=RSU7(16,64)
#每个层的输出
self.side1=Conv2D(out_ch,kernel_size=(3,3),padding="same")
self.side2=Conv2D(out_ch,kernel_size=(3,3),padding="same")
self.side3=Conv2D(out_ch,kernel_size=(3,3),padding="same")
self.side4=Conv2D(out_ch,kernel_size=(3,3),padding="same")
self.side5=Conv2D(out_ch,kernel_size=(3,3),padding="same")
self.side6=Conv2D(out_ch,kernel_size=(3,3),padding="same")
#最终输出
self.outconv=Conv2D(out_ch,kernel_size=(1,1))
def call(self, inputs, training=None, mask=None):
hx1=self.stage1(inputs)
hx=self.pool1_1(hx1)
hx2=self.stage2(hx)
hx=self.pool2_1(hx2)
hx3=self.stage3(hx)
hx=self.pool3_1(hx3)
hx4=self.stage4(hx)
hx=self.pool4_1(hx4)
hx5=self.stage5(hx)
hx=self.pool5_1(hx5)
hx6=self.stage6(hx)
hx6_up=UpSampling2D((2,2),interpolation="bilinear")(hx6)
#decode
hx5d=self.stage5d(tf.concat((hx6_up,hx5),axis=3))
hx5d_up=UpSampling2D((2,2),interpolation="bilinear")(hx5d)
hx4d=self.stage4d(tf.concat((hx5d_up,hx4),axis=3))
hx4d_up=UpSampling2D((2,2),interpolation="bilinear")(hx4d)
hx3d=self.stage3d(tf.concat((hx4d_up,hx3),axis=3))
hx3d_up=UpSampling2D((2,2),interpolation="bilinear")(hx3d)
hx2d=self.stage2d(tf.concat((hx3d_up,hx2),axis=3))
hx2d_up=UpSampling2D((2,2),interpolation="bilinear")(hx2d)
hx1d=self.stage1d(tf.concat((hx2d_up,hx1),axis=3))
# side out
d1=self.side1(hx1d)
d2=self.side2(hx2d)
d2=UpSampling2D((2,2),interpolation="bilinear")(d2)
d3=self.side3(hx3d)
d3=UpSampling2D((4,4),interpolation="bilinear")(d3)
d4=self.side4(hx4d)
d4=UpSampling2D((8,8),interpolation="bilinear")(d4)
d5=self.side5(hx5d)
d5=UpSampling2D((16,16),interpolation="bilinear")(d5)
d6=self.side6(hx6)
d6=UpSampling2D((32,32),interpolation="bilinear")(d6)
out=self.outconv(tf.concat((d1,d2,d3,d4,d5,d6),axis=3))
sig=k.activations.sigmoid #定义激活函数
return sig(out),sig(d1),sig(d2),sig(d3),sig(d4),sig(d5),sig(d6)
以上便是本篇文章的全部内容，如需训练代码，或发现文章中有错误，欢迎在评论区留言。