トップ > Computer Vision > Stable Diffusionを使った顔動画の自由なスタイル変換

Stable Diffusionを使った顔動画の自由なスタイル変換

Computer Vision Style Transfer Image Generation

この記事では、Stable Diffusionによるスタイル変換とThin Plate Spline Motion Modelによるモーショントレースを組み合わせて、人の顔の動画を簡単にスタイル変換する方法を紹介します。

概要

Stable DIffusionのimage-to-imageは、任意のプロンプトに合わせて元画像をハイクオリティに変換することができます。ここでは、これを利用して、好きな動画の1フレーム目だけをimage-to-imageで変換後、元動画のモーションをThin Plate Spline Motion Modelでトレースすることで、動画全体を様々なスタイルに変換してみます。

image-to-image、Thin Plate Spline Motion Modelの詳細は、それぞれ以下の記事をご覧ください。

デモ

それでは、Google Colaboratoryを使って実際に実行していきます。

なお、記事内で紹介したコードをすべて含むノートブックは、以下のリンクから直接参照することができます。

Open In Colab

環境設定

はじめに、画面上部のメニューから、「ランタイム」、「ランタイムのタイプを変更」と進み、「ハードウェアアクセラレータ」を「GPU」に変更しておきます。これにより、推論時にGPUを利用することができます。

※ 2022年10月よりGoogle Colaboratoryの料金プランがこちらの内容に変更されました。ご利用の際はご注意ください。

次に、必要なライブラリのインストール、インポート、リポジトリのクローン、関数の定義を行います。

# ライブラリのインストール
!pip install -qq diffusers==0.3.0 transformers scipy ftfy
!pip install -qq imageio-ffmpeg

# リポジトリのクローン
!git clone https://github.com/yoyo-nb/Thin-Plate-Spline-Motion-Model.git

# ライブラリのインポート
import os
import shutil
import warnings
import numpy as np
from PIL import Image
import imageio
from skimage import img_as_ubyte
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from skimage.transform import resize
from IPython.display import HTML
import cv2
import torch
from torch import autocast
from diffusers import StableDiffusionImg2ImgPipeline
warnings.filterwarnings('ignore')

