ColabでStyle Transfer（画風変換）を行なってみる｜Tech Press

画風変換のアルゴリズムはいくつかありますが、今回はこちらを使用してみます。

https://github.com/ProGamerGov/neural-style-pt

AI Hub(旧：seedbank)にあるコードを利用してもいいのですが、pytorchを使いたかったので上記のコードを利用します。

style transferの元になっている考え方は基本同じなのであまり気にせず実装していきます。

準備

早速、colabを起動してランタイムのタイプをGPUに変更しておきます。

後でtorchとcudaのバージョンを指定して再インストールしますが、Cudaのバージョンを確認してtorchと整合が取れているのであればそのまま使用しても大丈夫です。

!nvcc --version

続いて、入力画像や出力画像、checkpointを取り扱うのでドライブをマウントしておきます。

学習からは行わず、既存のパラメータを使用します。pthファイルをドライブ置く必要があるので、colab proを使っている方はターミナルからwgetを使用した方が早いですが、そうでない場合は時間がかかりますがダウンロードしてアップしてください。

https://drive.google.com/file/d/1J5co7VB-f1ft_dBB3xn4fNL7fQ7YaxIu/view?usp=sharing

zipファイルになっているので、アップした後に解凍してください。そうした方が時間を短縮することができます。

ドライブ側で必要なディレクトリを作成しておきます。chepointsフォルダにアップロードして解凍している状態になっています。

inputフォルダには変換したい画像を置いてください。

style1フォルダには適応させたい画像をおきます。

実装

必要なファイル群をgitでcloneしておきます。

!git clone https://github.com/ProGamerGov/neural-style-pt.git

先程アップロードしたパラメータファイルをcloneしたフォルダにコピーします。

パスが違う方は適宜変更してください。

!cp -ri "/content/drive/MyDrive/NSPT/checkpoints/channel_pruning.pth" /content/neural-style-pt/models/
!cp -ri "/content/drive/MyDrive/NSPT/checkpoints/nin_imagenet.pth" /content/neural-style-pt/models/
!cp -ri "/content/drive/MyDrive/NSPT/checkpoints/nyud-fcn32s-color-heavy.pth" /content/neural-style-pt/models/
!cp -ri "/content/drive/MyDrive/NSPT/checkpoints/vgg16-00b39a1b.pth" /content/neural-style-pt/models/
!cp -ri "/content/drive/MyDrive/NSPT/checkpoints/vgg19-d01eb7cb.pth" /content/neural-style-pt/models/

続いて、torchとcudaを再インストールしておきます。

!pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio===0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

最後にstyle transferを実行します。

ここでもパスは適宜変更してください。

