【DragGAN】【StyleGAN-Human】DragGANで自前の人物画像を使ってみました(PyTorch=1.12.1+cu116)


はじめに

DragGANで自前の画像を使用する場合にはその画像に対して前処理が必要です。

今回の記事はそれに挑戦した記録です。

環境

PyTorch 2.0.1を使いたい人はこちらの環境構築を見てください。

DragGANとStyleGAN-Humanの両方を使いますが以下の環境で両方とも動きます。

Ubuntu 20.04 on WSL2
CUDA 11.6.2
cuDNN 8.4.0
Python 3.9

途中でPaddlePaddleを使いますがcuDNNをインストールしておく必要があります。
インストール方法はこちらを見て下さい。

pip install -r https://raw.githubusercontent.com/dai-ichiro/myEnvironments/main/StyleGAN-Human/requirements_cu116.txt
pip install paddlepaddle-gpu==2.4.2.post116 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html
git clone https://github.com/PaddlePaddle/PaddleSeg
cd PaddleSeg
git checkout release/2.5
pip install -U setuptools wheel
pip install -r requirements.txt
pip install -v -e .
pip install wandb==0.15.4

StyleGAN-Human(Invert real image with PTI )

まずはStyleGAN-Humanのリポジトリをクローンします。

git clone https://github.com/stylegan-human/StyleGAN-Human
cd StyleGAN-Human

「pretrained_models」フォルダを作成します。

mkdir pretrained_models

こちらから「stylegan_human_v2_1024.pkl」をダウンロードして、先ほど作った「pretrained_models」フォルダ内に保存します。

こちらから body_pose_model.pthをダウンロードして「StyleGAN-Human/openpose/model」フォルダ内に保存します。

こちらから「deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax.zip」をダウンロードして、解凍した後に「StyleGAN-Human/PP_HumanSeg/export_model」フォルダ内に保存します。
(補足)以下を実行するだけでもダウンロード、解凍が出来ます。

cd PP_HumanSeg/export_model
wget https://paddleseg.bj.bcebos.com/dygraph/humanseg/export/deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax.zip
unzip deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax.zip
rm deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax.zip



こちらから「deeplabv3p_resnet50_os8_humanseg_512x512_100k.zip」をダウンロードして、解凍した後に「StyleGAN-Human/PP_HumanSeg/pretrained_model」フォルダ内に保存します。
(補足)以下を実行するだけでもダウンロード、解凍が出来ます。

cd PP_HumanSeg/pretrained_model
wget https://paddleseg.bj.bcebos.com/dygraph/humanseg/train/deeplabv3p_resnet50_os8_humanseg_512x512_100k.zip
unzip deeplabv3p_resnet50_os8_humanseg_512x512_100k.zip
rm deeplabv3p_resnet50_os8_humanseg_512x512_100k.zip



こちらから「e4e_w+.pt」をダウンロードして「StyleGAN-Human/pti」フォルダ内に保存します。

ここで今回使う画像を用意しましょう。
写真はぱくたそから使わせて頂きました。こちらの画像になります。

6720x4480pxのLサイズをダウンロードして「girl.jpg」という名前で「StyleGAN-Human/img/test」フォルダ内に保存して下さい。

もともと「StyleGAN-Human/img/test」フォルダ内には「test.jpg」というサンプル画像が入っていますがそちらは削除して下さい。

画像が準備できたら「StyleGAN-Human」フォルダ直下で以下を実行します。

python alignment.py --image-folder img/test/ --output-folder aligned_image/

「StyleGAN-Human/aligned_image」フォルダ内に人物をトリミングした512x1024pxの「girl.png」が作成されているはずです。

続いて「StyleGAN-Human」フォルダ直下で以下を実行します。

python run_pti.py

「StyleGAN-Human/outputs/pti」フォルダに「girl.png」という画像が作成されています。

一番左が元画像、一番右が今回の学習後に作成された画像です。
(真ん中の画像はよくわかりません)
左の画像と右の画像がほとんど同じであればうまくいったということになります。
今回はうまくいっていません。

設定を変更します。
「StyleGAN-Human/pti/pti_configs/hyperparameters.py」を開いて19行目と20行目を以下のように変更します。

LPIPS_value_threshold = 0.005
max_pti_steps = 1000

変更後にもう一度「StyleGAN-Human」フォルダ直下で以下を実行します。

python run_pti.py

「StyleGAN-Human/outputs/pti」フォルダの「girl.png」をみてみると今度はこのようになっています。

うまくいきました。

「StyleGAN-Human/outputs/pti/checkpoints」フォルダ内の「model_girl.pkl」と「StyleGAN-Human/outputs/pti/embeddings/test/PTI/girl」フォルダ内の「0.pt」をこの後使用します。

DragGAN

まずはDragGANのリポジトリをクローンします。

git clone https://github.com/XingangPan/DragGAN

「DragGAN/legacy.py」を開いて40行目と41行目を削除するかコメントアウトする必要があります。

#assert isinstance(data['G'], torch.nn.Module)
#assert isinstance(data['D'], torch.nn.Module)



次に「DragGAN」フォルダ直下に先ほどStyleGAN-Humanを使って作成した「model_girl.pkl」と「0.pt」をコピーします。
その際に「model_girl.pkl」は「stylegan_human.pkl」に名前を変更する必要があります。
実際にはファイル名のどこかに「stylegan_human」という文字列が含まれていれば問題ないようですが今回は単純に「stylegan_human.pkl」としました。

「DragGAN」フォルダ直下で以下を実行します。

export MESA_GL_VERSION_OVERRIDE=3.3
python visualizer_drag.py stylegan_human.pkl

このようなGUIが表示されます。



上の赤枠のテキストボックスに「0.pt」と入力し、その下の「Load latest」ボタンを押せば画像がこのように変わるはずです。



これで自前の人物画像に対して好きなようにDragGANが使用できるようになりました。

関連記事

touch-sp.hatenablog.com
touch-sp.hatenablog.com
touch-sp.hatenablog.com