本記事では、PAMAと呼ばれる機械学習手法を用いて、任意の画像に任意の写真のスタイルを転送する方法をご紹介します。

出典: luoxuan-cs/PAMA

PAMA

概要

近年、詳細なスタイル転送のため、attentionを用いたスタイル転送方法が提案されています。
この手法は、スタイル転送のためコンテンツとスタイル間のポイントごとの類似性を操作しています。しかし、特徴点に基づくattentionメカニズムは特徴のmulti-manifold分布を無視するという問題がありました。この結果セマンティック領域の不整合を持つスタイル転送結果が出力される場合がありました。

Progressive Attentional Manifold Alignment(PAMA)では、attentionとspace-aware補間を繰り返し適用しています。 attention操作により、コンテンツ特徴量の空間分布に従って、スタイル特徴量を動的に再配置します(attentional manifold alignment:AMA)。 これにより、コンテンツとスタイルのmanifoldが特徴量マップ上で対応するようになり、space-aware補間により、両者のmanifold間を補完し、それらの類似性を高めていきます。
結果、PAMAはセマンティック領域の不整合を回避しながら最先端のパフォーマンスを実現しています。

出典: Consistent Style Transfer

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、スタイル転送を行っていきます。

デモ(Colaboratory)

それでは、実際に動かしながらスタイル転送を行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

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

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

初めにGithubからソースコードを取得します。

%cd /content

!git clone https://github.com/luoxuan-cs/PAMA.git

次にライブラリをインストールします。

%cd /content/PAMA

!pip install --upgrade gdown

最後にライブラリをインポートします。

%cd /content/PAMA

import os
import gdown
import shutil
import argparse
import glob

import torch
import torch.nn as nn
import torch.utils.data as data
from torchvision.utils import save_image
from PIL import Image, ImageFile
from net import Net
from utils import DEVICE, train_transform, test_transform, FlatFolderDataset, InfiniteSamplerWrapper, plot_grad_flow, adjust_learning_rate
Image.MAX_IMAGE_PIXELS = None  
ImageFile.LOAD_TRUNCATED_IMAGES = True

以上で環境セットアップは完了です。

学習済みモデルのセットアップ

学習済みモデルをGoogle Driveからダウンロードします。

%cd /content/PAMA
!mkdir checkpoints

# https://drive.google.com/file/d/1rPB_qnelVVSad6CtadmhRFi0PMI_RKdy/view
original = 'checkpoints/original_PAMA.zip'
if not os.path.exists(original):
  gdown.download('https://drive.google.com/uc?id='+'1rPB_qnelVVSad6CtadmhRFi0PMI_RKdy', original, quiet=False)
  shutil.unpack_archive(original, 'checkpoints')

# https://drive.google.com/file/d/1IrggOiutiZceJCrEb24cLnBjeA5I3N1D/view
wo_color = 'checkpoints/PAMA_without_color.zip'
if not os.path.exists(wo_color):
  gdown.download('https://drive.google.com/uc?id='+'1IrggOiutiZceJCrEb24cLnBjeA5I3N1D', wo_color, quiet=False)
  shutil.unpack_archive(wo_color, 'checkpoints')

# https://drive.google.com/file/d/1HXet2u_zk2QCVM_z5Llg2bcfvvndabtt/view
color = 'checkpoints/PAMA_1.5_color.zip'
if not os.path.exists(color):
  gdown.download('https://drive.google.com/uc?id='+'1HXet2u_zk2QCVM_z5Llg2bcfvvndabtt', color, quiet=False)
  shutil.unpack_archive(color, 'checkpoints')

# https://drive.google.com/file/d/13m7Lb9xwfG_DVOesuG9PyxDHG4SwqlNt/view
content = 'checkpoints/PAMA_1.5_content.zip'
if not os.path.exists(content):
  gdown.download('https://drive.google.com/uc?id='+"13m7Lb9xwfG_DVOesuG9PyxDHG4SwqlNt", content, quiet=False)
  shutil.unpack_archive(content, 'checkpoints')

checkpointsに各種モデルをダウンロードしています。

テスト画像のセットアップ

コンテンツ画像とスタイル画像をそれぞれダウンロードします。
本記事では、ぱくたそ様の画像などを使用させていただきます。

%cd /content/PAMA
!mkdir -p tests/contents tests/styles

!wget -c https://www.pakutaso.com/shared/img/thumb/20220227-A7401834_TP_V4.jpg \
      -O ./tests/contents/test_1.jpg
!wget -c https://www.pakutaso.com/shared/img/thumb/SAYA160312500I9A3721_TP_V4.jpg \
      -O ./tests/contents/test_2.jpg
!wget -c https://www.pakutaso.com/shared/img/thumb/unific528--8628_TP_V4.jpg \
      -O ./tests/contents/test_3.jpg
!wget -c https://www.publicdomainpictures.net/pictures/80000/nahled/animal-sketch-13919381209K9.jpg \
      -O ./tests/styles/style_1.jpg
!wget -c https://www.publicdomainpictures.net/pictures/390000/velka/the-starry-night-van-gogh.jpg \
      -O ./tests/styles/style_2.jpg
!wget -c https://jojo-animation.com/img/top/mv_2.jpg \
      -O ./tests/styles/style_3.jpg

スタイル転送

それでは、ダウンロードしたコンテンツ画像にスタイル画像のスタイルを転送していきます。

まず、使用するモデルを選択します。

model_type = "ORIGINAL" #@param ["ORIGINAL", "WO_COLOR", "COLOR", "CONTENT"]

%cd /content/PAMA
!rm -rf ./checkpoints/*.pth
if model_type == "ORIGINAL":
  !cp ./checkpoints/original_PAMA/*.pth ./checkpoints
if model_type == "WO_COLOR":
  !cp ./checkpoints/original_PAMA/encoder.pth ./checkpoints
  !cp ./checkpoints/PAMA_without_color/*.pth ./checkpoints
if model_type == "COLOR":
  !cp ./checkpoints/original_PAMA/encoder.pth ./checkpoints
  !cp ./checkpoints/PAMA_1.5_color/*.pth ./checkpoints
if model_type == "CONTENT":
  !cp ./checkpoints/original_PAMA/encoder.pth ./checkpoints
  !cp ./checkpoints/PAMA_1.5_content/*.pth ./checkpoints

続いて入力画像のパスなどを設定します。

args = argparse.ArgumentParser()
args.pretrained = True
args.requires_grad = True
args.training = False

args.outdir = '/content/PAMA/tests_result'

args.run_folder = True
args.content = '/content/PAMA/tests/contents'
args.style = '/content/PAMA/tests/styles'
# args.run_folder = False
# args.content = '/content/PAMA/tests/contents/test_1.jpg'
# args.style = '/content/PAMA/tests/styles/style_1.jpg'

スタイル転送を実行します。

def inference(model, device, content_path, style_path, out):
  Ic = tf(Image.open(content_path)).to(device)
  Is = tf(Image.open(style_path)).to(device)
  Ic = Ic.unsqueeze(dim=0)
  Is = Is.unsqueeze(dim=0)

  with torch.no_grad():
    Ics = model(Ic, Is)

  os.makedirs(out, exist_ok=True)
  filename = "res_" + os.path.splitext(os.path.basename(style_path))[0] + "_" + os.path.basename(content_path)
  name_cs = os.path.join(out, filename)
  save_image(Ics[0], name_cs)

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Net(args)
model.eval()

model = model.to(DEVICE)

tf = test_transform()
if args.run_folder == True:
  for content_path in glob.glob(os.path.join(args.content, "*.*")):
    for style_path in glob.glob(os.path.join(args.style, "*.*")):
      inference(model, DEVICE, content_path, style_path, args.outdir)
else:
  inference(model, DEVICE, args.content, args.style, args.outdir)

出力結果は以下の通りです。
左から、コンテンツ画像、スタイル画像、スタイル転送結果の順に並んでいます。

自然に白黒のスケッチが転送されています。

まとめ

本記事では、PAMAと呼ばれる機械学習手法を用いて、任意の画像のスタイルを任意の画像に転送する方法をご紹介しました。
スタイル画像を選択できる点自由度が高く汎用性がありますね。

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

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

参考文献

1.  論文 - Consistent Style Transfer

2. GitHub - luoxuan-cs/PAMA