https://github.com/JiahuiYu/generative_inpainting
Tip revision: 3a5324373ba52c68c79587ca183bc10b9e57b783 authored by JiahuiYu on 10 June 2020, 15:16:02 UTC
Fix typo.
Fix typo.
Tip revision: 3a53243
inpaint_ops.py
import logging
import math
import cv2
import numpy as np
import tensorflow as tf
from tensorflow.contrib.framework.python.ops import add_arg_scope
from PIL import Image, ImageDraw
from neuralgym.ops.layers import resize
from neuralgym.ops.layers import *
from neuralgym.ops.loss_ops import *
from neuralgym.ops.gan_ops import *
from neuralgym.ops.summary_ops import *
logger = logging.getLogger()
np.random.seed(2018)
@add_arg_scope
def gen_conv(x, cnum, ksize, stride=1, rate=1, name='conv',
padding='SAME', activation=tf.nn.elu, training=True):
"""Define conv for generator.
Args:
x: Input.
cnum: Channel number.
ksize: Kernel size.
Stride: Convolution stride.
Rate: Rate for or dilated conv.
name: Name of layers.
padding: Default to SYMMETRIC.
activation: Activation function after convolution.
training: If current graph is for training or inference, used for bn.
Returns:
tf.Tensor: output
"""
assert padding in ['SYMMETRIC', 'SAME', 'REFELECT']
if padding == 'SYMMETRIC' or padding == 'REFELECT':
p = int(rate*(ksize-1)/2)
x = tf.pad(x, [[0,0], [p, p], [p, p], [0,0]], mode=padding)
padding = 'VALID'
x = tf.layers.conv2d(
x, cnum, ksize, stride, dilation_rate=rate,
activation=None, padding=padding, name=name)
if cnum == 3 or activation is None:
# conv for output
return x
x, y = tf.split(x, 2, 3)
x = activation(x)
y = tf.nn.sigmoid(y)
x = x * y
return x
@add_arg_scope
def gen_deconv(x, cnum, name='upsample', padding='SAME', training=True):
"""Define deconv for generator.
The deconv is defined to be a x2 resize_nearest_neighbor operation with
additional gen_conv operation.
Args:
x: Input.
cnum: Channel number.
name: Name of layers.
training: If current graph is for training or inference, used for bn.
Returns:
tf.Tensor: output
"""
with tf.variable_scope(name):
x = resize(x, func=tf.image.resize_nearest_neighbor)
x = gen_conv(
x, cnum, 3, 1, name=name+'_conv', padding=padding,
training=training)
return x
@add_arg_scope
def dis_conv(x, cnum, ksize=5, stride=2, name='conv', training=True):
"""Define conv for discriminator.
Activation is set to leaky_relu.
Args:
x: Input.
cnum: Channel number.
ksize: Kernel size.
Stride: Convolution stride.
name: Name of layers.
training: If current graph is for training or inference, used for bn.
Returns:
tf.Tensor: output
"""
x = conv2d_spectral_norm(x, cnum, ksize, stride, 'SAME', name=name)
x = tf.nn.leaky_relu(x)
return x
def random_bbox(FLAGS):
"""Generate a random tlhw.
Returns:
tuple: (top, left, height, width)
"""
img_shape = FLAGS.img_shapes
img_height = img_shape[0]
img_width = img_shape[1]
maxt = img_height - FLAGS.vertical_margin - FLAGS.height
maxl = img_width - FLAGS.horizontal_margin - FLAGS.width
t = tf.random_uniform(
[], minval=FLAGS.vertical_margin, maxval=maxt, dtype=tf.int32)
l = tf.random_uniform(
[], minval=FLAGS.horizontal_margin, maxval=maxl, dtype=tf.int32)
h = tf.constant(FLAGS.height)
w = tf.constant(FLAGS.width)
return (t, l, h, w)
def bbox2mask(FLAGS, bbox, name='mask'):
"""Generate mask tensor from bbox.
Args:
bbox: tuple, (top, left, height, width)
Returns:
tf.Tensor: output with shape [1, H, W, 1]
"""
def npmask(bbox, height, width, delta_h, delta_w):
mask = np.zeros((1, height, width, 1), np.float32)
h = np.random.randint(delta_h//2+1)
w = np.random.randint(delta_w//2+1)
mask[:, bbox[0]+h:bbox[0]+bbox[2]-h,
bbox[1]+w:bbox[1]+bbox[3]-w, :] = 1.
return mask
with tf.variable_scope(name), tf.device('/cpu:0'):
img_shape = FLAGS.img_shapes
height = img_shape[0]
width = img_shape[1]
mask = tf.py_func(
npmask,
[bbox, height, width,
FLAGS.max_delta_height, FLAGS.max_delta_width],
tf.float32, stateful=False)
mask.set_shape([1] + [height, width] + [1])
return mask
def brush_stroke_mask(FLAGS, name='mask'):
"""Generate mask tensor from bbox.
Returns:
tf.Tensor: output with shape [1, H, W, 1]
"""
min_num_vertex = 4
max_num_vertex = 12
mean_angle = 2*math.pi / 5
angle_range = 2*math.pi / 15
min_width = 12
max_width = 40
def generate_mask(H, W):
average_radius = math.sqrt(H*H+W*W) / 8
mask = Image.new('L', (W, H), 0)
for _ in range(np.random.randint(1, 4)):
num_vertex = np.random.randint(min_num_vertex, max_num_vertex)
angle_min = mean_angle - np.random.uniform(0, angle_range)
angle_max = mean_angle + np.random.uniform(0, angle_range)
angles = []
vertex = []
for i in range(num_vertex):
if i % 2 == 0:
angles.append(2*math.pi - np.random.uniform(angle_min, angle_max))
else:
angles.append(np.random.uniform(angle_min, angle_max))
h, w = mask.size
vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h))))
for i in range(num_vertex):
r = np.clip(
np.random.normal(loc=average_radius, scale=average_radius//2),
0, 2*average_radius)
new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
vertex.append((int(new_x), int(new_y)))
draw = ImageDraw.Draw(mask)
width = int(np.random.uniform(min_width, max_width))
draw.line(vertex, fill=1, width=width)
for v in vertex:
draw.ellipse((v[0] - width//2,
v[1] - width//2,
v[0] + width//2,
v[1] + width//2),
fill=1)
if np.random.normal() > 0:
mask.transpose(Image.FLIP_LEFT_RIGHT)
if np.random.normal() > 0:
mask.transpose(Image.FLIP_TOP_BOTTOM)
mask = np.asarray(mask, np.float32)
mask = np.reshape(mask, (1, H, W, 1))
return mask
with tf.variable_scope(name), tf.device('/cpu:0'):
img_shape = FLAGS.img_shapes
height = img_shape[0]
width = img_shape[1]
mask = tf.py_func(
generate_mask,
[height, width],
tf.float32, stateful=True)
mask.set_shape([1] + [height, width] + [1])
return mask
def local_patch(x, bbox):
"""Crop local patch according to bbox.
Args:
x: input
bbox: (top, left, height, width)
Returns:
tf.Tensor: local patch
"""
x = tf.image.crop_to_bounding_box(x, bbox[0], bbox[1], bbox[2], bbox[3])
return x
def resize_mask_like(mask, x):
"""Resize mask like shape of x.
Args:
mask: Original mask.
x: To shape of x.
Returns:
tf.Tensor: resized mask
"""
mask_resize = resize(
mask, to_shape=x.get_shape().as_list()[1:3],
func=tf.image.resize_nearest_neighbor)
return mask_resize
def contextual_attention(f, b, mask=None, ksize=3, stride=1, rate=1,
fuse_k=3, softmax_scale=10., training=True, fuse=True):
""" Contextual attention layer implementation.
Contextual attention is first introduced in publication:
Generative Image Inpainting with Contextual Attention, Yu et al.
Args:
x: Input feature to match (foreground).
t: Input feature for match (background).
mask: Input mask for t, indicating patches not available.
ksize: Kernel size for contextual attention.
stride: Stride for extracting patches from t.
rate: Dilation for matching.
softmax_scale: Scaled softmax for attention.
training: Indicating if current graph is training or inference.
Returns:
tf.Tensor: output
"""
# get shapes
raw_fs = tf.shape(f)
raw_int_fs = f.get_shape().as_list()
raw_int_bs = b.get_shape().as_list()
# extract patches from background with stride and rate
kernel = 2*rate
raw_w = tf.extract_image_patches(
b, [1,kernel,kernel,1], [1,rate*stride,rate*stride,1], [1,1,1,1], padding='SAME')
raw_w = tf.reshape(raw_w, [raw_int_bs[0], -1, kernel, kernel, raw_int_bs[3]])
raw_w = tf.transpose(raw_w, [0, 2, 3, 4, 1]) # transpose to b*k*k*c*hw
# downscaling foreground option: downscaling both foreground and
# background for matching and use original background for reconstruction.
f = resize(f, scale=1./rate, func=tf.image.resize_nearest_neighbor)
b = resize(b, to_shape=[int(raw_int_bs[1]/rate), int(raw_int_bs[2]/rate)], func=tf.image.resize_nearest_neighbor) # https://github.com/tensorflow/tensorflow/issues/11651
if mask is not None:
mask = resize(mask, scale=1./rate, func=tf.image.resize_nearest_neighbor)
fs = tf.shape(f)
int_fs = f.get_shape().as_list()
f_groups = tf.split(f, int_fs[0], axis=0)
# from t(H*W*C) to w(b*k*k*c*h*w)
bs = tf.shape(b)
int_bs = b.get_shape().as_list()
w = tf.extract_image_patches(
b, [1,ksize,ksize,1], [1,stride,stride,1], [1,1,1,1], padding='SAME')
w = tf.reshape(w, [int_fs[0], -1, ksize, ksize, int_fs[3]])
w = tf.transpose(w, [0, 2, 3, 4, 1]) # transpose to b*k*k*c*hw
# process mask
if mask is None:
mask = tf.zeros([1, bs[1], bs[2], 1])
m = tf.extract_image_patches(
mask, [1,ksize,ksize,1], [1,stride,stride,1], [1,1,1,1], padding='SAME')
m = tf.reshape(m, [1, -1, ksize, ksize, 1])
m = tf.transpose(m, [0, 2, 3, 4, 1]) # transpose to b*k*k*c*hw
m = m[0]
mm = tf.cast(tf.equal(tf.reduce_mean(m, axis=[0,1,2], keep_dims=True), 0.), tf.float32)
w_groups = tf.split(w, int_bs[0], axis=0)
raw_w_groups = tf.split(raw_w, int_bs[0], axis=0)
y = []
offsets = []
k = fuse_k
scale = softmax_scale
fuse_weight = tf.reshape(tf.eye(k), [k, k, 1, 1])
for xi, wi, raw_wi in zip(f_groups, w_groups, raw_w_groups):
# conv for compare
wi = wi[0]
wi_normed = wi / tf.maximum(tf.sqrt(tf.reduce_sum(tf.square(wi), axis=[0,1,2])), 1e-4)
yi = tf.nn.conv2d(xi, wi_normed, strides=[1,1,1,1], padding="SAME")
# conv implementation for fuse scores to encourage large patches
if fuse:
yi = tf.reshape(yi, [1, fs[1]*fs[2], bs[1]*bs[2], 1])
yi = tf.nn.conv2d(yi, fuse_weight, strides=[1,1,1,1], padding='SAME')
yi = tf.reshape(yi, [1, fs[1], fs[2], bs[1], bs[2]])
yi = tf.transpose(yi, [0, 2, 1, 4, 3])
yi = tf.reshape(yi, [1, fs[1]*fs[2], bs[1]*bs[2], 1])
yi = tf.nn.conv2d(yi, fuse_weight, strides=[1,1,1,1], padding='SAME')
yi = tf.reshape(yi, [1, fs[2], fs[1], bs[2], bs[1]])
yi = tf.transpose(yi, [0, 2, 1, 4, 3])
yi = tf.reshape(yi, [1, fs[1], fs[2], bs[1]*bs[2]])
# softmax to match
yi *= mm # mask
yi = tf.nn.softmax(yi*scale, 3)
yi *= mm # mask
offset = tf.argmax(yi, axis=3, output_type=tf.int32)
offset = tf.stack([offset // fs[2], offset % fs[2]], axis=-1)
# deconv for patch pasting
# 3.1 paste center
wi_center = raw_wi[0]
yi = tf.nn.conv2d_transpose(yi, wi_center, tf.concat([[1], raw_fs[1:]], axis=0), strides=[1,rate,rate,1]) / 4.
y.append(yi)
offsets.append(offset)
y = tf.concat(y, axis=0)
y.set_shape(raw_int_fs)
offsets = tf.concat(offsets, axis=0)
offsets.set_shape(int_bs[:3] + [2])
# case1: visualize optical flow: minus current position
h_add = tf.tile(tf.reshape(tf.range(bs[1]), [1, bs[1], 1, 1]), [bs[0], 1, bs[2], 1])
w_add = tf.tile(tf.reshape(tf.range(bs[2]), [1, 1, bs[2], 1]), [bs[0], bs[1], 1, 1])
offsets = offsets - tf.concat([h_add, w_add], axis=3)
# to flow image
flow = flow_to_image_tf(offsets)
# # case2: visualize which pixels are attended
# flow = highlight_flow_tf(offsets * tf.cast(mask, tf.int32))
if rate != 1:
flow = resize(flow, scale=rate, func=tf.image.resize_bilinear)
return y, flow
def test_contextual_attention(args):
"""Test contextual attention layer with 3-channel image input
(instead of n-channel feature).
"""
import cv2
import os
# run on cpu
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
rate = 2
stride = 1
grid = rate*stride
b = cv2.imread(args.imageA)
b = cv2.resize(b, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_CUBIC)
h, w, _ = b.shape
b = b[:h//grid*grid, :w//grid*grid, :]
b = np.expand_dims(b, 0)
logger.info('Size of imageA: {}'.format(b.shape))
f = cv2.imread(args.imageB)
h, w, _ = f.shape
f = f[:h//grid*grid, :w//grid*grid, :]
f = np.expand_dims(f, 0)
logger.info('Size of imageB: {}'.format(f.shape))
with tf.Session() as sess:
bt = tf.constant(b, dtype=tf.float32)
ft = tf.constant(f, dtype=tf.float32)
yt, flow = contextual_attention(
ft, bt, stride=stride, rate=rate,
training=False, fuse=False)
y = sess.run(yt)
cv2.imwrite(args.imageOut, y[0])
def make_color_wheel():
RY, YG, GC, CB, BM, MR = (15, 6, 4, 11, 13, 6)
ncols = RY + YG + GC + CB + BM + MR
colorwheel = np.zeros([ncols, 3])
col = 0
# RY
colorwheel[0:RY, 0] = 255
colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))
col += RY
# YG
colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))
colorwheel[col:col+YG, 1] = 255
col += YG
# GC
colorwheel[col:col+GC, 1] = 255
colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))
col += GC
# CB
colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))
colorwheel[col:col+CB, 2] = 255
col += CB
# BM
colorwheel[col:col+BM, 2] = 255
colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))
col += + BM
# MR
colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
colorwheel[col:col+MR, 0] = 255
return colorwheel
COLORWHEEL = make_color_wheel()
def compute_color(u,v):
h, w = u.shape
img = np.zeros([h, w, 3])
nanIdx = np.isnan(u) | np.isnan(v)
u[nanIdx] = 0
v[nanIdx] = 0
# colorwheel = COLORWHEEL
colorwheel = make_color_wheel()
ncols = np.size(colorwheel, 0)
rad = np.sqrt(u**2+v**2)
a = np.arctan2(-v, -u) / np.pi
fk = (a+1) / 2 * (ncols - 1) + 1
k0 = np.floor(fk).astype(int)
k1 = k0 + 1
k1[k1 == ncols+1] = 1
f = fk - k0
for i in range(np.size(colorwheel,1)):
tmp = colorwheel[:, i]
col0 = tmp[k0-1] / 255
col1 = tmp[k1-1] / 255
col = (1-f) * col0 + f * col1
idx = rad <= 1
col[idx] = 1-rad[idx]*(1-col[idx])
notidx = np.logical_not(idx)
col[notidx] *= 0.75
img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))
return img
def flow_to_image(flow):
"""Transfer flow map to image.
Part of code forked from flownet.
"""
out = []
maxu = -999.
maxv = -999.
minu = 999.
minv = 999.
maxrad = -1
for i in range(flow.shape[0]):
u = flow[i, :, :, 0]
v = flow[i, :, :, 1]
idxunknow = (abs(u) > 1e7) | (abs(v) > 1e7)
u[idxunknow] = 0
v[idxunknow] = 0
maxu = max(maxu, np.max(u))
minu = min(minu, np.min(u))
maxv = max(maxv, np.max(v))
minv = min(minv, np.min(v))
rad = np.sqrt(u ** 2 + v ** 2)
maxrad = max(maxrad, np.max(rad))
u = u/(maxrad + np.finfo(float).eps)
v = v/(maxrad + np.finfo(float).eps)
img = compute_color(u, v)
out.append(img)
return np.float32(np.uint8(out))
def flow_to_image_tf(flow, name='flow_to_image'):
"""Tensorflow ops for computing flow to image.
"""
with tf.variable_scope(name), tf.device('/cpu:0'):
img = tf.py_func(flow_to_image, [flow], tf.float32, stateful=False)
img.set_shape(flow.get_shape().as_list()[0:-1]+[3])
img = img / 127.5 - 1.
return img
def highlight_flow(flow):
"""Convert flow into middlebury color code image.
"""
out = []
s = flow.shape
for i in range(flow.shape[0]):
img = np.ones((s[1], s[2], 3)) * 144.
u = flow[i, :, :, 0]
v = flow[i, :, :, 1]
for h in range(s[1]):
for w in range(s[1]):
ui = u[h,w]
vi = v[h,w]
img[ui, vi, :] = 255.
out.append(img)
return np.float32(np.uint8(out))
def highlight_flow_tf(flow, name='flow_to_image'):
"""Tensorflow ops for highlight flow.
"""
with tf.variable_scope(name), tf.device('/cpu:0'):
img = tf.py_func(highlight_flow, [flow], tf.float32, stateful=False)
img.set_shape(flow.get_shape().as_list()[0:-1]+[3])
img = img / 127.5 - 1.
return img
def image2edge(image):
"""Convert image to edges.
"""
out = []
for i in range(image.shape[0]):
img = cv2.Laplacian(image[i, :, :, :], cv2.CV_64F, ksize=3, scale=2)
out.append(img)
return np.float32(np.uint8(out))
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--imageA', default='', type=str, help='Image A as background patches to reconstruct image B.')
parser.add_argument('--imageB', default='', type=str, help='Image B is reconstructed with image A.')
parser.add_argument('--imageOut', default='result.png', type=str, help='Image B is reconstructed with image A.')
args = parser.parse_args()
test_contextual_attention(args)