kedzkiestの日記

普段の発見や開発の進捗について書いていきます。

Shap-Eを使ってローカル環境で3Dモデルを生成してみる

はじめに

OpenAI社が2023年の3月に論文を発表し、そこでShap-Eというモデルが登場しました。
論文へのリンク:
Shap-E: Generating Conditional 3D Implicit Functions


Shap-Eの概要(原文をDeepL翻訳しました)
我々は、3D資産のための条件付き生成モデルであるShap-Eを発表する。Shap-Eは、単一の出力表現を生成する3D生成モデルの最近の研究とは異なり、テクスチャメッシュと神経放射フィールドの両方としてレンダリング可能な暗黙関数のパラメータを直接生成します。Shap-Eの学習は2段階で行われます。1つ目は、3D資産を陰関数のパラメータに決定論的にマッピングするエンコーダの学習、2つ目は、エンコーダの出力に対して条件付き拡散モデルを学習します。3Dデータとテキストデータを組み合わせた大規模なデータセットで学習した結果、複雑で多様な3D資産を数秒で生成することができました。点群に対する明示的な生成モデルであるPoint-Eと比較すると、Shap-Eは高次元の多表現出力空間をモデル化しているにもかかわらず、収束が早く、サンプル品質も同等かそれ以上であることがわかります。モデルの重み、推論コード、サンプルは、以下のサイトで公開しています。


ここで紹介されている以下のサイトとはGitHub上に公開されているShap-Eのサイトのことです。
github.com


概要にある通り、実装やサンプルコードなどが紹介されています。

今回やることとして、サンプルコードのうち、指定したテキストをもとに3Dモデルを作成するものを試してみたいと思います!
また、GitHubの方では.ipynbファイルでサンプルコードが提供されており、クラウド環境でJupyter NotebookやGoogle Colaboratoryを使用してコードを実行することが想定されているようですが、今回はローカル環境pythonプログラムを実行して同じことをやってみたいと思います!

作成したプロジェクト

github.com

↑公開から2日で100スター以上付きました!!

環境

OS:Ubuntu 20.04.5 LTS
CPU:Intel(R) Core(TM) i7-6800K CPU @ 3.40GHz
GPUNVIDIA GeForce RTX 3090
Cuda version:12.1
pip version:23.1.2
conda version:23.5.0

コード実行までの手順

shap-e-localの方にも英語で手順を書いていますが、こちらでは日本語での手順を載せておきます。

1. shap-e-localをzipフォルダとしてダウンロードし、解凍する

2. Anacondaを使用して仮想環境を作る
 作れたら仮想環境を起動

conda create -n shap-e python=3.10
conda activate shap-e

3. カレントディレクトリを解凍したフォルダに変更する

// 例
cd shap-e-local-main

4. 必要なライブラリをインストール

pip install -e .
pip install pyyaml ipywidgets fvcore iopath

5. testTorch.pyを実行してtorchがcudaを認識しているか確かめる

python testTorch.py

testTorch.pyの出力が下のようになっていれば、torchがcudaを認識できています!
PyTorch is using: cuda
PyTorch version: 2.0.1+cu117
Result of tensor operation: tensor([ 8., 15.], device='cuda:0')

6. pytorch3dPrep.pyを実行して適切なpytorch3dのバージョンを見つける

python pytorch3dPrep.py

pytorch3dPrep.pyは例えば"py310_cu117_pyt201"のような文字列を出力します。
この文字列を覚えておいてください。

7. pytorch3dのライブラリをインストール

pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{output string}/download.html

// 例
pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu117_pyt201/download.html

8. 3Dモデルを生成する
下のコマンドを実行する前にtextTo3D.pyを編集することで、テキストプロンプトやバッチサイズを変更できます。

python textTo3D.py

指定したバッチサイズに応じて、いくつか.objファイルと.plyファイルが出力されます。

9. 生成された3Dモデルのプレビューを見たい方へ
まずは3Dモデルを表示するのに必要なライブラリをインストールします。

pip install vedo
pip install open3d

次に、下のコマンドを使って生成された3Dモデルのプレビューを見れます。
コマンドライン引数のfilenameは拡張子を除いたファイル名のことです。例えばpython displayPly.py example_mesh_0を実行すると、example_mesh_0.plyファイルのプレビューが表示されます。)

python displayObj.py {filename}
python displayPly.py {filename}

displayObj.pyやdisplayPly.pyといった実行プログラムはプレビューを見ようとしている.objファイルや.plyファイルと同じ階層にあるようにしてください(デフォルトでそうなっています)。

デモ

生成した3DモデルをBlenderに持ってきてみた

入力となるテキストインプットをそれぞれ"an apple"、"a shark"、"a table"としてりんご、サメ、机の3Dモデルを生成したあと、それらをBlenderにインポートして簡単なシーンを作ってみました。

床や壁はBlenderの平面オブジェクトです。

おわりに

Shap-E公式のGitHubサイトで提供されているtext-to-3dのサンプルコードを書き直して、ローカル環境で3Dモデルを生成することができました!

Shap-Eは従来の最適化ベースでの生成モデルに比べて非常に推論時間が短く、10~20秒ほどで指定した3Dモデルを作ってくれるので、試行錯誤を高速に繰り返せる点で優れていますね!

また、生成された3Dモデルをそのまま使うのではなく、Blenderなどで手直しする、別の最適化ベースの生成モデルの入力として使う、といった応用も考えられます。

最後になりますが、よろしければshap-e-localを試してみて、不具合報告や改善の提案などをissueにしていただけると嬉しいです!
github.com