Thin-Plate Spline Motion Model for Image Animation

In plain words

Give this system one still photo and a short video of someone moving, and it moves the photo to match. It needs no labels that say where a face or a body's joints are. It learns a small set of points on the object and how they move from the video.

Earlier methods described motion with simple straight-line pieces. This one uses a smoother bending model, which fits curved movements such as an open mouth better. Because part of the picture is hidden or stretched by the motion, the system also works out which regions are missing and fills them in. It uses masks at several scales to do this, and it adds extra training rules so that each part of the network does a distinct job.

TL;DR

Unsupervised image animation that replaces the local affine motion of earlier methods with a thin-plate spline (TPS) warp built from several keypoint groups. Multi-resolution occlusion masks feed a skip-connected inpainting network, and an equivariance loss plus two auxiliary losses (background and warp) supervise training. The paper reports better keypoint and motion metrics than MRAA on most benchmarks, with slightly worse identity scores (AED) on TED-talks and TaiChiHD. All numbers are author-reported.

Problem

Earlier unsupervised methods learn keypoints and turn them into local affine transformations to warp the source image §1 §2. A single affine piece per keypoint cannot represent curved, non-linear motion such as a lip curving open, so large pose gaps between source and driving image are hard to learn §1. The warped feature maps then leave large holes, and the network must inpaint them, which previous methods did with a single occlusion mask for all scales §1 §3.2.

Method

The model has four modules §3: a keypoint detector, a background motion predictor, a dense motion network and an inpainting network. Given a source image S and a driving image D it outputs a reconstruction of D fig. 2.

Keypoints and background. The keypoint detector predicts K groups of N = 5 keypoints for each image §3.1 keypoint_detector.py. The background predictor takes S and D stacked along channels and outputs one 2×3 affine matrix for the background §3.1 bg_motion_predictor.py.

TPS motion estimation. Each group of five source-driving keypoint pairs defines one TPS transformation. The TPS is the smoothest mapping that passes through those pairs, obtained by minimising bending energy eq. 1. Its closed form is an affine part plus a weighted sum of radial basis terms U(r) = r² log r² over the keypoints eq. 2 eq. 3. The background uses a pure affine map eq. 4.

Combining the K+1 transforms. The K TPS maps and the affine background are each used to warp S. The warped images are concatenated and fed to the dense motion hourglass, which predicts K+1 contribution maps. A per-pixel softmax makes them sum to one eq. 5. The optical flow is the contribution-weighted sum of the K+1 transforms eq. 6 §3.1.

Dropout on TPS fields. During early training, the softmax terms of the K TPS fields are dropped at random, with probability P, so the network does not rely on only a few fields §3.1. Kept terms are rescaled by 1/(1−P) so that the expected sum stays the same eq. 7 eq. 8 §3.1. The paper says dropout is removed once every field contributes, after several epochs §3.1.

Multi-resolution occlusion masks. The decoder layers of the dense motion network each predict an occlusion mask with an extra convolution, so the masks cover several scales rather than one §3.2 fig. 7. The code produces four masks dense_motion.py dense_motion.py.

Inpainting. The inpainting network encodes S with a hourglass encoder. At each level, the encoder features are warped by the optical flow, multiplied by the occlusion mask of that resolution, and joined to the decoder through a skip connection. Each decoder level applies two residual blocks and then upsamples. The output is the generated D̂ §3.2 fig. 3.

Training losses. The reconstruction loss is a VGG-19 perceptual loss over several image scales eq. 9 §3.3. The equivariance loss asks the keypoint detector to commute with a random TPS warp eq. 10 §3.3. The background loss asks the forward and reversed affine matrices to be inverses of each other, written as a difference from the identity eq. 11 eq. 12 §3.3. The warp loss asks the warped encoder features of S to match the encoder features of D eq. 13 §3.3. The total is the sum of these four terms, with no weights eq. 14 §3.3.

Animation at test time. The paper uses MRAA's "animation via disentanglement" mode §3.4. A shape encoder reads the source keypoints and a pose encoder reads the driving keypoints. A decoder rebuilds the keypoints from the two codes, and the rebuilt keypoints drive the TPS warps §3.4 avd_network.py.

# Forward pass for one (source, driving) pair  [§3] [eq. 1-6]
# Input:  source, driving (B, 3, H, W)
# Output: reconstruction of driving (B, 3, H, W)

def forward(source, driving):
    # K groups of 5 keypoints per image, values in [-1, 1]  [§3.1]
    kp_source = keypoint_net(source)                 # (B, K, 5, 2)
    kp_driving = keypoint_net(driving)               # (B, K, 5, 2)
    # One affine background matrix  [§3.1] [eq. 4]
    affine_bg = background_net(source, driving)      # (B, 3, 3)
    # Fit one TPS per keypoint group, from driving to source  [§3.1] [eq. 1-3]
    tps_grids = fit_tps(kp_driving, kp_source)       # K grids (B, H, W, 2)
    # Warp grid of the background, as the (K+1)-th transform  [eq. 4]
    grids = [affine_grid(affine_bg)] + tps_grids     # K+1 grids
    # Dense motion works on a 1/4-resolution source  [§3.1]
    flow, masks = dense_motion(source, grids, kp_source, kp_driving)
    return inpaint(source, flow, masks)

