!git clone https://github.com/AliaksandrSiarohin/first-order-model
# or  
!git clone https://codehub.devcloud.cn-north-4.huaweicloud.com/ai-pome-free00001/first-order-model.git

# 此处牛刀小试--用 Moxing 下载文件
import moxing as mox
# 此处需要替换您的 OBS 地址
mox.file.copy_parallel('obs://modelarts-lab/first-order-motion-model/first-order-motion-model-20210226T075740Z-001.zip' , 'first-order-motion-model.zip')
mox.file.copy_parallel('obs://modelarts-lab/first-order-motion-model/02.mp4' , '02.mp4')

# 解压
!unzip first-order-motion-model.zip  
# 模版视频
!mv 02.mp4 first-order-motion-model/

cd first-order-model
import imageio
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from skimage.transform import resize
from IPython.display import HTML
import warnings
warnings.filterwarnings("ignore")

# 此处替换为您的图片路径，图片最好为 256*256，这里默认为普京大帝
source_image_path = '/home/ma-user/work/first-order-motion-model/02.png'
source_image = imageio.imread(source_image_path)

# 此处可替换为您的视频路径，这里默认为“蚂蚁牙黑”
reader_path = '/home/ma-user/work/first-order-motion-model/02.mp4'
reader = imageio.get_reader(reader_path)
# 调整图片和视频大小为 256x256

source_image = resize(source_image, (256, 256))[..., :3]

fps = reader.get_meta_data()['fps']
driving_video = []
try:
    for im in reader:
        driving_video.append(im)
except RuntimeError:
    pass
reader.close()

driving_video = [resize(frame, (256, 256))[..., :3] for frame in driving_video]

def display(source, driving, generated=None):
    fig = plt.figure(figsize=(8 + 4 * (generated is not None), 6))

    ims = []
    for i in range(len(driving)):
        cols = [source]
        cols.append(driving[i])
        if generated is not None:
            cols.append(generated[i])
        im = plt.imshow(np.concatenate(cols, axis=1), animated=True)
        plt.axis('off')
        ims.append([im])

    ani = animation.ArtistAnimation(fig, ims, interval=50, repeat_delay=1000)
    plt.close()
    return ani
HTML(display(source_image, driving_video).to_html5_video())

from demo import load_checkpoints
generator, kp_detector = load_checkpoints(config_path='config/vox-256.yaml', 
                            checkpoint_path='/home/ma-user/work/first-order-motion-model/vox-cpk.pth.tar')

from demo import make_animation
from skimage import img_as_ubyte

predictions = make_animation(source_image, driving_video, generator, kp_detector, relative=True)

# 保存结果视频
imageio.mimsave('../generated.mp4', [img_as_ubyte(frame) for frame in predictions], fps=fps)
# 在 Notebook 根目录能找，/home/ma-user/work/

HTML(display(source_image, driving_video, predictions).to_html5_video())

# 安装视频剪辑神器 moviepy
!pip install moviepy

# 为生成的视频加上源视频声音
from moviepy.editor import *

videoclip_1 = VideoFileClip(reader_path)
videoclip_2 = VideoFileClip("../generated.mp4")

audio_1 = videoclip_1.audio

videoclip_3 = videoclip_2.set_audio(audio_1)

videoclip_3.write_videofile("../result.mp4", audio_codec="aac")

# 还可以给视频加水印
video = VideoFileClip("../result.mp4")
# 水印图片请自行上传
logo = (ImageClip("/home/ma-user/work/first-order-motion-model/water.png")
        .set_duration(video.duration) # 水印持续时间
        .resize(height=50) # 水印的高度，会等比缩放
        .margin(right=0, top=0, opacity=1) # 水印边距和透明度
        .set_pos(("left","top"))) # 水印的位置

final = CompositeVideoClip([video, logo])
final.write_videofile("../result_water.mp4", audio_codec="aac")

final_reader = imageio.get_reader("../result_water.mp4")

fps = final_reader.get_meta_data()['fps']
result_water_video = []
try:
    for im in final_reader:
        result_water_video.append(im)
except RuntimeError:
    pass
reader.close()
result_water_video = [resize(frame, (256, 256))[..., :3] for frame in result_water_video]
HTML(display(source_image, driving_video, result_water_video).to_html5_video())