PF-AFNで、モデルに衣服を仮装試着させる

cedro

3年前

1.はじめに

　衣服の仮装試着は、eコマースやファッション画像編集などで様々なアプリケーションが発表されており注目を集めています。今回ご紹介するのは、従来の WUTON という技術に改良を加えた PF-AFN という技術です。

＊この論文は、2021.5に提出されました。

2.PF-AFNとは？

　WUTON は、まず Perser-based Teacher Network を Parsing（Real Image をもとに身体を区分したデータ）と Clothes（衣服データ）を入力として Fake Image を出力するように訓練します。そして、Perser-free Student Network を Real Image と Clothes を入力として Fake Image を出力するように訓練（蒸留）します。言い換えると、学生ネットワークが教師ネットワークを模倣するわけです。但し、このモデルには弱点があって、学生ネットワークは教師ネットワークが持つ不具合も学習してしまうことです。

　そこで、PF-AFNの登場です。まず、Persing と Clothes から Fake Image を出力するネットワークを訓練するまでは同じですが、このネットワークを教師ネットワークではなく不具合な出力をすることもある家庭教師ネットワークとして扱います。なので、名称を Perser-based Tutor Network と言っています。そして、Parser-free student Network を Fake Image と Clothes on Real Image を入力とし、Appearance Flows（人物画像と衣服画像間の相関関係）のみ蒸留し、 Teacher Knowledge（Real Image）を出力するように訓練します。こうすることで、最終出力を Real Image と比較するため不具合を減らすことができます。

　では、コードを動かしてみましょう。

3.コード

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

　まず、セットアップを行います。ディレクトリは、PF-AFN/PF-AFN_test にしています。

# githubからコードを取得
! git clone https://github.com/geyuying/PF-AFN.git
%cd PF-AFN/PF-AFN_test

# 学習済みパラメータをダウンロード
import gdown
gdown.download('https://drive.google.com/u/0/uc?id=1CKB11poq6qcHXwf7WeoTrSc5jRICTJfd', 'checkpoints.zip', quiet=False)
! unzip checkpoints.zip

# フォルダー内表示関数
import matplotlib.pyplot as plt
from PIL import Image
import os
import numpy as np
%matplotlib inline

def display_pic(folder):
    fig = plt.figure(figsize=(30, 40))
    files = os.listdir(folder)
    files.sort()
    for i, file in enumerate(files):
        if file=='.ipynb_checkpoints':
           continue
        if file=='.DS_Store':
           continue
        img = Image.open(folder+'/'+file)    
        images = np.asarray(img)
        ax = fig.add_subplot(10, 10, i+1, xticks=[], yticks=[])
        image_plt = np.array(images)
        ax.imshow(image_plt)
        ax.set_xlabel(file, fontsize=15)               
    plt.show()
    plt.close()

# githubからコードを取得

! git clone https://github.com/geyuying/PF-AFN.git

%cd PF-AFN/PF-AFN_test

# 学習済みパラメータをダウンロード

import gdown

gdown.download('https://drive.google.com/u/0/uc?id=1CKB11poq6qcHXwf7WeoTrSc5jRICTJfd', 'checkpoints.zip', quiet=False)

! unzip checkpoints.zip

# フォルダー内表示関数

import matplotlib.pyplot as plt

from PIL import Image

import os

import numpy as np

%matplotlib inline

def display_pic(folder):

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

files = os.listdir(folder)

files.sort()

for i, file in enumerate(files):

if file=='.ipynb_checkpoints':

continue

if file=='.DS_Store':

continue

img = Image.open(folder+'/'+file)

images = np.asarray(img)

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

image_plt = np.array(images)

ax.imshow(image_plt)

ax.set_xlabel(file, fontsize=15)

plt.show()

plt.close()

　仮装試着をさせるためのデータを確認します。データは、test_clothes（衣服）、test_edge（マスク）、test_img（モデル）の３種類があり、test_clothes に対応する test_edge は同じファイル名になっていることに注意して下さい。

display_pic('dataset/test_clothes')
display_pic('dataset/test_edge')
display_pic('dataset/test_img')

display_pic('dataset/test_clothes')

display_pic('dataset/test_edge')

display_pic('dataset/test_img')

　それでは、仮装試着をしてみましょう。どの test_img にどの test_clothes を着用させるかの設定は、demo.txt に下記のように記載されています。

000066_0.jpg 017575_1.jpg
016962_0.jpg 003434_1.jpg
004912_0.jpg 014396_1.jpg
005510_0.jpg 006026_1.jpg
014834_0.jpg 019119_1.jpg
015794_0.jpg 010567_1.jpg

000066_0.jpg 017575_1.jpg

016962_0.jpg 003434_1.jpg

004912_0.jpg 014396_1.jpg

005510_0.jpg 006026_1.jpg

014834_0.jpg 019119_1.jpg

015794_0.jpg 010567_1.jpg

　demo.txt の設定に従って、コスチュームを入れ替えます。入れ替えた結果は、results/demo/PFAFN に保存されます

! python test.py

1	! python test.py

　保存された結果を表示してみましょう。

from IPython.display import Image, display_jpeg
import glob

files = glob.glob('results/demo/PFAFN/*.jpg')
files.sort()

for file in files:
    print(file)
    display_jpeg(Image(file))