提起三維重建技術，NeRF是一個絕對繞不過去的名字。這項逆天的技術，一經提出就被眾多研究者所重視，對該技術進行深入研究并提出改進已經成為一個熱點。不到兩年的時間，NeRF及其變種已經成為重建領域的主流。本文通過100行的Pytorch代碼實現最初的 NeRF 論文。 NeRF全稱為Neural Radiance Fields(神經輻射場)，是一項利用多目圖像重建三維場景的技術。該項目的作者來自于加州大學伯克利分校，Google研究院，以及加州大學圣地亞哥分校。NeRF使用一組多目圖作為輸入，通過優化一個潛在連續的體素場景方程來得到一個完整的三維場景。該方法使用一個全連接深度網絡來表示場景，使用的輸入是一個單連通的5D坐標(空間位置x,y,z以及觀察視角θ，),輸出為一個體素場景，可以以任意視角查看，并通過體素渲染技術，生成需要視角的照片。該方法同樣支持視頻合成。

該方法是一個基于體素重建的方法，通過在多幅圖片中的五維坐標建立一個由粗到細的對應，進而恢復出原始的三維體素場景。

01???NeRF 和神經渲染的基本概念

1.1 Rendering 渲染是從 3D 模型創建圖像的過程。該模型將包含紋理、陰影、陰影、照明和視點等特征，渲染引擎的作用是處理這些特征以創建逼真的圖像。 三種常見的渲染算法類型是光柵化，它根據模型中的信息以幾何方式投影對象，沒有光學效果；光線投射，使用基本的光學反射定律從特定角度計算圖像；和光線追蹤，它使用蒙特卡羅技術在更短的時間內獲得逼真的圖像。光線追蹤用于提高 NVIDIA GPU 中的渲染性能。 1.2 Volume Rendering 立體渲染使能夠創建 3D 離散采樣數據集的 2D 投影。 對于給定的相機位置，立體渲染算法為空間中的每個體素獲取 RGBα（紅色、綠色、藍色和 Alpha 通道），相機光線通過這些體素投射。RGBα 顏色轉換為 RGB 顏色并記錄在 2D 圖像的相應像素中。對每個像素重復該過程，直到呈現整個 2D 圖像。 1.3 View Synthesis 視圖合成與立體渲染相反——它涉從一系列 2D 圖像創建 3D 視圖。這可以使用一系列從多個角度顯示對象的照片來完成，創建對象的半球平面圖，并將每個圖像放置在對象周圍的適當位置。視圖合成函數嘗試在給定一系列描述對象不同視角的圖像的情況下預測深度。

02??NeRF是如何工作的