!python neural-style-pt/neural_style.py -style_image '/content/drive/MyDrive/NSPT/style/style1' -style_weight 1000 -style_scale 1 -content_image '/content/drive/MyDrive/NSPT/input/input.jpeg' -init random -learning_rate 1 -print_iter 50 -save_iter 250 -image_size 512 -num_iterations 1000 -model_file '/content/drive/MyDrive/NSPT/checkpoints/nyud-fcn32s-color-heavy.pth' -content_layers relu1_1,relu2_1,relu3_1,relu4_1,relu5_1 -style_layers relu1_1,relu2_1,relu3_1,relu4_1,relu5_1 -optimizer lbfgs -output_image '/content/drive/MyDrive/NSPT/output/A1.png' -tv_weight 0.00001 -original_colors 0 -backend cudnn 
!python neural-style-pt/neural_style.py -style_image '/content/drive/MyDrive/NSPT/style/style1' -style_weight 1000 -style_scale 1 -content_image '/content/drive/MyDrive/NSPT/input/input.jpeg' -content_weight 5 -init image -init_image '/content/drive/MyDrive/NSPT/output/A1.png' -print_iter 50 -save_iter 250 -image_size 768 -num_iterations 600 -model_file '/content/drive/MyDrive/NSPT/checkpoints/nyud-fcn32s-color-heavy.pth' -content_layers relu1_1,relu2_1,relu3_1,relu4_1,relu5_1 -style_layers relu1_1,relu2_1,relu3_1,relu4_1,relu5_1 -optimizer lbfgs -output_imag '/content/drive/MyDrive/NSPT/output/A2.png' -tv_weight 0 -gpu 0 -original_colors 0 -backend cudnn
!python neural-style-pt/neural_style.py -style_image '/content/drive/MyDrive/NSPT/style/style1' -style_weight 1000 -style_scale 1 -content_image '/content/drive/MyDrive/NSPT/input/input.jpeg' -content_weight 5 -init image -init_image '/content/drive/MyDrive/NSPT/output/A2.png' -print_iter 50 -save_iter 250 -image_size 1024 -num_iterations 1000 -model_file '/content/drive/MyDrive/NSPT/checkpoints/nyud-fcn32s-color-heavy.pth' -content_layers relu1_1,relu2_1,relu3_1,relu4_1,relu5_1 -style_layers relu1_1,relu2_1,relu3_1,relu4_1,relu5_1 -optimizer lbfgs -output_image '/content/drive/MyDrive/NSPT/output/A3.png' -tv_weight 0.00001 -gpu 0 -original_colors 0 -backend cudnn
!python neural-style-pt/neural_style.py -style_image '/content/drive/MyDrive/NSPT/style/style1' -style_weight 40000 -style_scale 1 -content_image '/content/drive/MyDrive/NSPT/input/input.jpeg' -content_weight 5 -init image -init_image '/content/drive/MyDrive/NSPT/output/A3.png' -print_iter 50 -save_iter 250 -image_size 1400 -num_iterations 1000 -model_file '/content/drive/MyDrive/NSPT/checkpoints/nyud-fcn32s-color-heavy.pth' -content_layers relu1_1,relu2_1,relu3_1,relu4_1,relu5_1 -style_layers relu1_1,relu2_1,relu3_1,relu4_1,relu5_1 -optimizer lbfgs -output_image '/content/drive/MyDrive/NSPT/output/A4.png' -tv_weight 0 -gpu 0 -original_colors 0 -backend cudnn
!python neural-style-pt/neural_style.py -style_image '/content/drive/MyDrive/NSPT/style/style1' -style_weight 80000 -style_scale 1 -content_image '/content/drive/MyDrive/NSPT/input/input.jpeg' -content_weight 5 -init image -init_image '/content/drive/MyDrive/NSPT/output/A4.png' -print_iter 50 -save_iter 250 -image_size 1400 -num_iterations 1000 -model_file '/content/drive/MyDrive/NSPT/checkpoints/nyud-fcn32s-color-heavy.pth' -content_layers relu1_1,relu1_2,relu2_1,relu2_2,relu3_1,relu3_2,relu3_3,relu3_4,relu4_1,relu4_2,relu4_3,relu4_4,relu5_1,relu5_2,relu5_3,relu5_4 -style_layers relu1_1,relu1_2,relu2_1,relu2_2,relu3_1,relu3_2,relu3_3,relu3_4,relu4_1,relu4_2,relu4_3,relu4_4,relu5_1,relu5_2,relu5_3,relu5_4 -optimizer lbfgs -output_image '/content/drive/MyDrive/NSPT/output/A5.png' -tv_weight 0 -gpu 0 -original_colors 0 -backend cudnn
!python neural-style-pt/neural_style.py -style_image '/content/drive/MyDrive/NSPT/style/style1' -style_weight 40000 -style_scale 1 -content_image '/content/drive/MyDrive/NSPT/input/input.jpeg' -content_weight 0 -init image -init_image '/content/drive/MyDrive/NSPT/output/A5.png' -print_iter 50 -save_iter 100 -image_size 1800 -num_iterations 200 -model_file '/content/drive/MyDrive/NSPT/checkpoints/nyud-fcn32s-color-heavy.pth' -content_layers relu1_1,relu1_2 -style_layers relu1_1,relu1_2 -optimizer lbfgs -output_image '/content/drive/MyDrive/NSPT/output/A6.png' -tv_weight 0 -gpu 0 -original_colors 0 -backend cudnn
!python neural-style-pt/neural_style.py -style_image '/content/drive/MyDrive/NSPT/style/style1' -style_weight 40000 -style_scale 1 -content_image '/content/drive/MyDrive/NSPT/input/input.jpeg' -content_weight 15 -init image -init_image '/content/drive/MyDrive/NSPT/output/A6.png' -learning_rate 1 -print_iter 50 -save_iter 100 -image_size 3600 -num_iterations 200 -model_file '/content/drive/MyDrive/NSPT/checkpoints/channel_pruning.pth' -optimizer adam -output_image '/content/drive/MyDrive/NSPT/output/A7.png' -tv_weight 0 -original_colors 0 -backend cudnn
!python neural-style-pt/neural_style.py -style_image '/content/drive/MyDrive/NSPT/style/style1' -style_weight 1500 -style_scale 0.5 -content_image '/content/drive/MyDrive/NSPT/input/input.jpeg' -content_weight 0 -init image -init_image '/content/drive/MyDrive/NSPT/output/A7.png' -learning_rate 1 -print_iter 50 -save_iter 0 -image_size 7200 -num_iterations 10 -model_file '/content/drive/MyDrive/NSPT/checkpoints/nin_imagenet.pth' -content_layers relu0,relu1 -style_layers relu0,relu1 -optimizer adam -output_image '/content/drive/MyDrive/NSPT/output/A8-FINAL.png' -tv_weight 0 -original_colors 0 -backend cudnn

実行完了には、P100の環境で20分くらいかかりました。

処理が終わるとoutputフォルダに途中の画像も生成されているので、確認してみてください。

実際にやってみるとこんな感じになります。