关于自动驾驶:100行Pytorch代码实现三维重建技术神经辐射场-NeRF

转载自：DeepHub IMBA
原文：100 行 Pytorch 代码实现三维重建技术神经辐射场 (NeRF)

提起三维重建技术，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)−bbA(x)−bb||2，这里的 AA 可微渲染，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)

代码的初始局部应用分层采样沿射线抉择 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 就实现了, 看看成果：

心愿本文对你有所帮忙，如果你对 NeRF 感兴趣能够看看这个我的项目：https://github.com/kwea123/ne…

1. 神经渲染技术在主动驾驶畛域利用
2. 六问 Nerf | 简略易懂的神经辐射场入门介绍
3. 事实虚拟化：从三维重建到逆渲染（Inverse Rendering）
4. 下载 | 101 篇 NeRF 优质论文举荐（含我的项目代码）
5. 精髓文稿|在非现实输出下 NeRF 的重建