Weekly Machine Learning #169 で紹介されていたNeRF がとても興味深く、

また、簡単に触ってみることができる状態でコードが公開されていたので、実際に触りつつ論文を読んでみました。

この論文では、いくつかの視点から撮影した物体の画像を再現するモデルを学習させます。

撮影する点 (x, y, z) と見る方向(θ, φ) の情報を入力として渡し、その視点、見る方向からどのような像が見えるか？という画像を得ることができます。

写真を何枚も撮影したんだけど、ちょうどこの角度の画像だけ撮ってなかった！！！みたいなときに使えそうです。

自分が面白いと感じたのは、

画像を扱うのにCNN を使っていない
視点と見る角度を入力し、RGBと物体の密度を予測する
学習には幾何学的なデータセットも使用したが、独自で作成したデータセットも使用した

です。

自分は、GAN や U-Net 以外の画像生成系のモデルを知らなかったので、こんなモデルも有るのか！と衝撃を受けました。

モデルの構造は、以下のようになっています。1層目から dense layer になっています。 tf_op_layer_concat で、 input を concat しているのがResidual Network みたいで面白いです。最も、ResNet の場合は、concat ではなく足し合わせですが。

Model: "model"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(None, 39)]         0                                            
__________________________________________________________________________________________________
dense (Dense)                   (None, 256)          10240       input_1[0][0]                    
__________________________________________________________________________________________________
dense_1 (Dense)                 (None, 256)          65792       dense[0][0]                      
__________________________________________________________________________________________________
dense_2 (Dense)                 (None, 256)          65792       dense_1[0][0]                    
__________________________________________________________________________________________________
dense_3 (Dense)                 (None, 256)          65792       dense_2[0][0]                    
__________________________________________________________________________________________________
dense_4 (Dense)                 (None, 256)          65792       dense_3[0][0]                    
__________________________________________________________________________________________________
tf_op_layer_concat (TensorFlowO [(None, 295)]        0           dense_4[0][0]                    
                                                                 input_1[0][0]                    
__________________________________________________________________________________________________
dense_5 (Dense)                 (None, 256)          75776       tf_op_layer_concat[0][0]         
__________________________________________________________________________________________________
dense_6 (Dense)                 (None, 256)          65792       dense_5[0][0]                    
__________________________________________________________________________________________________
dense_7 (Dense)                 (None, 256)          65792       dense_6[0][0]                    
__________________________________________________________________________________________________
dense_8 (Dense)                 (None, 4)            1028        dense_7[0][0]                    
==================================================================================================
Total params: 481,796
Trainable params: 481,796
Non-trainable params: 0
__________________________________________________________________________________________________

ネットワークの詳細は、p.5 に書いてあります。概要を記述します。

入力は、 (x, y, z, θ, φ) の5次元ベクトルで、 (x, y, z) は視点の位置表し、(θ, φ) は、見ている方向を表しています。  
これらを入力し、 (c, σ) を出力します。 c = (r, g, b), σ は物体の密度となっています。  
全結合層が役に立つのは、 (x, y, z, θ, φ) => (c, σ) への変換です。  
変換を行ったあと、物体の密度とRGBの値を参考に、画像を層状にレンダリングしていきます。

今回参考にしたNotebookは、tiny_nerf ということで、ネットワークの構造が少し論文に記述されているものと違いますが、論文では、視点の情報を入れることにより、視点情報がない場合には再現が難しかった反射を再現することが可能になったと報告されています。

論文中 Fig.3 は、視点をずらしたときに、同じ点の見え方がどのように移り変わるかを表した図になっており、光の反射具合まで考慮できていることがわかります。

nerf_paper_fig_3

モデルを実際に手元で動かしてみるのはとても簡単です。 Colab のリンクを開くと特にエラーもなく実行できます。

ちなみに、入力は、レゴを別の位置から見た画像と、その姿勢パラメータのようです。データは100セット(学習ではそのうち64個を使用)あり、デフォルトでは 1000 エポック回しています。i_plot 等、いくつかのパラメータをいじることで学習の進捗や結果の可視化方法を調整できます。

この実装の肝は、こちらです。

images = data['images']
# 4x4 の配列
poses = data['poses']
# スカラー。デフォルトでは 138.88887889922103 
focal = data['focal']
# 高さ 幅
H, W = images.shape[1:3]

rays_o, rays_d = get_rays(H, W, focal, pose)
with tf.GradientTape() as tape:
  
    rgb, depth, acc = render_rays(model, rays_o, rays_d, near=2., far=6., N_samples=N_samples, rand=True)
    loss = tf.reduce_mean(tf.square(rgb - target))
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))

特に、get_rays と render_rays が重要になっています。最初に面白いな！と思ったのは、モデルへの入力パラメータを作成する get_rays というメソッドに元画像を入れていないことです。ここが、GANっぽい。だけど、GANとも違う。

