def transform_frame(self, frame):
grid = make_coordinate_grid(frame.shape[2:])
grid = F.reshape(grid, (1, frame.shape[2] * frame.shape[3], 2))
grid = self.warp_coordinates(grid)
grid = F.reshape(grid, (self.bs, frame.shape[2], frame.shape[3], 2))
return F.warp_by_grid(frame, grid, padding_mode="reflect", align_corners=True)
def gaussian2kp(heatmap):
shape = heatmap.shape
heatmap = F.reshape(heatmap, shape + (1, ), inplace=False)
grid = make_coordinate_grid(shape[2:])
grid = F.reshape(grid, (1, 1) + grid.shape)
value = F.sum(heatmap * grid, axis=(2, 3))
kp = {'value': value}
return kp
def __init__(self, bs, **kwargs):
noise = F.randn(mu=0, sigma=kwargs['sigma_affine'], shape=(bs, 2, 3))
self.theta = noise + \
nn.Variable.from_numpy_array(
np.array([[[1., 0., 0.], [0., 1., 0.]]]))
self.bs = bs
if ('sigma_tps' in kwargs) and ('points_tps' in kwargs):
self.tps = True
self.control_points = make_coordinate_grid(
(kwargs['points_tps'], kwargs['points_tps']))
self.control_points = F.reshape(
self.control_points, (1,) + self.control_points.shape)
self.control_params = F.randn(
mu=0, sigma=kwargs['sigma_tps'], shape=(bs, 1, kwargs['points_tps'] ** 2))
else:
self.tps = False
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