PULSEで低解像度の顔画像を高解像度に変換する

1.はじめに

　皆さんは、低解像度の顔画像を高解像度の画像に変換したいと思ったことはありませんか。

　2020年6月デューク大は従来8倍が最大であった高解像度変換を64倍にまで引き上げる PULSE という手法を発表しました。これにより、目元も口元もハッキリしない様な顔画像が毛穴や質感が分かるくらいの高画質に変換できます。

2.アルゴリズム

　PULSEは、StyleGAN学習済みモデルを使っていますので、まずStyleGANの説明から行います。

　StyleGAN 学習済みモデルは、大量の顔画像を学習していて様々な高画質の顔画像を生成できます。そして、その画像生成プロセスは、低解像度から高解像度まで9段階で画像生成を行います。

　PULSEは、低解像度のターゲット画像から顔画像を切り出し・リサイズした低解像度画像とStyleGANの低解像度画像との差をロスとして、ロスを最小化するStyleGANのパラメータを求めることによって、高解像度の顔画像を得ます。

　それでは、実際に動かしてみましょう。

3.コード

　コードはGoogle Colabで動かす形にしてGithubに上げてありますので、それに沿って説明して行きます。自分で動かしてみたい方は、この「リンク」をクリックし表示されたノートブックの先頭にある「Colab on Web」ボタンをクリックすると動かせます

　まず、セットアップを行います。

# githubからコードのコピー
!git clone https://github.com/cedro3/pulse.git

1 2	# githubからコードのコピー !git clone https://github.com/cedro3/pulse.git

# ディレクトリをpulseへ移動
%cd pulse

1 2	# ディレクトリをpulseへ移動 %cd pulse

# フォルダーの設定
import os
os.makedirs('out', exist_ok=True)
os.makedirs('input', exist_ok=True)
os.makedirs('runs', exist_ok=True)

# フォルダーの設定

import os

os.makedirs('out', exist_ok=True)

os.makedirs('input', exist_ok=True)

os.makedirs('runs', exist_ok=True)

　必要な関数を定義します。

# フォルダー内の画像表示
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import glob

def display(DIR):
  files = glob.glob(DIR+'/*.png')
  files.sort()
  fig = plt.figure(figsize=(30, 30))    
  for i, file in enumerate(files):
      img = Image.open(file)
      img_plt = np.asarray(img)
      ax = fig.add_subplot(10, 10, i+1, xticks=[], yticks=[])
      ax.imshow(img_plt)
      ax.set_xlabel(file.replace(DIR+'/',''), fontsize=12)        
  plt.show()
  plt.close()

# フォルダー内の画像表示

import numpy as np

from PIL import Image

import matplotlib.pyplot as plt

import glob

def display(DIR):

files = glob.glob(DIR+'/*.png')

files.sort()

fig = plt.figure(figsize=(30, 30))

for i, file in enumerate(files):

img = Image.open(file)

img_plt = np.asarray(img)

ax = fig.add_subplot(10, 10, i+1, xticks=[], yticks=[])

ax.imshow(img_plt)

ax.set_xlabel(file.replace(DIR+'/',''), fontsize=12)

plt.show()

plt.close()

　sample/picの中には、下記の小さなサンプル画像が3枚保存されています。今回は、これらの画像から顔画像を切り出し、高解像度の画像に変換します。

　まず、3枚の画像から顔画像を切り出します。切り出し方は、STyleGANに学習させた画像と同じにします。そして、StyleGAN 学習済みモデルでロスを計算する解像度はデフォルトの32×32とします。

# サンプル画像から顔画像を切り出す
!python align_face.py -input_dir 'sample/pic' -output_dir 'input' -output_dir2 'out' -output_size=32
display('./out')

# サンプル画像から顔画像を切り出す

!python align_face.py -input_dir 'sample/pic' -output_dir 'input' -output_dir2 'out' -output_size=32

display('./out')

　顔画像を切り出したものがこれで、さすがに小さな画像から切り出したので、ボケボケですよね。この画像を高解像度にします。

　32×32の画像から1024×1024の画像へ、steps=500で高解像度化します。

# 高解像度化 (inputフォルダーの画像を高解像度化し、runsフォルダーへ保存)
!python run.py -input_dir 'input' -output_dir 'runs' -steps 500

# 画像表示
display('./out')
display('./runs')

# 高解像度化 (inputフォルダーの画像を高解像度化し、runsフォルダーへ保存)

!python run.py -input_dir 'input' -output_dir 'runs' -steps 500

# 画像表示

display('./out')

display('./runs')

　上段がサンプル画像から切り出した低解像度画像、下段がそれを高解像度化した画像です。解像度は上がりましたが、再現性がイマイチですよね。何か外人っぽさがある様な感じです。

　それでは、StyleGAN 学習済みモデルでロスを計算する解像度を32×32→64×64、stepsも500→1,000へ上げてやってみましょう。

!python align_face.py -input_dir 'sample/test' -output_dir 'input/test' -output_dir2 'out/test' -output_size=64
!python run.py -input_dir 'input/test' -output_dir 'runs/test' -steps 1000
display('./out/test')
display('./runs/test')