魔性！Python生成全网爆火的“蚂蚁呀嘿”

Posted 2022-06-16 走召大爷

前段时间火爆的“蚂蚁呀嘿”，将一个人的说话动作和表情迁移到另一张静态图中，让静态图中的人脸做出指定动作表情，主要基于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 源码获取

1. 关注公众号：Python学习实战
2. 公众号聊天界面回复：表情迁移，获取完整源码。

如果您觉得本文有帮助，辛苦您点个不需花钱的赞，您的举手之劳将对我提供了无限的写作动力！ 也欢迎关注我的公众号：Python学习实战， 第一时间获取最新文章。