Talk-to-Editと呼ばれる機械学習手法を使って一枚の顔写真から笑顔の写真やひげの生えた写真など様々な顔写真を生成する方法をご紹介します。

Talk-to-Editとは

Talk-to-Editは対話形式できめの細かい顔の属性操作を実現する顔編集フレームワークです。

従来技術では、ほくそ笑むような表情から大笑いへの変化のような、連続的できめの細かい顔の編集ができませんでした。Talk-to-Editはこの問題をGAN潜在空間(GAN latent space)における連続的な「セマンティックフィールド」をモデル化することによって解決しています。

出典: https://github.com/yumingj/Talk-to-Edit

大まかな理解としては、対話形式で顔編集できる技術です。

デモ(Colaboratory)

それでは早速、Google Colaboratoryで動かしていきます。

なお、これから紹介するソースコードは全てこちらのGitHubに掲載しております。以下のボタンをクリックするとColaboratoryを開くことも可能です。

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

まず、Pytorchや、その他ライブラリをインストールします。

# ライブラリダウンロード
!pip install ninja facenet-pytorch scikit-image==0.15.*
!pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 -f https://download.pytorch.org/whl/torch_stable.html 

# GitHubからコードを取得
!git clone https://github.com/yumingj/Talk-to-Edit.git

# 学習済みモデルのダウンロード
%cd /content/Talk-to-Edit/download
!wget -c https://www.dropbox.com/s/wlq8plp2qg3sw7i/pretrained_models.zip?dl=0 -O pretrained_models.zip
!unzip pretrained_models.zip
!rm -rf pretrained_models.zip

この時、scikit-imageのバージョンにはご注意ください。

models/archs/stylegan2/lpips/__init__.pyにてfrom skimage.measure import compare_ssimを使用している箇所があり、compare_ssimがバージョン0.16以降で別の関数に置き換えられているためエラーになります。
scikit-imageの最新バージョンを使用したい場合は、ソースコードを変更してください。

学習済みモデルのダウンロード

続いて学習モデルをダウンロードします。

# 学習済みモデルのダウンロード 
%cd /content/Talk-to-Edit/download/pretrained_models/
!wget -c http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
!bzip2 -d shape_predictor_68_face_landmarks.dat.bz2

顔編集を行う画像のアップロード

次に、顔編集を行う画像のアップロードを行います。
本記事では、JAPAN RUGBY掲載の稲垣啓太選手のプロフィール写真を使用させていただきます。

%cd /content/Talk-to-Edit/download/

# 画像アップロード
from google.colab import files
uploaded = files.upload()
uploaded = list(uploaded.keys())
print(uploaded)
file_name = uploaded[0]

モデルのセットアップ

Configファイルを読込みモデルをセットアップします。

%cd /content/Talk-to-Edit

# Config設定
import argparse
import json
import logging
import os.path
import random

import numpy as np
import torch
import torchvision

from models import create_model
from models.utils import save_image
from utils.editing_utils import edit_target_attribute
from utils.logger import get_root_logger
from utils.options import (dict2str, dict_to_nonedict, parse,
                           parse_opt_wrt_resolution)
from utils.util import make_exp_dirs
from utils.inversion_utils import inversion


from PIL import Image

opt = './configs/editing/editing_wo_dialog.yml'
args = argparse.ArgumentParser(description='')
opt = parse(opt, is_train=False)
opt['img_res'] = 1024
opt = parse_opt_wrt_resolution(opt)
try:
  make_exp_dirs(opt)
except Exception as e:
  print(e)

# set up logger
save_log_path = f'{opt["path"]["log"]}'
editing_logger = get_root_logger(
    logger_name='editing',
    log_level=logging.INFO,
    log_file=f'{save_log_path}/editing.log')
editing_logger.info(dict2str(opt))

save_image_path = f'{opt["path"]["visualization"]}'
os.makedirs(save_image_path, exist_ok=True)
  
# ---------- create model ----------
opt['predictor_ckpt'] = './download/pretrained_models/predictor_1024.pth.tar'
field_model = create_model(opt)

画像のセットアップ

アップデートした画像からlatent codeを生成します。

latent_type='upload_image' #@param ['upload_image', 'random'] {allow-input: true}

if latent_type == 'random':
  latent_code = torch.randn(1, 512, device=torch.device('cuda'))
  with torch.no_grad():
    latent_code = field_model.stylegan_gen.get_latent(latent_code)
  latent_code = latent_code.cpu().numpy()

else:
  file_path = "/content/Talk-to-Edit/download/" + file_name
  latent_code = inversion(opt, field_model)

np.save(f'{opt["path"]["visualization"]}/latent_code.npz.npy', latent_code)

from models.utils import save_image
with torch.no_grad():
    start_image, start_label, start_score = \
            field_model.synthesize_and_predict(torch.from_numpy(latent_code).to(torch.device('cuda'))) # noqa

    save_image(start_image, f'{opt["path"]["visualization"]}/start_image.png')

import matplotlib.image as mpimg
import matplotlib.pyplot as plt
plt.figure()
plt.imshow(
    mpimg.imread(f'{opt["path"]["visualization"]}/start_image.png'))
plt.axis('off')
plt.show()

顔の属性変更(顔編集)

最後に顔の属性を変更します。なお、変更可能な属性は以下の通りです。

# ---------- decide which attribute to edit ----------
target_attribute='No_Beard' #@param ['Bangs', 'Eyeglasses', 'No_Beard', 'Smiling', 'Young'] {allow-input: false}
target_attr_value="5" #@param [0, 1, 2, 3, 4, 5] {allow-input: false}

# initialize attribtue_dict
attribute_dict = {
    "Bangs": start_label[0],
    "Eyeglasses": start_label[1],
    "No_Beard": start_label[2],
    "Smiling": start_label[3],
    "Young": start_label[4],
}

edit_label = {
    'attribute': target_attribute, 
    'target_score': int(target_attr_value)
    }

print_intermediate_result = False
round_idx = 0

edited_latent_code = None

設定した属性を反映します。

attribute_dict, exception_mode, _, edited_latent_code = edit_target_attribute(
    opt, attribute_dict, edit_label, round_idx, latent_code,
    edited_latent_code, field_model, None,
    print_intermediate_result, display_img=True
    )

if exception_mode != 'normal':
    if exception_mode == 'already_at_target_class':
        editing_logger.info("This attribute is already at the degree that you want. Let's try a different attribute degree or another attribute.")
    elif exception_mode == 'max_edit_num_reached':
        editing_logger.info("Sorry, we are unable to edit this attribute. Perhaps we can try something else.")

Smilingを5に設定した、大笑いの結果は以下の通りです。

口元のほうれい線は非常に自然に表現されています。口も自然に開いていますがもう少し口角を上げてほしいところですね。

また前髪と髭を生やしてみた結果は以下の通りです。

髪質が非常に自然で、もはやこういう人がいると言われても納得してしまいそうです。
本物の写真と区別がつきにくいからこそ、厳格な倫理観を持って技術を使っていきたいものです。

まとめ

本記事では、Talk-to-Editを使った顔の属性変更の方法をご紹介しました。
画像処理、特に顔や表情の編集は、スマホのアプリなどに適用されており身近な技術になりつつあります。
一方で身近であるからこそ技術の扱い方には気をつけてください。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
機械学習について学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - Talk-to-Edit: Fine-Grained Facial Editing via Dialog

2. GitHub - yumingj/Talk-to-Edit