def generate_shading_image(self, permanent_factor, lighting_factor,
azimuth_factor):
"""Generate log shading images from the input factors.
Given a set of input factors, decode a log shading image that can be used
to relight a given scene. The generated shading image is comprised of a
standard shading image as well as predictions of two global illuminations.
The system can model lighter and darker areas as well as mixed exposure to
different global illuminants such as being illuminated by the yellow
sun and the blue sky.
Args:
permanent_factor: [B, H // 8, W // 8, D] A representation that encodes the
scene property that we wish to relight. H and W are dimensions of the
panorama stack used to encode permanent_factor.
lighting_factor: [B, lighting_dim] A style vector sampled from the
multivariate normal distribution predicted by the illumination encoder.
azimuth_factor: [B] An angle in radians that describe the offset of the
sun's position from the center of the image.
Returns:
A log shading image, of shape [B, H, W, D], of a scene encoded by
permanent_factor illuminated by lighting_factor and azimuth_factor.
"""
UPSAMPLING_FACTOR = 8
factor_h, factor_w = permanent_factor.shape.as_list()[1:3]
output_h = UPSAMPLING_FACTOR * factor_h
output_w = UPSAMPLING_FACTOR * factor_w
pad_fn = lambda x, y: utils.pad_panorama_for_convolutions(
x, y, "reflect")
filters_up = [256, 128, 128, 64]
bsz = tf.shape(permanent_factor)[0]
# Rotate the permanent_factor by azimuth_factor such that the generator's
# input is azimuth-normalized representation.
rot_permanent = pano_transformer.rotate_pano_horizontally(
permanent_factor, azimuth_factor)
# Positional encoding to break the generator's shift invariance.
radian_vec = tf.linspace(
-np.pi, np.pi, output_w + 1)[tf.newaxis, :-1] + tf.zeros([bsz, 1])
cos_vec = tf.cos(radian_vec)
sin_vec = tf.sin(radian_vec)
circular_embed = tf.stack([cos_vec, sin_vec], axis=-1)[:, None, :, :]
circular_embed = tf.tile(circular_embed, [1, output_h, 1, 1])
spade_normalization = lambda inp, scope: spade.sn_spade_normalize(
inp,
rot_permanent,
circular_embed,
scope,
is_training=self.is_training)
postconvfn_actnorm = lambda inp, scope: tf.nn.leaky_relu(
spade_normalization(inp, scope))
with tf.compat.v1.variable_scope("decomp_internal",
reuse=tf.AUTO_REUSE):
with tf.compat.v1.variable_scope("0_shade"):
net = sn_ops.snlinear(lighting_factor,
2 * self.lighting_dim,
is_training=self.is_training,
name="g_sn_enc0")
net = tf.nn.leaky_relu(net)
net = sn_ops.snlinear(net,
4 * self.lighting_dim,
is_training=self.is_training,
name="g_sn_enc1")
net = tf.nn.leaky_relu(net)
net = sn_ops.snlinear(net,
8 * self.lighting_dim,
is_training=self.is_training,
name="g_sn_enc2")
net = tf.nn.leaky_relu(net)
net = sn_ops.snlinear(
net, (self.intial_reshape[0] * self.intial_reshape[1] *
self.lighting_dim),
is_training=self.is_training,
name="g_sn_enc3")
# Decode global illuminations from the style vector.
rgb1 = sn_ops.snlinear(tf.nn.leaky_relu(net),
3,
is_training=self.is_training,
name="g_sn_rgb")
rgb2 = sn_ops.snlinear(tf.nn.leaky_relu(net),
3,
is_training=self.is_training,
name="g_sn_rgb2")
rgb1 = tf.reshape(rgb1, [bsz, 1, 1, 3])
rgb2 = tf.reshape(rgb2, [bsz, 1, 1, 3])
# Reshape the vector to a spatial representation.
netreshape = tf.reshape(
net, (bsz, self.intial_reshape[0], self.intial_reshape[1],
self.lighting_dim))
net = sn_ops.snconv2d(postconvfn_actnorm(netreshape, "enc"),
256,
1,
1,
is_training=self.is_training,
name="g_sn_conv0")
net = utils.upsample(net)
for i in range(len(filters_up)):
# Because k is different at each block, we need to learn skip.
skip = postconvfn_actnorm(net, "skiplayer_%d_0" % i)
skip = pad_fn(
skip,
self.kernel_size,
)
skip = sn_ops.snconv2d(skip,
filters_up[i],
self.kernel_size,
1,
is_training=self.is_training,
name="g_skip_sn_conv%d_0" % i)
net = postconvfn_actnorm(net, "layer_%d_0" % i)
net = pad_fn(
net,
self.kernel_size,
)
net = sn_ops.snconv2d(net,
filters_up[i],
self.kernel_size,
1,
is_training=self.is_training,
name="g_sn_conv%d_0" % i)
net = postconvfn_actnorm(net, "layer_%d_1" % i)
# net = pad_fn(net, self.kernel_size,)
net = sn_ops.snconv2d(net,
filters_up[i],
1,
1,
is_training=self.is_training,
name="g_sn_conv%d_1" % i)
net = utils.upsample(net + skip)
net = pad_fn(net, 5)
# Predict a standard gray-scale log shading.
monochannel_shading = sn_ops.snconv2d(
net,
1,
5,
1,
name="output1",
is_training=self.is_training,
use_bias=False)
# Predicts the influence of each global color illuminant at every pixel.
mask_shading = tf.nn.sigmoid(
sn_ops.snconv2d(net,
1,
5,
1,
is_training=self.is_training,
name="output2"))
mixed_lights = mask_shading * rgb1 + (1 - mask_shading) * rgb2
# Restore the original orientation of the sun position.
log_shading = pano_transformer.rotate_pano_horizontally(
monochannel_shading + mixed_lights, -azimuth_factor)
return log_shading
def extract_illumination(self, aligned_stack):
"""Predicts a style vector and sun azimuth distribution for aligned_stack.
The illumination factors are predicted from individual panoramas. For each
panorama we predict a distribution over style vectors using a VAE and a
distribution over sun positions in the panorama with a fully-convolutional
encoder.
Args:
aligned_stack: [B, H, W, 3] An aligned panorama stack.
Returns:
A dictionary containing (mu, logvar) style vectors of shape [B, D], where
D is self.lighting_dim, and azimuth_distribution of shape [B, 40].
"""
postconvfn_actnorm = lambda x, y: tf.nn.leaky_relu(
self.normalization_illumination(x, y))
pad_fn = lambda x, y: utils.pad_panorama_for_convolutions(
x, y, "reflect")
with tf.compat.v1.variable_scope("decomp_internal",
reuse=tf.AUTO_REUSE):
with tf.compat.v1.variable_scope("shading"):
net = layers.conv2d(pad_fn(aligned_stack, self.kernel_size),
16,
self.kernel_size,
2,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
net = pad_fn(postconvfn_actnorm(net, "enc1"), self.kernel_size)
net = layers.conv2d(net,
32,
self.kernel_size,
2,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
net = pad_fn(postconvfn_actnorm(net, "enc2"), self.kernel_size)
net = layers.conv2d(net,
64,
self.kernel_size,
2,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
net = pad_fn(postconvfn_actnorm(net, "enc3"), self.kernel_size)
net = layers.conv2d(net,
128,
self.kernel_size,
2,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
# Compute horizontal azimuth bins from a feature map by learning a
# convolution layer that reduces over the height of the panorama. This
# distribution is fully convolutional along the width of the panoramas.
feature_map_height = net.shape.as_list()[1]
with tf.compat.v1.variable_scope("azimuth"):
hm_net = layers.conv2d(tf.nn.leaky_relu(net),
1, [feature_map_height, 1],
1,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
azimuth_distribution = tf.nn.softmax(hm_net, axis=-2)
azimuth_distribution = azimuth_distribution[:, 0, :, 0]
net = pad_fn(postconvfn_actnorm(net, "enc4"), self.kernel_size)
net = layers.conv2d(net,
256,
self.kernel_size,
2,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
# Transform the spatial map into a global style vector.
mean_net = tf.reduce_mean(net, axis=[1, 2])
net = layers.fully_connected(mean_net, 128, activation_fn=None)
net = layers.fully_connected(tf.nn.leaky_relu(net),
64,
activation_fn=None)
x = tf.nn.leaky_relu(net)
# Predict parameters of a multivariate normal distribution.
mu = layers.fully_connected(x,
self.lighting_dim,
activation_fn=None)
logvar = layers.fully_connected(x,
self.lighting_dim,
activation_fn=None)
return {
"mu": mu,
"logvar": logvar,
"azimuth_distribution": azimuth_distribution
}
def extract_permanent(self, aligned_stack):
"""Encode aligned_stack to individual permanent factors.
While a stack is expected to have a shared permanent factor, the permanent
encoder predicts a permanent factor from each individual frame. The choice
of extracting a single shared representation from multiple permanent factors
is determined outside of this function.
Args:
aligned_stack: [B, H, W, 3] An aligned panorama stack.
Returns:
A [B, H//8, W//8, D] feature map where D is self.permanent_dim.
"""
postconvfn_actnorm = lambda x, y: tf.nn.leaky_relu(
self.normalization_permanent(x, y))
pad_fn = lambda x, y: utils.pad_panorama_for_convolutions(
x, y, "reflect")
with tf.compat.v1.variable_scope("decomp_internal",
reuse=tf.AUTO_REUSE):
with tf.compat.v1.variable_scope("geometry"):
net = layers.conv2d(pad_fn(aligned_stack, self.kernel_size),
16,
self.kernel_size,
2,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
net = pad_fn(postconvfn_actnorm(net, "enc1"), self.kernel_size)
net = layers.conv2d(net,
32,
self.kernel_size,
2,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
net = pad_fn(postconvfn_actnorm(net, "enc2"), self.kernel_size)
net = layers.conv2d(net,
64,
self.kernel_size,
2,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
for i in range(2):
net_pad = pad_fn(postconvfn_actnorm(net, "res1_%d" % i),
self.kernel_size)
res_net = layers.conv2d(net_pad,
128,
self.kernel_size,
1,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
res_net_pad = pad_fn(
postconvfn_actnorm(res_net, "res2_%d" % i),
self.kernel_size)
res_net = layers.conv2d(res_net_pad,
128,
self.kernel_size,
1,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
skip_pad = pad_fn(postconvfn_actnorm(net, "skip_%d" % i),
self.kernel_size)
skip = layers.conv2d(skip_pad,
128,
self.kernel_size,
1,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
net = res_net + skip
old_net = pad_fn(net, self.kernel_size)
net = layers.conv2d(old_net,
self.permanent_dim,
self.kernel_size,
1,
padding="VALID",
activation_fn=None,
normalizer_fn=None)
return net
def sn_spade_normalize(tensor,
condition,
second_condition=None,
scope="spade",
is_training=False):
"""A spectral normalized version of SPADE.
Performs SPADE normalization (Park et al.) on a tensor based on condition. If
second_condition is defined, concatenate second_condition to condition.
These inputs are separated because they encode conditioning of different
things.
Args:
tensor: [B, H, W, D] a tensor to apply SPADE normalization to
condition: [B, H', W', D'] A tensor used to predict SPADE's normalization
parameters.
second_condition: [B, H'', W'', D''] A tensor used to encode another kind of
conditioning. second_condition is provided in case its dimensions do not
natively match condition's dimension.
scope: (str) The scope of the SPADE convolutions.
is_training: (bool) used to control the spectral normalization update
schedule. When true apply an update, else freeze updates.
Returns:
A SPADE normalized tensor of shape [B, H, W, D].
"""
# resize condition to match input spatial
n_tensor = layers.instance_norm(tensor,
center=False,
scale=False,
trainable=False,
epsilon=1e-4)
unused_b, h_tensor, w_tensor, feature_dim = n_tensor.shape.as_list()
with tf.compat.v1.variable_scope(scope, reuse=tf.AUTO_REUSE):
resize_condition = diff_resize_area(condition, [h_tensor, w_tensor])
if second_condition is not None:
second_condition = diff_resize_area(second_condition,
[h_tensor, w_tensor])
resize_condition = tf.concat([resize_condition, second_condition],
axis=-1)
resize_condition = utils.pad_panorama_for_convolutions(
resize_condition, 3, "symmetric")
condition_net = tf.nn.relu(
sn_ops.snconv2d(resize_condition,
32,
3,
1,
is_training=is_training,
name="intermediate_spade"))
condition_net = utils.pad_panorama_for_convolutions(
condition_net, 3, "symmetric")
gamma_act = sn_ops.snconv2d(condition_net,
feature_dim,
3,
1,
is_training=is_training,
name="g_spade")
mu_act = sn_ops.snconv2d(condition_net,
feature_dim,
3,
1,
is_training=is_training,
name="b_spade")
return n_tensor * (1 + gamma_act) - mu_act
