概要

colabで公開されているPIFuHD DemoをAI PlatformのJupyterLabで実行してみました。
今回実行するコードはGitHubに入れてあります。

経緯

colabのdemoを実行して遊んでいたところ、colabのコメントに「解像度を上げるには自分のマシンで実行してね」と書かれていたので手持ちのMacで実行してみることにしました。

ところが、「CUDAがないよ！」とMacがエラーを吐いたため、CUDAといえばAI Platform（？）ということで、AI Platformで実行してみることにしました。
AssertionError: Torch not compiled with CUDA enabled

手順

GCPのアカウント作成等は割愛します。GCP Consoleの左のナビゲーションからAI Platformを選択してください。

(1) AI Platform画面の上部メニューから「NEW INSTANCE」を選択。インスタンスはPyTorch 1.4 with 1 NVIDIA Tesla T4とします。（当初CUDAといえばCUDA ToolKitだろうと思って試してみたのですが、torch.cuda.is_available()がFALSEになっていました）

(2) 「New Notebook Instance」ダイアログではRegionを指定します。Tokyo(asia-northeast1)で良いと思います。また、「Install NVIDIA GPU driver automatically for me」にチェックを入れておきます。

(3) そのままではメモリが足りないのでCUSTOMOIZEボタンを押して30GB(Machine Typeをn1-standard-8)にしてください。

(4) インスタンスが作成されたら「OPEN JUPYTERLAB」ボタンを押下します。

(5) Terminalを開いてGitHubからコードをcloneしてください。
git clone git@github.com:kurihiroi/pifuhd-for-ai-platform.git

(6) pifuhd-for-ai-platformの中にあるsetup.ipynbを上から順番に実行します。
途中numpyのアップグレードをしている箇所があるのですが、これをしないと3D生成スクリプト実行時にNotebookのカーネルがリセットされてしまってうまく実行できませんでした。なおアップグレード時にERRORが表示されますが、動作はするので気にしないことにします・・・

(7) pifuhd-for-ai-platformの中にあるimage2obj.ipynbを上から順番に実行します。
入力画像はpifuhd/sample_imagesフォルダに入っているtest.pngを使用します。変更する場合はimage2obj.ipynbの一番上のセルのfile_nameを変更してください。
以下のコマンドの-rオプションの値(256)を512に変更すると高解像度の3Dファイル(.obj)が生成されます。
!python -m apps.simple_test -r 256 --use_rect -i $image_dir
解像度を1024にするにはインスタンスのメモリが足りないのでn1-highmem-16(16 vCPU 104 GB memory)にする必要があります。試してみたのですがtest.pngの男性の右半身しか出力されませんでした。

(8) 実験が終わったらインスタンスを停止または削除してください。従量課金なので起動したままにはしない方が良いです。

出力結果

以下が解像度を256と512で試してみたときの比較画像です。若干ですが滑らかになっていますね。

写真には写っていない背中側もかなり自然に生成されています。すごいですね。

3Dプリンタで出力

ついでに3Dプリンタで出力してみました。ちゃんと立体になりました。

まとめ

写真１枚から簡単に立体が生成できてしまいました。colabでも十分な解像度だと思いますが、より高解像度を目指す方は試してみてください。
また、画像を差し替えてみても楽しいと思います。いろいろ試してみたところ頭から爪先までが写っている画像で、なるべく正面を向いている方が良いようです。
出力された3Dデータに色がついていないので、今後はテクスチャも貼ってみたいと思います。