NeRFのレンダリング表現を3D情報として得る方法

こんにちは,Ridge-iの関口です.本記事ではNeural Radiance Fields (NeRF) [1]の技術で場を学習したネットワークから3D情報を取得する方法についてご紹介したいと思います.といっても著者らのオリジナル研究のプロジェクトページ上にあるスクリプトの解説に近いのですが,ウェブ上にあまり情報がなかったので書き記しておきたいと思います.

NeRFは任意視点の画像をレンダリングできますが,あくまで2Dとしての見え方を推定しているものです(正確に言えば,NNに対して座標とその座標を覗き込む角度を入力として,そのときの点の色と密度を推定する作業を繰り返して表面座標の情報を得る,ということを画素分繰り返して画像を得ています).では学習したNeRFから3Dモデルを得ることはできるのでしょうか? 3Dデータとして出力できればゲームやVR,3Dプリンティングといった色々なことに使えそうです.NeRFは場を学習したNNですから3D表現を獲得しているはずなので,なんとかしてここから3Dデータを得たいところです.ではどのようにすればよいでしょうか? 本記事ではその方法について述べます.

今回解説するもととなるスクリプトや学習済みの重みは著者のGitHub[2]で公開されているのでぜひお試しになってみてください.

導入

本記事で対象とするのはオリジナルのNeRF論文[1]で提案されているニューラルネットワークから3Dのボリューメトリックデータを取得して,ポリゴンデータ(メッシュデータ)に変換する方法です.NeRFについての解説は省略しますので,NeRF自体の技術をお知りになりたい方は原論文をご参照ください.

今回は問題設定をシンプルにして,色なしの3Dのメッシュデータを得る方法について述べたいと思います. 本記事では以下に示すレゴのおもちゃを学習させたNeRFから3Dメッシュを抽出します.

Photo by [2]

方法

基本的な考え方としてはいたってシンプルで,NeRFを使って場を学習したNNに対して様々な視点からレイを投げてレンダリングを行い物体表面の点群を得るというものです.

それではコードとともに解説したいと思います.全てのコードではなく,学習済みのモデルを読み込むところから掲載します.

# ネットワークを定義し,重みを読み込む
model = MLP()
model.load_state_dict(torch.load(WEIGHT_PATH))

# 様々な方向から見たレイを定義する
N = 256
t = np.linspace(-1.2, 1.2, N+1)
query_pts = np.stack(np.meshgrid(t, t, t), -1).astype(np.float32)
sh = query_pts.shape
flat = query_pts.reshape([-1,3])
flat = torch.Tensor(flat)
flat = positional_encoder(flat)

# chunk数ごとにレイを分割してNNに入力する
chunk = 1024*64
with torch.no_grad():
    raw = np.concatenate([model(flat[i:i+chunk]).detach().numpy().copy() for i in range(0, flat.shape[0], chunk)], 0)

# レイ上で最大の密度を持つ位置を点とする
raw = np.reshape(raw, list(sh[:-1]) + [-1])
sigma = np.maximum(raw[...,-1], 0.)

# 密度の分布を描画する
plt.hist(np.maximum(0,sigma.ravel()), log=True)
plt.show()

1つめのブロックではNeRFのNNをインスタンス化し,学習済みの重みを読み込んでいます.2つめのブロックでは,世界座標系中の点とその点を覗き込むレイの角度を数多く定義しています.3つめのブロックにおいて,ここで数多く定義した,点とその角度をNeRFに入力し,その位置に物体表面座標が存在する密度を計算します. 4つめのブロックで示しているのは,あるレイの線上でもっとも大きい値を持つ密度の点を抽出する処理です.これにより,一つのレイが物体表面にぶつかった(と考えられる)点,すなわち物体表面の座標を得ることができます.ここまでで,変数sigmaに世界座標系での点群座標が格納されます. 最後のブロックでは,sigmaの分布を描画しています.分布は以下のようになります.

密度σの分布

ここまでできると,点群からメッシュデータに変換するだけなので基本的にはお好みの方法を使うことができます.ここでは著者らの実装にそってマーチングキューブ法[3][4]でメッシュ化します.コードは極めてシンプルで以下に示すものです.

import mcubes

threshold = 25
vertices, triangles = mcubes.marching_cubes(sigma, threshold)
mcubes.export_obj(vertices, triangles, "lego.obj")

先の図で示した通り,今回の例ではsigmaはおよそ0~200の間で分布していますが,値の小さい点は物体表面の点の推定として確信度が小さく,ノイズが多く含まれる可能性があるので,しきい値thresholdを設けてそういった点をカットしてからメッシュにしています.

こうして出来上がったメッシュ(ポリゴンデータ)は以下のようになります.

NeRFから出力したメッシュデータ

今回は時短のため学習ステップ数を小さくしたので学習が収束しきっておらず,地面より下にノイズが出ているのとディテールは出ていませんがレゴのおもちゃとして十分認識できる形になっているかと思います.

まとめ

本記事では学習済みのNeRFから3Dメッシュデータを抽出する方法をご紹介しました.NeRFは場を学習するニューラルネットワークとして最近注目されている技術であり,今後様々な方面で研究がなされ実用化が進んでいくことが予想されます.興味のある方は継続的に研究の動向に注目していくと面白いと思います.

さいごに

Ridge-iでは様々なポジションで積極採用中です. カジュアル面談も可能ですのでご興味がある方は是非ご連絡ください.

ridge-i.com

参考文献

[1] Ben Mildenhall, Pratul P. Srinivasan, Matthew Tancik, Jonathan T. Barron, Ravi Ramamoorthi and Ren Ng, "NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis," In 2020, ECCV.

[2] https://github.com/bmild/nerf

[3] William E. Lorensen and Harvey E. Cline, "Marching cubes: A high resolution 3D surface construction algorithm," In 1987, SIGGRAPH.

[4] https://github.com/pmneila/PyMCubes