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

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

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

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

02??NeRF是如何工作的

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

下面是如何從一個(gè)特定的視點(diǎn)生成一個(gè)NeRF:

通過移動(dòng)攝像機(jī)光線穿過場(chǎng)景生成一組采樣的3D點(diǎn)

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

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

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

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

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

 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)有了一個(gè)可以從3D模型生成2D圖像的可微分模擬器，下面就是實(shí)現(xiàn)NeRF模型。 根據(jù)上面的介紹，NeRF非常的復(fù)雜，但實(shí)際上NeRF模型只是多層感知器(MLPs)。但是具有ReLU激活函數(shù)的mlp傾向于學(xué)習(xí)低頻信號(hào)。當(dāng)試圖用高頻特征建模物體和場(chǎng)景時(shí)，這就出現(xiàn)了一個(gè)問題。為了抵消這種偏差并允許模型學(xué)習(xí)高頻信號(hào)，使用位置編碼將神經(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)單，因?yàn)樗彩潜O(jiān)督學(xué)習(xí)。我們可以直接最小化預(yù)測(cè)顏色和實(shí)際顏色之間的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 測(cè)試 訓(xùn)練過程完成，NeRF模型就可以用于從任何角度生成圖像。測(cè)試函數(shù)通過使用來自測(cè)試圖像的射線數(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)

這樣一個(gè)簡(jiǎn)單的NeRF就完成了,看看效果：

編輯：黃飛

?