# Dropout on the K TPS logits, training only  [§3.1] [eq. 7-8]
# logits: (B, K+1, h, w); channel 0 is the affine background

def dropout_softmax(logits, p):
    # Keep the background always; drop each TPS field with prob. p  [§3.1]
    keep = bernoulli(1 - p, shape=(B, K))
    keep = concat([ones(B, 1), keep], axis=1)
    # Rescale the kept TPS terms by 1 / (1 - p)  [eq. 7]
    # (the reference code does not apply this rescale; see Paper vs code)
    scale = concat([ones(B, 1), ones(B, K) / (1 - p)], axis=1)
    weights = exp(logits) * scale * broadcast(keep)
    return weights / sum(weights, axis=1)            # [eq. 8]

# Training objective for one pair  [§3.3] [eq. 9-14]
# The reference code weights each term by 10 and uses means, not sums.

def training_loss(source, driving):
    recon = forward(source, driving)
    # Perceptual loss, VGG-19 features at several scales  [eq. 9]
    loss_rec = perceptual(recon, driving)
    # Equivariance: keypoints of a warped driving frame  [eq. 10]
    loss_eq = equivariance(driving)
    # Forward and reversed affine matrices should cancel  [eq. 11-12]
    loss_bg = mean_abs(affine_fwd @ affine_rev - eye(3))
    # Warped encoder features of source vs driving features  [eq. 13]
    loss_warp = warp_match(source, driving, flow)
    # Unweighted sum  [eq. 14]
    return loss_rec + loss_eq + loss_bg + loss_warp

Hyperparameters

Values come from the paper where it states them, and from the released training configs otherwise. The paper states only K, the epoch count, the GPU setup and the training duration. Config values are for the TaiChiHD config unless a row says otherwise.

Results

Author-reported. The setup is reconstruction on held-out videos, with all methods trained for the same number of epochs at K = 10 §4.2 tab. 1. Lower is better for L1, AKD, MKR and AED. Bold marks the paper's best value per column.

"—" marks metrics that tab. 1 does not report for that dataset: MKR for VoxCeleb, and everything except L1 for MGif. On TaiChiHD the text reports relative gains of 15.5% for AKD and 28.0% for MKR over MRAA §4.2. These follow from the table values (5.41 → 4.57 and 0.025 → 0.018).

Comparison with K matched by dimension tab. 3, TaiChiHD, L1 / (AKD, MKR) / AED:

User study tab. 2 §4.2: 20 randomly chosen video pairs from TaiChiHD, TED-talks and VoxCeleb. Share of users preferring the proposed method over MRAA.

Ablation on TaiChiHD tab. 4, added in turn: TPS replaces the affine motion; then dropout; then multi-resolution masks; then L_bg and L_warp.

The paper's own reading: multi-resolution masks help L1 and AED but raise AKD, and the two auxiliary losses bring AKD back down §4.3.

Image animation is shown only as qualitative figures against MRAA fig. 6. No quantitative animation table exists.

Paper vs code

Repository: yoyo-nb/Thin-Plate-Spline-Motion-Model [code], MIT licence, linked from the paper's abstract footnote. Checked at code_commit. Every item below is a difference between the paper text and the code at that commit.