# 描画関数の定義
def image_grid(imgs, rows, cols):
    assert len(imgs) == rows * cols

    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols * w, rows * h))
    grid_w, grid_h = grid.size
    
    for i, img in enumerate(imgs):
        grid.paste(img, box=(i % cols * w, i // cols * h))

    return grid
    
# 前処理関数の定義
def preprocess(image):
    w, h = image.size
    w, h = map(lambda x: x - x % 32, (w, h))
    image = image.resize((w, h), resample=Image.LANCZOS)
    image = np.array(image).astype(np.float32) / 255.0
    image = image[None].transpose(0, 3, 1, 2)
    image = torch.from_numpy(image)

    return 2. * image - 1.

続いて、Stable Diffusionの学習済みモデルを利用するために、Hugging Faceにログインしておきます。以下のセルを実行後、表示される入力欄にアクセストークンを入力します。アクセストークンの取得方法はこちらの記事をご確認ください。

from huggingface_hub import notebook_login
notebook_login()

対象動画の準備

対象の動画をアップロードします。以下のコードでは正方形に近いものを想定しています。

uploaded = files.upload()

動画をフレームに分割します。

# 動画ファイルパスの取得
video_path = os.path.join('/content', list(uploaded.keys())[0])

# 保存ディレクトリの作成
frame_dir = '/content/frames/'
if os.path.exists(frame_dir):
    shutil.rmtree(frame_dir)
if not os.path.exists(frame_dir):
    os.makedirs(frame_dir)    

# 動画のフレーム分割
cmd = "ffmpeg -i %s -start_number 0 -vsync 0 %s/%%03d.png" % (
            video_path,
            frame_dir,
        )
os.system(cmd)

今回の対象動画の1フレーム目はこちら。この記事よりお借りしています。

これで動画の準備は完了です。

Stable Diffusionによるimage-to-image

では、image-to-imageを実施します。ここではアニメ調に変換してみます。

# プロンプトの設定
prompt = 'Portrait of a female, defined facial features, highly detailed, animation cel, official Kyoto Animation and Studio Ghibli anime screenshot, by Makoto Shinkai and Range Murata'

# 学習済みモデルのダウンロード
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    'CompVis/stable-diffusion-v1-4',
    revision='fp16', 
    torch_dtype=torch.float16,
    use_auth_token=True
).to('cuda')

# 出力ディレクトリの作成
output_dir = '/content/output'
!mkdir -p $output_dir
if os.path.exists(output_dir):
    shutil.rmtree(output_dir)
if not os.path.exists(output_dir):
    os.makedirs(output_dir)   

# 生成画像数の設定
num = 30

# 初期画像の読込
image_file = '/content/frames/000.png'
original_image = Image.open(image_file).convert('RGB')
resized_image = original_image.resize((512, 512))
init_image = preprocess(resized_image)
plt.imshow(resized_image)
plt.axis('off')
plt.title('before processing')

# 画像の生成
images = []
for i in range(num):
    with autocast('cuda'):
        image = pipe(prompt=prompt, init_image=init_image, strength=0.7)['images'][0]
        image.save(os.path.join(output_dir, f'{i:03d}.png'))
        images.append(image)

# 結果の描画
image_grid(images, 6, 5)

上記のコードでは30枚を生成しています。枚数は求める完成度に応じて増減してください。

ここでは、以下の1枚を採用しました。

# 画像番号の指定
image_idx = '009'  # 選択した画像の通番

# 選択画像の読込
image_file = os.path.join(output_dir, f'{image_idx}.png')
image = Image.open(image_file).convert('RGB')
plt.imshow(image)
plt.axis('off')
plt.title('selected image');

Thin Plate Spline Motion Modelによるモーショントレース

変換した1フレームに対して、元動画からモーションをコピーします。

# パッケージのインポート
%cd Thin-Plate-Spline-Motion-Model
from demo import load_checkpoints
from demo import make_animation

# パラメータ設定
device = torch.device('cuda:0')
dataset_name = 'vox'
source_image_path = image_file
driving_video_path = video_path
output_video_path = '/content/result.mp4'
config_path = 'config/vox-256.yaml'
checkpoint_path = 'checkpoints/vox.pth.tar'
predict_mode = 'relative'
find_best_frame = False
pixel = 256

# トレース先の画像・トレース元の動画の読込
source_image = imageio.imread(source_image_path)
reader = imageio.get_reader(driving_video_path)
fps = reader.get_meta_data()['fps']
source_image = resize(source_image, (pixel, pixel))[..., :3]

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

driving_video = [resize(frame, (pixel, pixel))[..., :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

# 学習済みパラメータのダウンロード
!mkdir -p checkpoints
!wget -c https://cloud.tsinghua.edu.cn/f/da8d61d012014b12a9e4/?dl=1 -O checkpoints/vox.pth.tar

# 学習済みパラメータのロード
inpainting, kp_detector, dense_motion_network, avd_network = load_checkpoints(config_path = config_path, checkpoint_path = checkpoint_path, device = device)

# トレースの実行
predictions = make_animation(source_image, driving_video, inpainting, kp_detector, dense_motion_network, avd_network, device = device, mode = predict_mode)

# 結果の保存
result_dir = '/content/result'
if os.path.exists(result_dir):
    shutil.rmtree(result_dir)
if not os.path.exists(result_dir):
    os.makedirs(result_dir)   
for i, frame in enumerate(predictions):
    cv2.imwrite(os.path.join(result_dir, f'{i:03d}.png'), img_as_ubyte(frame)[:, :, ::-1])
imageio.mimsave(output_video_path, [img_as_ubyte(frame) for frame in predictions], fps=fps)

# 描画
HTML(display(source_image, driving_video, predictions).to_html5_video())

結果は、フレームごとに/content/resultに保存されます。

別途実行した、ヒューマノイド風、少女風も含めた最終的な結果がこちらです。

まとめ

動画スタイル変換の数ある手法のうちの一つですが、1枚目のみをスタイル変換の基準とするため、Stable Diffusionの結果を厳選しやすい点、全体的に高速である点でオススメです。リンク先のcolabファイルでは、動画をアップロードするだけで上記のすべての処理を再現可能なので、ぜひ試してみてください。

参考文献

High-Resolution Image Synthesis with Latent Diffusion Models

GitHub - yoyo-nb/Thin-Plate-Spline-Motion-Model: [CVPR 2022] Thin-Plate Spline Motion Model for Image Animation.