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

該方法是一個基于體素重建的方法，通過在多幅圖片中的五維坐標(biāo)建立一個由粗到細(xì)的對應(yīng)，進(jìn)而恢復(fù)出原始的三維體素場景。

01???NeRF 和神經(jīng)渲染的基本概念

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

02??NeRF是如何工作的

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

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

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

將采樣點及其相應(yīng)的2D觀察方向輸入神經(jīng)網(wǎng)絡(luò)，生成密度和顏色的輸出集

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

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

03??Pytorch代碼實現(xiàn)

3.1 渲染 神經(jīng)輻射場的一個關(guān)鍵組件，是一個可微分渲染，它將由NeRF模型表示的3D表示映射到2D圖像。該問題可以表述為一個簡單的重構(gòu)問題：，這里的A是可微渲染，x是NeRF模型，b是目標(biāo)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模型和來自相機(jī)的一些光線作為輸入，并使用立體渲染返回與每個光線相關(guān)的顏色。 代碼的初始部分使用分層采樣沿射線選擇3D點。然后在這些點上查詢神經(jīng)輻射場模型(連同射線方向)以獲得密度和顏色信息。模型的輸出可以用蒙特卡羅積分計算每條射線的線積分。 累積透射率(論文中Ti)用下面的專用函數(shù)中單獨計算。

 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

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

?

 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 訓(xùn)練

訓(xùn)練循環(huán)也很簡單，因為它也是監(jiān)督學(xué)習(xí)。我們可以直接最小化預(yù)測顏色和實際顏色之間的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 測試 訓(xùn)練過程完成，NeRF模型就可以用于從任何角度生成圖像。測試函數(shù)通過使用來自測試圖像的射線數(shù)據(jù)集進(jìn)行操作，然后使用渲染函數(shù)和優(yōu)化的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就完成了,看看效果：

編輯：黃飛

?