「 First Order Motion Model for Image Animation 代码阅读 」

FOMM论文阅读笔记

代码仓库

demo .py

main

# read source_iamge and driving_frame
source_image = imageio.imread(opt.source_image)
reader = imageio.get_reader(opt.driving_video)
fps = reader.get_meta_data()['fps']
driving_video = []
try:
	for im in reader:
		driving_video.append(im)
except RuntimeError:
	pass
reader.close()
source_image = resize(source_image, (256, 256))[..., :3]
driving_video = [resize(frame, (256, 256))[..., :3] for frame in driving_video]

首先读取source_image和driving_video，然后把它们分别resize成 (256, 256) ，这里需要注意，input的图像和视频最好都是正方形的，不然resize时不可避免地会出现形变，影响最终效果。

generator, kp_detector = load_checkpoints(config_path=opt.config, checkpoint_path=opt.checkpoint, cpu=opt.cpu)

读取模型。

# FOMM要求S1和D1的pose大体相同，才能获得最好的效果
if opt.find_best_frame or opt.best_frame is not None:
	i = opt.best_frame if opt.best_frame is not None else find_best_frame(source_image, driving_video, cpu=opt.cpu)
	print ("Best frame: " + str(i))
	driving_forward = driving_video[i:] #best_frame之后的，把best_frame作为D1
	driving_backward = driving_video[:(i+1)][::-1] #同理，把best_frame作为D1
	predictions_forward = make_animation(source_image, driving_forward, generator, kp_detector, relative=opt.relative, adapt_movement_scale=opt.adapt_scale, cpu=opt.cpu)
	predictions_backward = make_animation(source_image, driving_backward, generator, kp_detector, relative=opt.relative, adapt_movement_scale=opt.adapt_scale, cpu=opt.cpu)
	predictions = predictions_backward[::-1] + predictions_forward[1:]
else:
	predictions = make_animation(source_image, driving_video, generator, kp_detector, relative=opt.relative, adapt_movement_scale=opt.adapt_scale, cpu=opt.cpu)

由于FOMM生成动画时是把 $T_{D_t \gets D_1}(p)$ 施加到 $S_1$ 上，所以就要求S1和D1的pose大体相同，但如果实在不同，可以使用 --find_best_frame 参数让程序自动选出最优的一帧作为D1，或者使用 --best_frame int 来手动指定D1帧。

如果不使用best_frame的相关参数，就直接进入 make_animation() 进行视频帧的预测，这个函数返回一个每帧的list。

如果使用了相关参数，选出第i帧为D1帧，那么这时就会分为两步进行预测，分别是 forward = make_animation(driving_video = driving_video[i:]) 和 backward = make_animation(driving_video = driving_video[:(i+1)][::-1]) ，其实就是把i作为第一帧，预测从它开始到结尾；以及预测从它作为第一帧返回到开头。最后把后者倒序一下，再把两者连接一下。

make_animation()

kp_source = kp_detector(source)
kp_driving_initial = kp_detector(driving[:, :, 0]) #D1

首先通过kp_detector获取相关帧的kp，这里获取kp的参数是在 load_checkpoints() 就通过 /config/{model_name}.yaml 中读取的。比如对于人脸的vox-256.yaml中就有相关定义：

model_params:
  common_params:
    num_kp: 10
    num_channels: 3
    estimate_jacobian: True
  kp_detector_params:
     temperature: 0.1
     block_expansion: 32
     max_features: 1024
     scale_factor: 0.25
     num_blocks: 5

所以此时我们对于人脸工作，得到的kp就是一个 {'value': (1, 10, 2), 'jacobian': (1, 10, 2, 2)} 的dict，value所对应的就是预测出来的关键点，其坐标分布在 [(-1, 1), (-1, 1)] 的范围，jacobian是对应的雅可比矩阵。

for frame_idx in tqdm(range(driving.shape[2])):
	driving_frame = driving[:, :, frame_idx]
	if not cpu:
		driving_frame = driving_frame.cuda()
	kp_driving = kp_detector(driving_frame)
	kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving, kp_driving_initial=kp_driving_initial, use_relative_movement=relative, use_relative_jacobian=relative, adapt_movement_scale=adapt_movement_scale)
	out = generator(source, kp_source=kp_source, kp_driving=kp_norm)
	predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])

