def adjust_kp(kp_source,
kp_driving,
kp_driving_initial,
adapt_movement_scale=1,
use_relative_movement=False,
use_relative_jacobian=False):
kp_new = {k: v for k, v in kp_driving.items()}
if use_relative_movement:
kp_value_diff = (kp_driving['value'] - kp_driving_initial['value'])
kp_value_diff *= adapt_movement_scale
kp_new['value'] = kp_value_diff + kp_source['value']
if use_relative_jacobian:
jacobian_diff = F.batch_matmul(
kp_driving['jacobian'],
F.reshape(
F.batch_inv(
F.reshape(kp_driving_initial['jacobian'], (-1, ) +
kp_driving_initial['jacobian'].shape[-2:],
inplace=False)),
kp_driving_initial['jacobian'].shape))
kp_new['jacobian'] = F.batch_matmul(jacobian_diff,
kp_source['jacobian'])
return kp_new
def transform(point, center, scale, resolution, invert=False):
"""Generate and affine transformation matrix.
Given a set of points, a center, a scale and a targer resolution, the
function generates and affine transformation matrix. If invert is ``True``
it will produce the inverse transformation.
Arguments:
point {numpy.array} -- the input 2D point
center {numpy.array} -- the center around which to perform the transformations
scale {float} -- the scale of the face/object
resolution {float} -- the output resolution
Keyword Arguments:
invert {bool} -- define wherever the function should produce the direct or the
inverse transformation matrix (default: {False})
"""
point.append(1)
h = 200.0 * scale
t = F.matrix_diag(F.constant(1, [3]))
t.d[0, 0] = resolution / h
t.d[1, 1] = resolution / h
t.d[0, 2] = resolution * (-center[0] / h + 0.5)
t.d[1, 2] = resolution * (-center[1] / h + 0.5)
if invert:
t = F.reshape(F.batch_inv(F.reshape(t, [1, 3, 3])), [3, 3])
_pt = nn.Variable.from_numpy_array(point)
new_point = F.reshape(F.batch_matmul(
F.reshape(t, [1, 3, 3]), F.reshape(_pt, [1, 3, 1])), [3, ])[0:2]
return new_point.d.astype(int)
def batch_det_backward(inputs):
"""
Args:
inputs (list of nn.Variable): Incomming grads/inputs to/of the forward function.
kwargs (dict of arguments): Dictionary of the corresponding function arguments.
Return:
list of Variable: Return the gradients wrt inputs of the corresponding function.
"""
dy = inputs[0]
x0 = inputs[1]
x0_inv_T = F.transpose(F.batch_inv(x0), [0, 2, 1])
b = x0.shape[0]
dy = F.reshape(dy, [b, 1, 1])
y0 = get_output(x0, "BatchDet")
y0 = F.reshape(y0, [b, 1, 1], inplace=False)
dx0 = dy * y0 * x0_inv_T
return dx0
def equivariance_jacobian_loss(kp_driving_jacobian, arithmetic_jacobian,
trans_kp_jacobian, weight):
jacobian_transformed = F.batch_matmul(arithmetic_jacobian,
trans_kp_jacobian)
normed_driving = F.reshape(
F.batch_inv(
F.reshape(kp_driving_jacobian,
(-1, ) + kp_driving_jacobian.shape[-2:])),
kp_driving_jacobian.shape)
normed_transformed = jacobian_transformed
value = F.batch_matmul(normed_driving, normed_transformed)
eye = nn.Variable.from_numpy_array(np.reshape(np.eye(2), (1, 1, 2, 2)))
jacobian_loss = F.mean(F.absolute_error(eye, value))
loss = weight * jacobian_loss
return loss
def create_sparse_motions(source_image, kp_driving, kp_source, num_kp):
bs, _, h, w = source_image.shape
identity_grid = make_coordinate_grid((h, w))
identity_grid = F.reshape(identity_grid,
(1, 1, h, w, 2)) # (1, 1, h, w, 2)
coordinate_grid = identity_grid - \
F.reshape(kp_driving['value'], (bs, num_kp, 1, 1, 2), inplace=False)
if 'jacobian' in kp_driving:
jacobian = F.batch_matmul(
kp_source['jacobian'],
F.reshape(
F.batch_inv(
F.reshape(kp_driving['jacobian'],
(-1, ) + kp_driving['jacobian'].shape[-2:],
inplace=False)), kp_driving['jacobian'].shape))
# what it does
# batched_driving_jacobian = F.reshape(kp_driving['jacobian'], (-1) + kp_driving['jacobian'].shape[-2:])
# batched_inverse_jacobian = F.batch_inv(batched_driving_jacobian)
# inverse_jacobian = F.reshape(batched_inverse_jacobian, kp_driving['jacobian'].shape)
jacobian = F.reshape(
jacobian, jacobian.shape[:-2] + (1, 1) + jacobian.shape[-2:])
jacobian = F.broadcast(
jacobian, jacobian.shape[:2] + (h, w) + jacobian.shape[-2:])
coordinate_grid = F.batch_matmul(
jacobian, F.reshape(coordinate_grid,
coordinate_grid.shape + (1, )))
coordinate_grid = F.reshape(coordinate_grid,
coordinate_grid.shape[:-1])
driving_to_source = coordinate_grid + \
F.reshape(kp_source['value'], (bs, num_kp, 1, 1, 2), inplace=False)
# background feature
identity_grid = F.broadcast(identity_grid, (bs, 1, h, w, 2))
sparse_motions = F.concatenate(identity_grid, driving_to_source, axis=1)
return sparse_motions
def invertible_conv(x, reverse, rng, scope):
r"""Invertible 1x1 Convolution Layer.
Args:
x (nn.Variable): Input variable.
reverse (bool): Whether it's a reverse direction.
rng (numpy.random.RandomState): A random generator.
scope (str): The scope.
Returns:
nn.Variable: The output variable.
"""
batch_size, c, n_groups = x.shape
with nn.parameter_scope(scope):
# initialize w by an orthonormal matrix
w_init = np.linalg.qr(rng.randn(c, c))[0][None, ...]
W_var = get_parameter_or_create("W", (1, c, c), w_init, True, True)
W = F.batch_inv(W_var) if reverse else W_var
x = F.convolution(x, F.reshape(W, (c, c, 1)), None, stride=(1, ))
if reverse:
return x
log_det = batch_size * n_groups * F.log(F.abs(F.batch_det(W)))
return x, log_det
