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