接下来对于D的每一帧处理，先获得其对应的kp，然后进行一个标准化 normalize_kp() ，这个其实就是对 --relative --adapt_scale 等参数的实现，来决定是否要让target图像与source保持相同的缩放、是否进行相对运动等参数。默认都是false的，所以一般可以不用考虑。

最后将source、kp_source、kp_driving丢进generator，进行预测，然后由于前面对图像矩阵进行了置换，这里将其置换为真实图像，然后返回，大体框架到这里结束。

dense_motion .py

在generator中第一步就是预测图片的dense motion，所以就先从这部分开始看起。

forward()

heatmap_representation = self.create_heatmap_representations(source_image, kp_driving, kp_source)
sparse_motion = self.create_sparse_motions(source_image, kp_driving, kp_source)
deformed_source = self.create_deformed_source_image(source_image, sparse_motion)
...
if self.occlusion:
	occlusion_map = torch.sigmoid(self.occlusion(prediction))

暂时略过简单的变换部分，这个模型的核心就是这四个函数。 create_heatmap_representations() 对应的是论文中的Eq 6. $H_k(z)=exp(\frac{(T_{D \gets R}(p_k)-z)^2}{\sigma})-exp(\frac{(T_{S \gets R}(p_k)-z)^2}{\sigma})$ ， create_sparse_motions() 对应Eq 4. $T_{S \gets D}(z)=T_{S \gets R}(p_k)+J_k(z-T_{D \gets R}(p_k))$ ， create_deformed_source_image() 对应 $\hat T_{S \gets D}(z)=M_0 z+\displaystyle\sum^K_{k=1} M_k(T_{S \gets R}(p_k)+J_k(z-T_{D \gets R}(p_k)))$ ，occlusion() 则是论文中Sec. 3.2用mask $\hat O$ 和变换 $\hat T$ 来预测最终图像的部分。

create_sparse_motions()

bs, _, h, w = source_image.shape
identity_grid = make_coordinate_grid((h, w), type=kp_source['value'].type())
identity_grid = identity_grid.view(1, 1, h, w, 2)
coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 2)
if 'jacobian' in kp_driving:
	jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian']))
	jacobian = jacobian.unsqueeze(-3).unsqueeze(-3)
	jacobian = jacobian.repeat(1, 1, h, w, 1, 1)
	coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
	coordinate_grid = coordinate_grid.squeeze(-1)

driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 2)

#adding background feature
identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1)
return sparse_motions

$\begin{array}{cc} T_{S \gets D}(z)=T_{S \gets R}\circ T_{D \gets R}^{-1}\approx T_{S \gets R}(p_k)+J_k(z-T_{D \gets R}(p_k)) \\ J_k \approx (\frac{d}{dp}T_{S \gets R}(p)|_{p=p_k})(\frac{d}{dp}T_{D \gets R}(p)|_{p=p_k})^{-1} \end{array}$

首先 util.py > make_coordinate_grid() 返回一张与source_image相同比例的 ([-1, 1], [-1, 1]) 图片，用来映射kp（这与kp的坐标范围相同）。

coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 2) 对应公式中的 $(z-T_{D \gets R}(p_k))$ ，从这里就可以看出，KPDetector预测出的关键点其实已经是 $T_{X \gets R}$ 的形式了。

jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian'])) 计算出 $J_k$ ，然后和上一步算出的 $(z-T_{D \gets R}(p_k))$ 相乘。

driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 2) 最后计算 $T_{S \gets D}(z)=T_{S \gets R}(p_k)+J_k(z-T_{D \gets R}(p_k))$ ，然后把 $T_{S \gets D}$ 和 identity_grid 一起返回。

然后就是把返回值丢到 create_deformed_source_image 中进行对齐，然后与heatmap一起丢进Hourglass中训练，得到 $M_k$ ，然后预测出mask，最后通过occlusion预测出图像。