def non_local_block(input_tensor, computation_compression=2, is_training=True):
input_shape = input_tensor.get_shape().as_list()
print('input_shape', input_shape)
batchsize, dim1, dim2, channels = input_shape
print(input_shape, '-------------------')
theta = cb.conv2d(input_tensor, channels, 1, name='theta')
theta = cb.layerbn(theta, is_training=is_training, name='bn1')
theta = cb.relu(theta, name='relu1')
phi = cb.conv2d(input_tensor, channels, 1, name='phi')
phi = cb.layerbn(phi, is_training=is_training, name='bn2')
phi = cb.relu(phi, name='relu2')
# phi = tf.reshape(phi,shape=[batchsize, dim2, dim1, channels])
f = tf.matmul(theta, phi, transpose_b=True)
f = tf.nn.softmax(f)
g = cb.conv2d(input_tensor, channels, 1, name='g')
g = cb.layerbn(g, is_training=is_training, name='bn3')
g = cb.relu(g, name='relu3')
y = tf.matmul(f, g)
print('y=', y)
y = cb.conv2d(y, channels, kernel_size=3, name='y')
print('y=', y)
residual = input_tensor + y
return residual
def NA(x, is_training=True):
input_shape = x.get_shape().as_list()
print('input_shape', input_shape)
batchsize, dim1, dim2, channels = input_shape
print(input_shape, '-------------------')
y = tf.nn.depthwise_conv2d(x,
filter=[3, 3, channels, channels],
strides=1,
padding='SAME',
name='dp_conv')
y = cb.layerbn(y, is_training=is_training, name='bn1')
y = cb.relu(y, name='relu1')
y = cb.conv2d(y, channels, 1, name='conv')
y = tf.nn.sigmoid(y)
return x * y #tf.multiply(x, y)
def SmoothDilatedResidualBlock(x, dilation=1, index=1, reuse=False):
with tf.variable_scope("SmoothDilatedResidualBlock_%d" % index,
reuse=reuse):
y = dilated._smoothed_dilated_conv2d_SSC(x,
3,
64,
dilation * 2 - 1,
name='dilated_conv_1')
y = cb.instancenorm(y, name='norm_1')
y1 = cb.lrelu(y, name='relu_1')
y = dilated._smoothed_dilated_conv2d_SSC(y1,
3,
64,
dilation * 2 - 1,
name='dilated_conv_2')
y2 = cb.instancenorm(y, name='norm_2')
y = cb.lrelu(y1 + y2, name='relu_2')
return y
def ResidualBlock(x, dilation=1, reuse=False):
with tf.variable_scope("ResidualBlock", reuse=reuse):
y = cb.dilation_conv(x, 3, 64, dilation, name='diconv_1')
y = cb.instancenorm(y, name='norm_1')
y1 = cb.lrelu(y, name='relu_1')
y = cb.dilation_conv(x, 3, 64, dilation, name='diconv_2')
y2 = cb.instancenorm(y, name='norm_2')
y = cb.lrelu(y1 + y2, name='relu_2')
return y
def GCANet(x, out_c=3, only_residual=False, reuse=False):
with tf.variable_scope("GCANet", reuse=reuse):
y = cb.conv2d(x,
64,
kernel_size=3,
padding='SAME',
stride=1,
use_bias=False,
name='conv_1')
y = cb.instancenorm(y, name='norm_1')
y = cb.lrelu(y, name='relu_1')
y = cb.conv2d(y,
64,
kernel_size=3,
padding='SAME',
stride=1,
use_bias=False,
name='conv_2')
y = cb.instancenorm(y, name='norm_2')
y = cb.lrelu(y, name='relu_2')
y = cb.conv2d(y,
64,
kernel_size=3,
padding='SAME',
stride=2,
use_bias=False,
name='conv_3')
y = cb.instancenorm(y, name='norm_3')
y1 = cb.lrelu(y, name='relu_3')
y = SmoothDilatedResidualBlock(y1, 2, 1)
y = SmoothDilatedResidualBlock(y, 2, 2)
y = SmoothDilatedResidualBlock(y, 2, 3)
y2 = SmoothDilatedResidualBlock(y, 4, 4)
y = SmoothDilatedResidualBlock(y2, 4, 5)
y = SmoothDilatedResidualBlock(y, 4, 6)
y3 = ResidualBlock(y, 1)
gates = tf.concat([y1, y2, y3], axis=-1)
gates = cb.conv2d(gates,
3,
kernel_size=3,
padding='SAME',
stride=1,
use_bias=True,
name='conv_4')
gated_y = y1 * gates[:, :, :,
0:-2] + y2 * gates[:, :, :,
1:-1] + y3 * gates[:, :, :, 2:]
y = cb.deconv2d(gated_y,
out_channel=64,
kernel_size=4,
padding='SAME',
stride=2,
name='deconv1')
y = cb.instancenorm(y, name='norm_4')
y = cb.lrelu(y, name='relu_4')
y = cb.deconv2d(y,
out_channel=64,
kernel_size=3,
padding='SAME',
stride=1,
name='deconv2')
y = cb.instancenorm(y, name='norm_5')
y = cb.lrelu(y, name='relu_5')
if only_residual:
y = cb.deconv2d(y,
out_channel=out_c,
kernel_size=1,
padding='SAME',
stride=1,
name='deconv3')
else:
y = cb.deconv2d(y,
out_channel=out_c,
kernel_size=1,
padding='SAME',
stride=1,
name='deconv3')
y = cb.lrelu(y, name='relu_6')
return y
def recursive_block5(rain_img, input_, res1, res2, res3, res4, c_dim, output_channel=16, recursive_num=3, index=5,
stride=1, is_train=True):
with tf.variable_scope("rec_block_%d" % index):
print('block5')
conv1_1 = cb.conv2d(input_, output_channel, kernel_size=3, padding='SAME', stride=1, use_bias=True,
name='conv1_1')
conv1_1 = cb.lrelu(conv1_1, name='relu_1')
dia_conv1 = cb.dilation_conv(input_tensor=conv1_1, k_size=3, out_dims=output_channel, rate=2,
padding='SAME', use_bias=False, name='dia_conv_1')
relu_7 = cb.lrelu(dia_conv1, name='relu_7')
dia_conv2 = cb.dilation_conv(input_tensor=relu_7, k_size=3, out_dims=output_channel, rate=4,
padding='SAME', use_bias=False, name='dia_conv_2')
relu_8 = cb.lrelu(dia_conv2, name='relu_8')
dia_conv3 = cb.dilation_conv(input_tensor=relu_8, k_size=3, out_dims=output_channel, rate=8,
padding='SAME', use_bias=False, name='dia_conv_3')
relu_9 = cb.lrelu(dia_conv3, name='relu_9')
dia_conv4 = cb.dilation_conv(input_tensor=relu_9, k_size=3, out_dims=output_channel, rate=16,
padding='SAME', use_bias=False, name='dia_conv_4')
relu_10 = cb.lrelu(dia_conv4, name='relu_10')
print('non_local_block start', conv1_1.shape)
conv1_1 = non_local_block(relu_10)
# print('non_local_block End', conv1_1.shape)
conv_temp = tf.concat([conv1_1, res1, res2, res3, res4], 3)
for i in range(recursive_num):
conv_temp = cb.conv2d(conv_temp, output_channel, kernel_size=3, padding='SAME', stride=1, use_bias=True,
name='conv_temp_1_%d'%(i+1))
conv_temp = cb.lrelu(conv_temp, name='relu_2_%d'%(i+1))
conv_temp = tf.concat([conv_temp, res1, res2, res3, res4], 3)
conv_temp = cb.conv2d(conv_temp, output_channel, kernel_size=3, padding='SAME', stride=1, use_bias=True,
name='conv_temp_2_%d'%(i+1))
conv_temp = cb.lrelu(conv_temp, name='relu_3_%d'%(i+1))
conv_temp = tf.concat([conv_temp, res1, res2, res3, res4], 3)
conv1_2 = conv_temp
res = cb.conv2d(conv1_2, c_dim, kernel_size=3, padding='SAME', stride=1, use_bias=True, name='res')
res = cb.lrelu(res, name='relu_4')
return rain_img + res, res