def make_feed_dict(data, init=False, **params):
if type(data) is tuple:
x, y = data
else:
x = data
y = None
x = np.cast[np.float32]((x - 127.5) / 127.5) ## preprocessing
if args.use_coordinates:
g = grid.generate_grid((x.shape[1], x.shape[2]), batch_size=x.shape[0])
xg = np.concatenate([x, g], axis=-1)
xg, _ = uf.random_crop_images(xg,
output_size=(args.input_size,
args.input_size))
x, g = xg[:, :, :, :3], xg[:, :, :, 3:]
else:
x, _ = uf.random_crop_images(x,
output_size=(args.input_size,
args.input_size))
if 'mask_generator' in params:
mgen = params['mask_generator']
ms = mgen.gen(x.shape[0])
x_masked = x * uf.broadcast_mask(ms, 3)
x_masked = np.concatenate(
[x_masked, uf.broadcast_mask(ms, 1)], axis=-1)
# global conditioning
if args.global_conditional:
global_lv = []
if 'z' in params:
global_lv.append(params['z'])
global_lv = np.concatenate(global_lv, axis=-1)
# spatial conditioning
if args.spatial_conditional:
spatial_lv = []
if args.context_conditioning:
spatial_lv.append(x_masked)
spatial_lv = np.concatenate(spatial_lv, axis=-1)
if init:
feed_dict = {x_init: x}
if args.global_conditional:
feed_dict.update({gh_init: global_lv})
if args.spatial_conditional:
feed_dict.update({sh_init: spatial_lv})
if args.use_coordinates:
c1 = g
c2 = grid.zoom_batch(c1, [obs_shape[0] // 2, obs_shape[1] // 2])
c4 = grid.zoom_batch(c1, [obs_shape[0] // 4, obs_shape[1] // 4])
feed_dict.update({ch_1_init: c1})
feed_dict.update({ch_2_init: c2})
feed_dict.update({ch_4_init: c4})
else:
x = np.split(x, args.nr_gpu)
feed_dict = {xs[i]: x[i] for i in range(args.nr_gpu)}
if args.global_conditional:
global_lv = np.split(global_lv, args.nr_gpu)
feed_dict.update(
{ghs[i]: global_lv[i]
for i in range(args.nr_gpu)})
if args.spatial_conditional:
spatial_lv = np.split(spatial_lv, args.nr_gpu)
feed_dict.update(
{shs[i]: spatial_lv[i]
for i in range(args.nr_gpu)})
if args.use_coordinates:
c1 = g
c2 = grid.zoom_batch(c1, [obs_shape[0] // 2, obs_shape[1] // 2])
c4 = grid.zoom_batch(c1, [obs_shape[0] // 4, obs_shape[1] // 4])
c1 = np.split(c1, args.nr_gpu)
c2 = np.split(c2, args.nr_gpu)
c4 = np.split(c4, args.nr_gpu)
feed_dict.update({ch_1[i]: c1[i] for i in range(args.nr_gpu)})
feed_dict.update({ch_2[i]: c2[i] for i in range(args.nr_gpu)})
feed_dict.update({ch_4[i]: c4[i] for i in range(args.nr_gpu)})
return feed_dict
def complete(sess, data, mask, **params):
if type(data) is tuple:
x, y = data
else:
x = data
y = None
x = np.cast[np.float32]((x - 127.5) / 127.5) ## preprocessing
# mask images
masks = uf.broadcast_mask(mask, 3, x.shape[0])
x *= masks
if 'x_hats' in params:
x_hats = params['x_hats']
x_hats = (x_hats * 2.) - 1.
x_ret = np.split(x, args.nr_gpu)
# global conditioning
if args.global_conditional:
global_lv = []
if 'z' in params:
global_lv.append(params['z'])
global_lv = np.concatenate(global_lv, axis=-1)
global_g = grid.generate_grid((x.shape[1], x.shape[2]),
batch_size=x.shape[0])
if args.global_conditional:
global_lv = np.split(global_lv, args.nr_gpu)
feed_dict.update({ghs[i]: global_lv[i] for i in range(args.nr_gpu)})
while True:
# find the next pixel and the corresonding window
p = uf.find_next_missing_pixel(mask)
if p is None:
break
window = uf.find_maximally_conditioned_window(mask, 32, p)
print(p, window)
[[h0, h1], [w0, w1]] = window
g = global_g[:, h0 - 2:h1 + 2, w0 - 2:w1 + 2, :]
# mw = mask[h0:h1, w0:w1]
# xw = x[:, h0:h1, w0:w1, :]
x_hatsw = x_hats[:, h0 - 2:h1 + 2, w0 - 2:w1 + 2, :]
x_hatsws = np.split(x_hatsw, args.nr_gpu)
yi, xi = p[0] - h0, p[1] - w0
# spatial conditioning
if args.spatial_conditional:
spatial_lv = []
if 'use_coordinates' in params and params['use_coordinates']:
spatial_lv.append(g)
if 'x_hats' in params:
spatial_lv.append(x_hatsw)
spatial_lv = np.concatenate(spatial_lv, axis=-1)
if args.spatial_conditional:
spatial_lv = np.split(spatial_lv, args.nr_gpu)
feed_dict.update(
{shs[i]: spatial_lv[i]
for i in range(args.nr_gpu)})
x_gen = [
x_ret[i][:, h0:h1, w0:w1, :].copy() for i in range(args.nr_gpu)
] # np.split(xw, args.nr_gpu)
feed_dict.update({xs[i]: x_gen[i] for i in range(args.nr_gpu)})
new_x_gen_np = sess.run(new_x_gen, feed_dict=feed_dict)
for i in range(args.nr_gpu):
x_ret[i][:, p[0], p[1], :] = new_x_gen_np[i][:, yi, xi, :]
mask[p[0], p[1]] = 1
return np.concatenate(x_ret, axis=0)
def sample_from_model(sess, data=None, **params):
if type(data) is tuple:
x, y = data
else:
x = data
y = None
x = np.cast[np.float32]((x - 127.5) / 127.5) ## preprocessing
if args.use_coordinates:
g = grid.generate_grid((x.shape[1], x.shape[2]), batch_size=x.shape[0])
xg = np.concatenate([x, g], axis=-1)
xg, _ = uf.random_crop_images(xg,
output_size=(args.input_size,
args.input_size))
x, g = xg[:, :, :, :3], xg[:, :, :, 3:]
else:
x, _ = uf.random_crop_images(x,
output_size=(args.input_size,
args.input_size))
if 'mask_generator' in params:
mgen = params['mask_generator']
ms = mgen.gen(x.shape[0])
x_masked = x * uf.broadcast_mask(ms, 3)
x_masked = np.concatenate(
[x_masked, uf.broadcast_mask(ms, 1)], axis=-1)
# global conditioning
if args.global_conditional:
global_lv = []
if 'z' in params:
global_lv.append(params['z'])
global_lv = np.concatenate(global_lv, axis=-1)
# spatial conditioning
if args.spatial_conditional:
spatial_lv = []
if args.context_conditioning:
spatial_lv.append(x_masked)
spatial_lv = np.concatenate(spatial_lv, axis=-1)
feed_dict = {} ##
# coordinates conditioning:
if args.use_coordinates:
c1 = g
c2 = grid.zoom_batch(c1, [obs_shape[0] // 2, obs_shape[1] // 2])
c4 = grid.zoom_batch(c1, [obs_shape[0] // 4, obs_shape[1] // 4])
c1 = np.split(c1, args.nr_gpu)
c2 = np.split(c2, args.nr_gpu)
c4 = np.split(c4, args.nr_gpu)
feed_dict.update({ch_1[i]: c1[i] for i in range(args.nr_gpu)})
feed_dict.update({ch_2[i]: c2[i] for i in range(args.nr_gpu)})
feed_dict.update({ch_4[i]: c4[i] for i in range(args.nr_gpu)})
if args.global_conditional:
global_lv = np.split(global_lv, args.nr_gpu)
feed_dict.update({ghs[i]: global_lv[i] for i in range(args.nr_gpu)})
if args.spatial_conditional:
spatial_lv = np.split(spatial_lv, args.nr_gpu)
feed_dict.update({shs[i]: spatial_lv[i] for i in range(args.nr_gpu)})
if 'mask_generator' in params:
x_gen = np.split(x_masked[:, :, :, :3], args.nr_gpu)
else:
x_gen = [np.zeros_like(x) for i in range(args.nr_gpu)]
for yi in range(obs_shape[0]):
for xi in range(obs_shape[1]):
if ('mask_generator' not in params) or ms[0][yi, xi] == 0:
feed_dict.update({xs[i]: x_gen[i] for i in range(args.nr_gpu)})
new_x_gen_np = sess.run(new_x_gen, feed_dict=feed_dict)
for i in range(args.nr_gpu):
x_gen[i][:, yi, xi, :] = new_x_gen_np[i][:, yi, xi, :]
return np.concatenate(x_gen, axis=0)
def complete(sess, data, mask, **params):
if type(data) is tuple:
x, y = data
else:
x = data
y = None
x = np.cast[np.float32]((x - 127.5) / 127.5) ## preprocessing
# mask images
x_ret = x * uf.broadcast_mask(mask, 3, x.shape[0])
x_ret = np.split(x_ret, args.nr_gpu)
x_masked = np.concatenate([
np.concatenate(x_ret, axis=0),
uf.broadcast_mask(mask, 1, x.shape[0])
],
axis=-1)
# global conditioning
if args.global_conditional:
global_lv = []
if 'z' in params:
global_lv.append(params['z'])
global_lv = np.concatenate(global_lv, axis=-1)
if args.global_conditional:
global_lv = np.split(global_lv, args.nr_gpu)
feed_dict.update({ghs[i]: global_lv[i] for i in range(args.nr_gpu)})
global_g = grid.generate_grid((x.shape[1], x.shape[2]),
batch_size=x.shape[0])
while True:
# find the next pixel and the corresonding window
p = uf.find_next_missing_pixel(mask)
if p is None:
break
window = uf.find_maximally_conditioned_window(mask, 32, p)
print(p, window)
[[h0, h1], [w0, w1]] = window
g = global_g[:, h0:h1, w0:w1, :]
x_masked_w = x_masked[:, h0:h1, w0:w1, :]
yi, xi = p[0] - h0, p[1] - w0
# spatial conditioning
if args.spatial_conditional:
spatial_lv = []
if args.context_conditioning:
spatial_lv.append(x_masked_w)
spatial_lv = np.concatenate(spatial_lv, axis=-1)
if args.spatial_conditional:
spatial_lv = np.split(spatial_lv, args.nr_gpu)
feed_dict.update(
{shs[i]: spatial_lv[i]
for i in range(args.nr_gpu)})
# coordinates conditioning:
if args.use_coordinates:
c1 = g
c2 = grid.zoom_batch(c1, [obs_shape[0] // 2, obs_shape[1] // 2])
c4 = grid.zoom_batch(c1, [obs_shape[0] // 4, obs_shape[1] // 4])
c1 = np.split(c1, args.nr_gpu)
c2 = np.split(c2, args.nr_gpu)
c4 = np.split(c4, args.nr_gpu)
feed_dict.update({ch_1[i]: c1[i] for i in range(args.nr_gpu)})
feed_dict.update({ch_2[i]: c2[i] for i in range(args.nr_gpu)})
feed_dict.update({ch_4[i]: c4[i] for i in range(args.nr_gpu)})
x_gen = [x_ret[i][:, h0:h1, w0:w1, :] for i in range(args.nr_gpu)]
feed_dict.update({xs[i]: x_gen[i] for i in range(args.nr_gpu)})
new_x_gen_np = sess.run(new_x_gen, feed_dict=feed_dict)
for i in range(args.nr_gpu):
x_ret[i][:, p[0], p[1], :] = new_x_gen_np[i][:, yi, xi, :]
mask[p[0], p[1]] = 1
#x_masked = np.concatenate([np.concatenate(x_ret, axis=0), uf.broadcast_mask(mask, 1, x.shape[0])], axis=-1)
return np.concatenate(x_ret, axis=0)