• Direction of the TPS constraint. The paper writes the constraint as T(P^S_i) = P^D_i eq. 1. The code fits the TPS from driving keypoints to source keypoints, so it maps driving coordinates to source coordinates, and that map is used directly as a sampling grid for the source image util.py dense_motion.py dense_motion.py.
• Ridge term in the TPS solve. The code adds 0.01 times the identity to the linear system before inverting it util.py. The paper's constraint is exact interpolation. Also, the radial term is computed as r² log(r² + 1e-9), with an epsilon inside the log util.py util.py.
• Resolution of the flow and the masks. The paper defines the contribution maps on the full H×W grid eq. 5. The code runs the dense motion network on a 0.25-downscaled source, so its maps and flow are at 1/4 resolution dense_motion.py dense_motion.py taichi-256.yaml. The flow is bilinearly resized to each feature size in the inpainting network inpainting_network.py.
• Dropout rescaling. The code computes the exponential first, then divides the logits by (1 − P), so the division never reaches the exponentials that the partition sums. The kept terms are therefore not rescaled as eq. 7 requires dense_motion.py. The background channel is always kept dense_motion.py.
• Dropout sampling. The paper speaks of K fields "in this mini-batch" §3.1. The code draws an independent keep/drop flag for each image and each field, and shares it across pixels dense_motion.py.
• Dropout schedule. The paper gives no value for P and says dropout is removed "after training several epochs" §3.1. The code ramps P linearly and switches it off at epoch 35 model.py taichi-256.yaml.
• Loss weights and reduction. The paper's total loss is an unweighted sum with absolute-value sums eq. 9-14. The code multiplies each term by 10 and takes the mean of the absolute values over all elements taichi-256.yaml model.py model.py model.py model.py. The four weights are 10 for perceptual, equivariance, warp and background.
• Equivariance input. The paper's equivariance loss uses the source image S eq. 10. The code applies the random TPS warp to the driving frame model.py.
• Pyramid scales. The paper says the reconstruction loss uses images downsampled j times, without giving scales eq. 9. The code uses scales 1, 0.5, 0.25 and 0.125 taichi-256.yaml.
• Background motion network. The paper does not name the architecture. The code uses a ResNet-18 with six input channels, initialised so that the output is the identity bg_motion_predictor.py.
• Keypoint network. The paper does not name the architecture. The code uses a ResNet-18 regression head with a sigmoid mapped to [−1, 1] keypoint_detector.py.
• Dense motion inputs. The paper says the dense network receives the K+1 warped source images §3.1. The code also concatenates a heatmap of the difference between Gaussian keypoint maps (variance 0.01) dense_motion.py dense_motion.py.
• Output blend. The paper's inpainting output comes from the last layer §3.2. The code blends the sigmoid output with the warped source image, weighted by the last occlusion mask inpainting_network.py.
• Gradient stopping. The code detaches encoder features and occlusion masks in the warp-loss path inpainting_network.py inpainting_network.py inpainting_network.py. The paper does not mention this. As a result, the warp loss appears to change only the flow.
• Animation mode of the demo. The paper says it uses the animation via disentanglement mode §3.4. The demo's default is "relative" mode demo.py.
• Animation-stage keypoint noise. The paper says keypoints are "randomly deformed". The code scales each keypoint's x and y independently by a uniform factor in [0.75, 1.25] train_avd.py. The paper does not give the range.
• Augmentation. The code applies horizontal flip, time flip and colour jitter during training taichi-256.yaml. The paper does not mention augmentation.
• Frame sampling. The paper says two frames are sampled from a video §1. The code draws two sorted frame indices with replacement from the same video frames_dataset.py.

Checked and consistent: the TPS kernel r² log r² eq. 3 util.py; the softmax over K+1 maps eq. 5 dense_motion.py; the flow as the weighted sum of transforms eq. 6 dense_motion.py; K = 10 and N = 5 taichi-256.yaml; and the form of the background consistency term eq. 12 model.py.

Pitfalls

• Dimension formula contradicts its own number. The paper states the motion dimension as K·(6 + 5·2) + 6, which gives 38, 70 and 134 for K = 2, 4 and 8. Its text gives 118 for K = 8 §4.3. The formula gives 134. The comparison in tab. 3 depends on this count, so check which one holds before quoting it.
• Dropout rescale. Eq. 7 divides the kept terms by 1−P. The released code does not apply that division to the exponentials, so a faithful reimplementation of the text and a copy of the code will train differently. Pick one and say which one you chose.
• Loss scale. The paper writes absolute-value sums with an unweighted total eq. 14. The released code takes means, weights each term by 10, and clips gradients at 10 train.py. Following the text literally therefore gives a different loss scale from the released run.
• Hidden hyperparameters. The paper leaves out learning rate, optimiser, batch size, dropout probability, pyramid scales, the resolution of each occlusion mask, and the architecture of the three modules. Each value in the table above that says "not stated" must be taken from the code or chosen by the implementer.
• Dataset-specific settings. Background start, dropout ramp and batch size differ between datasets in the released configs vox-256.yaml ted-384.yaml mgif-256.yaml. MGif turns the background off mgif-256.yaml.
• Ridge and epsilon are numerically sensitive. The 0.01 ridge changes the fitted warp, and the 1e-9 epsilon inside the log keeps r² log r² finite at zero. Removing either changes the flow.
• Two-stage training. The animation stage needs a trained generator checkpoint and raises an error without one train_avd.py. It is a separate training run, with its own epochs and learning rate.
• Extreme identity mismatch. The paper says its method does not work well when the identity mismatch is extreme (for example, a puppy driven by a giraffe) §4.2 §5.
• Cost. The TPS fit inverts an (N+3)×(N+3) system per group and per sample, and the paper reports 8 days of training on TED-talks §4.2. The paper gives no per-frame inference time.
• Human judgements are thin. The user study uses 20 video pairs. It does not report the number of raters, the video selection procedure or a significance test §4.2 tab. 2.
• Qualitative claims. The paper's claims of fewer ghosts on VoxCeleb and better temporal continuity rest on figures and the user study. They do not have a quantitative table of their own fig. 4 fig. 5 fig. 6 §4.2.
• Metric definitions. AKD and MKR use pretrained pose detectors, and AED uses pretrained re-identification networks, all taken from MRAA. The paper cites these detectors by reference only and gives no versions, so exact reproduction of the numbers depends on MRAA's evaluation setup §4.2.

Lineage

• Builds on: first-order motion model (FOMM) and MRAA, both cited in §2. This paper replaces their local affine motion with TPS fields and keeps MRAA's background transform and its animation mode §2 §3.1 §3.4. The paper gives no arXiv ids for them.
• Compared against: X2Face, Monkey-Net, FOMM and MRAA tab. 1.
• Superseded by: not known from this card.