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

Posted 走召大爷

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了魔性!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对象,需要根据GPUCPU环境,由读者指定使用具体的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库,读取mp4gif文件中的每一帧。具体代码如下所示,函数返回列表,列表内容为视频帧:

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

因为模型的约束,这里将每一帧resize256*256

2.3 调用模型

模型调用非常简单,只需读取静态图和动态图的每一帧,并对静态图和动态图第一帧调用Model类的get_kp函数获取关键点即可。遍历动态图的每一帧,将动图帧、静态图、静态图关键点、动图第一帧关键点一起传给Modelinfer函数即可得到生成的帧。具体代码如下所示。

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学习实战, 第一时间获取最新文章。

以上是关于魔性!Python生成全网爆火的“蚂蚁呀嘿”的主要内容,如果未能解决你的问题,请参考以下文章

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

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

使用PaddlePaddle实现火爆全球的“蚂蚁呀嘿”视频

曾被网友疯狂恶搞的「蚂蚁呀嘿」项目开源上过GitHub热榜

从刷屏的“蚂蚁呀嘿”到“蚂呀完了”,谁在挖坑谁在跳?

Python爆火的六大原因