前段时间火爆的“蚂蚁呀嘿”，将一个人的说话动作和表情迁移到另一张静态图中，让静态图中的人脸做出指定动作表情，主要基于FOMM（First Order Motion model）技术。这已经是2年前的技术了，在一些场景中生成的效果并不理想。近期，清华大学团队在CVPR2022发布最新表情动作迁移论文Thin-Plate Spline Motion Model for Image Animation。本文不具体讲论文原理，而是直接将其开源的模型down下来用。效果如下：
第1张图是静态照片，第二张是gif驱动动画，第三张是生成的结果。
本文目的：将开源模型打包封装成一个单独接口，读者只需简单传入一张图片和一个动画（gif或mp4），即可生成表情迁移动画（mp4）。
0 环境准备
读者需要安装好pytorch环境，可前往https://pytorch.org/get-started/locally/根据实际的硬件环境，选择GPU或cpu版本。
安装imageio-ffmpeg库，用于读取mp4文件。
读者可以直接跳到最后，获取源码，下载源码资源包后，将文件夹中的对应图片替换成自己的图片，即可一键生成“蚂蚁呀嘿”。
1 模型封装
将模型导出为pt后，再创建Model类，将接口封装到infer函数， 具体代码如下：
class Model():
def __init__(self, kp="models/kp.pt", aio="models/aio.pt",
device=torch.device('cpu')):
self.device = device
self.kp = torch.jit.load(kp, map_location=device).eval()
self.aio = torch.jit.load(aio, map_location=device).eval()
def relative_kp(self, kp_source, kp_driving, kp_driving_initial):
source_area = ConvexHull(kp_source[0].data.cpu().numpy()).volume
driving_area = ConvexHull(
kp_driving_initial[0].data.cpu().numpy()).volume
adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
kp_new = kp_driving
kp_value_diff = (kp_driving - kp_driving_initial)
kp_value_diff *= adapt_movement_scale
kp_new = kp_value_diff + kp_source
return kp_new
def get_kp(self, src):
src = np.expand_dims(src, 0).transpose(0, 3, 1, 2)
src = torch.from_numpy(src).float().to(self.device)
return self.kp(src)
def infer(self, src, driving, src_kp, init_kp):
src = np.expand_dims(src, 0).transpose(0, 3, 1, 2)
src = torch.from_numpy(src).float().to(self.device)
driving = np.expand_dims(driving, 0).transpose(0, 3, 1, 2)
driving = torch.from_numpy(driving).float().to(self.device)
kp_driving = self.kp(driving)
kp_norm = self.relative_kp(kp_source=src_kp,
kp_driving=kp_driving,
kp_driving_initial=init_kp)
with torch.no_grad():
out = self.aio(src, src_kp, kp_norm)
out = out[0].cpu().numpy()
out = out.transpose(1, 2, 0)
return out
其中，get_kp函数用于获取脸部关键点数据。infer函数中，src表示静态图，driving表示动态图中的某一帧，src_kp表示静态图的关键点，init_kp表示动态图中的第一帧的关键点。
2 调用模型
整个调用流程可以拆分为4步：创建模型对象、读取动图的每一帧、调用模型、生成帧导出mp4。
2.1 创建模型对象
前面定义好了Model对象，需要根据GPU和CPU环境，由读者指定使用具体的pytorch版本，具体代码如下所示。
def create_model(use_gpu):
if use_gpu:
device = torch.device('cuda')
else:
device = torch.device('cpu')
model = Model(device=device)
return model
上面代码中，use_gpu是个boolean类型，用于判断是否使用GPU版本，读者根据自己的时间情况设置。
2.2 读取动图的每一帧
调用imageio-ffmpeg库，读取mp4或gif文件中的每一帧。具体代码如下所示，函数返回列表，列表内容为视频帧：
def read_mp4_or_gif(path):
reader = imageio.get_reader(path)
if path.lower().endswith('.mp4'):
fps = reader.get_meta_data().get('fps')
elif path.lower().endswith('.gif'):
fps = 1000 / Image.open(path).info['duration']
driving_video = []
try:
for im in reader:
im = resize(im, (256, 256))[..., :3]
driving_video.append(im)
except RuntimeError:
pass
reader.close()
return driving_video, fps
因为模型的约束，这里将每一帧resize到256*256。
2.3 调用模型
模型调用非常简单，只需读取静态图和动态图的每一帧，并对静态图和动态图第一帧调用Model类的get_kp函数获取关键点即可。遍历动态图的每一帧，将动图帧、静态图、静态图关键点、动图第一帧关键点一起传给Model的infer函数即可得到生成的帧。具体代码如下所示。
def run(use_gpu, src_path, driving_path):
src = imageio.imread(src_path)
src = resize(src, (256, 256))[..., :3]
driving_video, fps = read_mp4_or_gif(driving_path)
model = create_model(use_gpu)
src_kp = model.get_kp(src)
init_kp = model.get_kp(driving_video[0])
outs = []
for driving in driving_video:
out = model.infer(src, driving, src_kp, init_kp)
out = img_as_ubyte(out)
outs.append(out)
return outs, fps
2.4 生成帧导出mp4
这里继续调用imageio-ffmpeg库，将视频帧组装成mp4文件，代码如下：
def write_mp4(out_path, frames, fps):
imageio.mimsave(out_path, frames, fps=fps)
2.5 pipeline调用
整个调用流水线如下：
src_path = 'assets/source.png'
driving_path = 'assets/driving2.gif'
frames, fps = run(True, src_path, driving_path)
write_mp4("out.mp4", frames, fps)
3 源码获取
- 关注公众号：
Python学习实战 - 公众号聊天界面回复：
表情迁移，获取完整源码。
如果您觉得本文有帮助，辛苦您点个不需花钱的赞，您的举手之劳将对我提供了无限的写作动力！ 也欢迎关注我的公众号：Python学习实战， 第一时间获取最新文章。