NeRF使用一組稀疏的輸入視圖來優化連續的立體場景函數。這種優化的結果是能夠生成復雜場景的新視圖。 NeRF使用一組多目圖作為輸入: 輸入為一個單連通的5D坐標(空間位置x,y,z以及觀察視角(θ; Φ) 輸出為一個體素場景 c = (r; g; b) 和體積密度 (α)。

下面是如何從一個特定的視點生成一個NeRF:

通過移動攝像機光線穿過場景生成一組采樣的3D點

將采樣點及其相應的2D觀察方向輸入神經網絡，生成密度和顏色的輸出集

通過使用經典的立體渲染技術，將密度和顏色累積到2D圖像中

上述過程深度的全連接、多層感知器(MLP)進行優化，并且不需要使用卷積層。它使用梯度下降來最小化每個觀察到的圖像和從表示中呈現的所有相應視圖之間的誤差。

03??Pytorch代碼實現

3.1 渲染 神經輻射場的一個關鍵組件，是一個可微分渲染，它將由NeRF模型表示的3D表示映射到2D圖像。該問題可以表述為一個簡單的重構問題：，這里的A是可微渲染，x是NeRF模型，b是目標2D圖像。 代碼如下：

 def render_rays(nerf_model, ray_origins, ray_directions, hn=0, hf=0.5, nb_bins=192):
     device = ray_origins.device
     t = torch.linspace(hn, hf, nb_bins, device=device).expand(ray_origins.shape[0], nb_bins)
     # Perturb sampling along each ray.
     mid = (t[:, :-1] + t[:, 1:]) / 2.
     lower = torch.cat((t[:, :1], mid), -1)
     upper = torch.cat((mid, t[:, -1:]), -1)
     u = torch.rand(t.shape, device=device)
     t = lower + (upper - lower) * u  # [batch_size, nb_bins]
     delta = torch.cat((t[:, 1:] - t[:, :-1], torch.tensor([1e10], device=device).expand(ray_origins.shape[0], 1)), -1)
 
     x = ray_origins.unsqueeze(1) + t.unsqueeze(2) * ray_directions.unsqueeze(1)   # [batch_size, nb_bins, 3]
     ray_directions = ray_directions.expand(nb_bins, ray_directions.shape[0], 3).transpose(0, 1)
 
     colors, sigma = nerf_model(x.reshape(-1, 3), ray_directions.reshape(-1, 3))
     colors = colors.reshape(x.shape)
     sigma = sigma.reshape(x.shape[:-1])
 
     alpha = 1 - torch.exp(-sigma * delta)  # [batch_size, nb_bins]
     weights = compute_accumulated_transmittance(1 - alpha).unsqueeze(2) * alpha.unsqueeze(2)
     c = (weights * colors).sum(dim=1)  # Pixel values
     weight_sum = weights.sum(-1).sum(-1)  # Regularization for white background
     return c + 1 - weight_sum.unsqueeze(-1)

渲染將NeRF模型和來自相機的一些光線作為輸入，并使用立體渲染返回與每個光線相關的顏色。 代碼的初始部分使用分層采樣沿射線選擇3D點。然后在這些點上查詢神經輻射場模型(連同射線方向)以獲得密度和顏色信息。模型的輸出可以用蒙特卡羅積分計算每條射線的線積分。 累積透射率(論文中Ti)用下面的專用函數中單獨計算。

 def compute_accumulated_transmittance(alphas):
     accumulated_transmittance = torch.cumprod(alphas, 1)
     return torch.cat((torch.ones((accumulated_transmittance.shape[0], 1), device=alphas.device),
                       accumulated_transmittance[:, :-1]), dim=-1)

?

?

3.2 NeRF

我們已經有了一個可以從3D模型生成2D圖像的可微分模擬器，下面就是實現NeRF模型。 根據上面的介紹，NeRF非常的復雜，但實際上NeRF模型只是多層感知器(MLPs)。但是具有ReLU激活函數的mlp傾向于學習低頻信號。當試圖用高頻特征建模物體和場景時，這就出現了一個問題。為了抵消這種偏差并允許模型學習高頻信號，使用位置編碼將神經網絡的輸入映射到高維空間。

?

 class NerfModel(nn.Module):
     def __init__(self, embedding_dim_pos=10, embedding_dim_direction=4, hidden_dim=128):
         super(NerfModel, self).__init__()
 
         self.block1 = nn.Sequential(nn.Linear(embedding_dim_pos * 6 + 3, hidden_dim), nn.ReLU(),
                                     nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
                                     nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
                                     nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), )
 
         self.block2 = nn.Sequential(nn.Linear(embedding_dim_pos * 6 + hidden_dim + 3, hidden_dim), nn.ReLU(),
                                     nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
                                     nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
                                     nn.Linear(hidden_dim, hidden_dim + 1), )
 
         self.block3 = nn.Sequential(nn.Linear(embedding_dim_direction * 6 + hidden_dim + 3, hidden_dim // 2), nn.ReLU(), )
         self.block4 = nn.Sequential(nn.Linear(hidden_dim // 2, 3), nn.Sigmoid(), )
 
         self.embedding_dim_pos = embedding_dim_pos
         self.embedding_dim_direction = embedding_dim_direction
         self.relu = nn.ReLU()
 
     @staticmethod
     def positional_encoding(x, L):
         out = [x]
         for j in range(L):
             out.append(torch.sin(2 ** j * x))
             out.append(torch.cos(2 ** j * x))
         return torch.cat(out, dim=1)
 
     def forward(self, o, d):
         emb_x = self.positional_encoding(o, self.embedding_dim_pos)
         emb_d = self.positional_encoding(d, self.embedding_dim_direction)
         h = self.block1(emb_x)
         tmp = self.block2(torch.cat((h, emb_x), dim=1))
         h, sigma = tmp[:, :-1], self.relu(tmp[:, -1])
         h = self.block3(torch.cat((h, emb_d), dim=1))
         c = self.block4(h)
         return c, sigma

?

?

3.3 訓練

訓練循環也很簡單，因為它也是監督學習。我們可以直接最小化預測顏色和實際顏色之間的L2損失。

?

 def train(nerf_model, optimizer, scheduler, data_loader, device='cpu', hn=0, hf=1, nb_epochs=int(1e5),
           nb_bins=192, H=400, W=400):
     training_loss = []
     for _ in tqdm(range(nb_epochs)):
         for batch in data_loader:
             ray_origins = batch[:, :3].to(device)
             ray_directions = batch[:, 3:6].to(device)
             ground_truth_px_values = batch[:, 6:].to(device)
 
             regenerated_px_values = render_rays(nerf_model, ray_origins, ray_directions, hn=hn, hf=hf, nb_bins=nb_bins)
             loss = ((ground_truth_px_values - regenerated_px_values) ** 2).sum()
 
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
             training_loss.append(loss.item())
         scheduler.step()
 
         for img_index in range(200):
             test(hn, hf, testing_dataset, img_index=img_index, nb_bins=nb_bins, H=H, W=W)
 
     return training_loss

3.4 測試 訓練過程完成，NeRF模型就可以用于從任何角度生成圖像。測試函數通過使用來自測試圖像的射線數據集進行操作，然后使用渲染函數和優化的NeRF模型為這些射線生成圖像。

 @torch.no_grad()
 def test(hn, hf, dataset, chunk_size=10, img_index=0, nb_bins=192, H=400, W=400):
     ray_origins = dataset[img_index * H * W: (img_index + 1) * H * W, :3]
     ray_directions = dataset[img_index * H * W: (img_index + 1) * H * W, 3:6]
 
     data = []
     for i in range(int(np.ceil(H / chunk_size))):
         ray_origins_ = ray_origins[i * W * chunk_size: (i + 1) * W * chunk_size].to(device)
         ray_directions_ = ray_directions[i * W * chunk_size: (i + 1) * W * chunk_size].to(device)
 
         regenerated_px_values = render_rays(model, ray_origins_, ray_directions_, hn=hn, hf=hf, nb_bins=nb_bins)
         data.append(regenerated_px_values)
     img = torch.cat(data).data.cpu().numpy().reshape(H, W, 3)
 
     plt.figure()
     plt.imshow(img)
     plt.savefig(f'novel_views/img_{img_index}.png', bbox_inches='tight')
     plt.close()

所有的部分都可以很容易地組合起來。

 if __name__ == 'main':
     device = 'cuda'
     training_dataset = torch.from_numpy(np.load('training_data.pkl', allow_pickle=True))
     testing_dataset = torch.from_numpy(np.load('testing_data.pkl', allow_pickle=True))
     model = NerfModel(hidden_dim=256).to(device)
     model_optimizer = torch.optim.Adam(model.parameters(), lr=5e-4)
     scheduler = torch.optim.lr_scheduler.MultiStepLR(model_optimizer, milestones=[2, 4, 8], gamma=0.5)
 
     data_loader = DataLoader(training_dataset, batch_size=1024, shuffle=True)
     train(model, model_optimizer, scheduler, data_loader, nb_epochs=16, device=device, hn=2, hf=6, nb_bins=192, H=400,
           W=400)

這樣一個簡單的NeRF就完成了,看看效果：

編輯：黃飛

?