魔性！Python生成全网爆火的“蚂蚁呀嘿”
前段时间火爆的“蚂蚁呀嘿”，将一个人的说话动作和表情迁移到另一张静态图中，让静态图中的人脸做出指定动作表情，主要基于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学习实战 - 公众号聊天界面回复：
表情迁移，获取完整源码。
