inpainting里面重要的一个工作是Nvidia的Partial Convolution部分卷积,该卷积包含两个分支,一个是更新mask图像,另一个则是根据mask图进行部分的卷积。
下面是他的1.x版本的实现代码如下:
def pc_conv(x, mask, cnum, ksize, stride, test ,name='pconv'):
b, h, w, c = int(x.shape[0]), int(x.shape[1]), int(x.shape[2]), int(x.shape[3])
with tf.variable_scope(name_or_scope=name):
w = tf.get_variable(name='conv2d/kernel', shape=[ksize, ksize, c, cnum], dtype=tf.float32)
b = tf.get_variable(name='conv2d/bias', shape=[cnum], dtype=tf.float32)
# conv(x*m)
x1 = tf.nn.conv2d(x*mask, w, strides=[1, stride, stride, 1], padding='SAME')
# no grad conv(m)
w_noGrad = tf.ones_like(w, dtype=tf.float32)
sum_m = tf.nn.conv2d(mask, w_noGrad, strides=[1, stride, stride, 1], padding='SAME')
mask_one = tf.ones_like(mask)
sum_1 = tf.nn.conv2d(mask_one,w_noGrad, strides=[1, stride, stride, 1], padding='SAME')
# feature output
output = tf.nn.bias_add((x1*sum_1)/(sum_m+1e-8), b, name='feature_output')
# update mask
mask_ouput = tf.clip_by_value(sum_m, 0, 1,name='mask_output')
return output, mask_ouput