get_raysとrender_raysの中はまだ全て追えていませんが、論文中で書いてある式が愚直に実装されているような感じでした。実際にColabを実行してみると、学習の過程が見れたりインタラクティブなデモもついていたりしてとても面白いです。たったこれだけのコードで！と驚かれると思います。ぜひ一度お試しください！

また、読んでいて躓いたところについてもご紹介します。パット見た感じ、論文内の式を愚直に実装しているようなのですが、見たことのないメソッドを使ってたりするので(tf.meshgrid とか)。。。

ご存じの方は、読み飛ばしてください。\ \ コンテンツは以下

...
スライスの中のnp.newaxis や None
tf.meshgrid
tf.math.cumprod

...

何箇所か、 ... という書き方が出てきます。

import numpy as np
x = np.arange(8).reshape((2, 2, 2))
x.shape #=> (2, 2, 2)
x #=> array([
  #     [[0, 1],[2, 3]],
  #     [[4, 5],[6, 7]],
  #   ])
x[..., 1] #=> 
  # array([
  #  [1, 3],[5, 7]
  # ])
x[:, :, 1] #=> 
    # array([
    #  [1, 3],[5, 7]
    # ])

こちら、読み方は、Elipsisというものらしいです。x[..., 1] の意味は、x[:, :, 1] だそうです。: をたくさん書かなくて良いから便利！

ちなみに、 x[...,1,...] みたいな書き方はできません。次元が決まらないからでしょう。2x2x2 行列の場合、上記の書き方でも行けそうな気がしたのですが。。。

スライスの中のnp.newaxis や None

配列に次元を追加する書き方みたいです。論文中に、dirs[..., np.newaxis, :] みたいな書き方や rays_d[...,None,:] みたいな書き方が出てきます。

これらは同じ意味で、

import numpy as np
x = np.arange(8).reshape((2, 2, 2))
x[..., np.newaxis, :]
#=> array([
# [
#   [[0, 1]],
#   [[2, 3]]
# ],
# [
#   [[4, 5]],
#   [[6, 7]]]
# ])
x[..., np.newaxis, :].shape
#=> (2, 2, 1, 2)
x[..., None, :].shape
#=> (2, 2, 1, 2)

という感じらしいです。 reshapeでも同じことができますが、こちらのほうが直感的に書ける場合もあるなと感じました。

tf.meshgrid

中身が同じ行列を作成する方法みたいです。numpy にも同じメソッドがあるらしく、グラフの描画などに便利だそうです。知らなかった。。。

これ、説明が難しいので、以下の例を御覧くださいませ。

import numpy as np
import tensorflow as tf

a = np.arange(9).reshape((-3, 3))
b = np.arange(6).reshape((-3, 2)) + 100
c, d = tf.meshgrid(a, b)
c #=>
# <tf.Tensor: shape=(6, 9), dtype=int64, numpy=
# array([[0, 1, 2, 3, 4, 5, 6, 7, 8],
# [0, 1, 2, 3, 4, 5, 6, 7, 8],
# [0, 1, 2, 3, 4, 5, 6, 7, 8],
# [0, 1, 2, 3, 4, 5, 6, 7, 8],
# [0, 1, 2, 3, 4, 5, 6, 7, 8],
# [0, 1, 2, 3, 4, 5, 6, 7, 8]])>
d #=>
# <tf.Tensor: shape=(6, 9), dtype=int64, numpy=
# array([[100, 100, 100, 100, 100, 100, 100, 100, 100],
# [101, 101, 101, 101, 101, 101, 101, 101, 101],
# [102, 102, 102, 102, 102, 102, 102, 102, 102],
# [103, 103, 103, 103, 103, 103, 103, 103, 103],
# [104, 104, 104, 104, 104, 104, 104, 104, 104],
# [105, 105, 105, 105, 105, 105, 105, 105, 105]])>

tf.math.cumprod

累積積を計算するメソッドみたいです。こちらも numpy にあるらしいです。全然知りませんでした。。。。

一応注意ですが、TensorFlow と numpy で動作が違うので注意が必要です。

a #=> array([[0, 1, 2],
# [3, 4, 5],
# [6, 7, 8]])

tf.math.cumprod(a)
#=> <tf.Tensor: shape=(3, 3), dtype=int64, numpy=
# array([[ 0,  1,  2],
#        [ 0,  4, 10],
#        [ 0, 28, 80]])>

np.cumprod(a)
#=> array([0, 0, 0, 0, 0, 0, 0, 0, 0])

以上です。とても面白かったので、実装が公開されている論文があったらまた記事を書こうと思います。

Fusic Tech Blog

# 論文

[論文読み] NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis