100行Pytorch代码实现三维重建技术神经辐射场 (NeRF)

简介: 提起三维重建技术,NeRF是一个绝对绕不过去的名字。这项逆天的技术,一经提出就被众多研究者所重视,对该技术进行深入研究并提出改进已经成为一个热点。不到两年的时间,NeRF及其变种已经成为重建领域的主流。本文通过100行的Pytorch代码实现最初的 NeRF 论文。

NeRF全称为Neural Radiance Fields(神经辐射场),是一项利用多目图像重建三维场景的技术。该项目的作者来自于加州大学伯克利分校,Google研究院,以及加州大学圣地亚哥分校。NeRF使用一组多目图作为输入,通过优化一个潜在连续的体素场景方程来得到一个完整的三维场景。该方法使用一个全连接深度网络来表示场景,使用的输入是一个单连通的5D坐标(空间位置x,y,z以及观察视角θ,),输出为一个体素场景,可以以任意视角查看,并通过体素渲染技术,生成需要视角的照片。该方法同样支持视频合成。

该方法是一个基于体素重建的方法,通过在多幅图片中的五维坐标建立一个由粗到细的对应,进而恢复出原始的三维体素场景。

NeRF 和神经渲染的基本概念

Rendering

渲染是从 3D 模型创建图像的过程。该模型将包含纹理、阴影、阴影、照明和视点等特征,渲染引擎的作用是处理这些特征以创建逼真的图像。

三种常见的渲染算法类型是光栅化,它根据模型中的信息以几何方式投影对象,没有光学效果;光线投射,使用基本的光学反射定律从特定角度计算图像;和光线追踪,它使用蒙特卡罗技术在更短的时间内获得逼真的图像。光线追踪用于提高 NVIDIA GPU 中的渲染性能。

Volume Rendering

立体渲染使能够创建 3D 离散采样数据集的 2D 投影。

对于给定的相机位置,立体渲染算法为空间中的每个体素获取 RGBα(红色、绿色、蓝色和 Alpha 通道),相机光线通过这些体素投射。RGBα 颜色转换为 RGB 颜色并记录在 2D 图像的相应像素中。对每个像素重复该过程,直到呈现整个 2D 图像。

View Synthesis

视图合成与立体渲染相反——它涉从一系列 2D 图像创建 3D 视图。这可以使用一系列从多个角度显示对象的照片来完成,创建对象的半球平面图,并将每个图像放置在对象周围的适当位置。视图合成函数尝试在给定一系列描述对象不同视角的图像的情况下预测深度。

NeRF是如何工作的

NeRF使用一组稀疏的输入视图来优化连续的立体场景函数。这种优化的结果是能够生成复杂场景的新视图。

NeRF使用一组多目图作为输入:

输入为一个单连通的5D坐标(空间位置x,y,z以及观察视角(θ; Φ)

输出为一个体素场景 c = (r; g; b) 和体积密度 (α)。

下面是如何从一个特定的视点生成一个NeRF:

  • 通过移动摄像机光线穿过场景生成一组采样的3D点
  • 将采样点及其相应的2D观察方向输入神经网络,生成密度和颜色的输出集
  • 通过使用经典的立体渲染技术,将密度和颜色累积到2D图像中

上述过程深度的全连接、多层感知器(MLP)进行优化,并且不需要使用卷积层。它使用梯度下降来最小化每个观察到的图像和从表示中呈现的所有相应视图之间的误差。

Pytorch代码实现

渲染

神经辐射场的一个关键组件,是一个可微分渲染,它将由NeRF模型表示的3D表示映射到2D图像。该问题可以表述为一个简单的重构问题

这里的A是可微渲染,x是NeRF模型,b是目标2D图像。

代码如下:

 defrender_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
     returnc+1-weight_sum.unsqueeze(-1)

渲染将NeRF模型和来自相机的一些光线作为输入,并使用立体渲染返回与每个光线相关的颜色。

代码的初始部分使用分层采样沿射线选择3D点。然后在这些点上查询神经辐射场模型(连同射线方向)以获得密度和颜色信息。模型的输出可以用蒙特卡罗积分计算每条射线的线积分。

累积透射率(论文中Ti)用下面的专用函数中单独计算。

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

NeRF

我们已经有了一个可以从3D模型生成2D图像的可微分模拟器,下面就是实现NeRF模型。

根据上面的介绍,NeRF非常的复杂,但实际上NeRF模型只是多层感知器(MLPs)。但是具有ReLU激活函数的mlp倾向于学习低频信号。当试图用高频特征建模物体和场景时,这就出现了一个问题。为了抵消这种偏差并允许模型学习高频信号,使用位置编码将神经网络的输入映射到高维空间。

 classNerfModel(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
     defpositional_encoding(x, L):
         out= [x]
         forjinrange(L):
             out.append(torch.sin(2**j*x))
             out.append(torch.cos(2**j*x))
         returntorch.cat(out, dim=1)
 
     defforward(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)
         returnc, sigma

训练

训练循环也很简单,因为它也是监督学习。我们可以直接最小化预测颜色和实际颜色之间的L2损失。

 deftrain(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_intqdm(range(nb_epochs)):
         forbatchindata_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()
 
         forimg_indexinrange(200):
             test(hn, hf, testing_dataset, img_index=img_index, nb_bins=nb_bins, H=H, W=W)
 
     returntraining_loss

测试

训练过程完成,NeRF模型就可以用于从任何角度生成图像。测试函数通过使用来自测试图像的射线数据集进行操作,然后使用渲染函数和优化的NeRF模型为这些射线生成图像。

 @torch.no_grad()
 deftest(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= []
     foriinrange(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://avoid.overfit.cn/post/3d89b7ed625b437993e3fde57f36c70a

目录
相关文章
|
18天前
|
存储 物联网 PyTorch
基于PyTorch的大语言模型微调指南:Torchtune完整教程与代码示例
**Torchtune**是由PyTorch团队开发的一个专门用于LLM微调的库。它旨在简化LLM的微调流程,提供了一系列高级API和预置的最佳实践
125 59
基于PyTorch的大语言模型微调指南:Torchtune完整教程与代码示例
|
2月前
|
机器学习/深度学习 PyTorch 算法框架/工具
CNN中的注意力机制综合指南:从理论到Pytorch代码实现
注意力机制已成为深度学习模型的关键组件,尤其在卷积神经网络(CNN)中发挥了重要作用。通过使模型关注输入数据中最相关的部分,注意力机制显著提升了CNN在图像分类、目标检测和语义分割等任务中的表现。本文将详细介绍CNN中的注意力机制,包括其基本概念、不同类型(如通道注意力、空间注意力和混合注意力)以及实际实现方法。此外,还将探讨注意力机制在多个计算机视觉任务中的应用效果及其面临的挑战。无论是图像分类还是医学图像分析,注意力机制都能显著提升模型性能,并在不断发展的深度学习领域中扮演重要角色。
92 10
|
11天前
|
机器学习/深度学习 监控 PyTorch
深度学习工程实践:PyTorch Lightning与Ignite框架的技术特性对比分析
在深度学习框架的选择上,PyTorch Lightning和Ignite代表了两种不同的技术路线。本文将从技术实现的角度,深入分析这两个框架在实际应用中的差异,为开发者提供客观的技术参考。
32 7
|
5月前
|
机器学习/深度学习 PyTorch 算法框架/工具
【从零开始学习深度学习】26.卷积神经网络之AlexNet模型介绍及其Pytorch实现【含完整代码】
【从零开始学习深度学习】26.卷积神经网络之AlexNet模型介绍及其Pytorch实现【含完整代码】
|
5月前
|
机器学习/深度学习 PyTorch 算法框架/工具
【从零开始学习深度学习】28.卷积神经网络之NiN模型介绍及其Pytorch实现【含完整代码】
【从零开始学习深度学习】28.卷积神经网络之NiN模型介绍及其Pytorch实现【含完整代码】
|
1月前
|
机器学习/深度学习 PyTorch 算法框架/工具
聊一聊计算机视觉中常用的注意力机制以及Pytorch代码实现
本文介绍了几种常用的计算机视觉注意力机制及其PyTorch实现,包括SENet、CBAM、BAM、ECA-Net、SA-Net、Polarized Self-Attention、Spatial Group-wise Enhance和Coordinate Attention等,每种方法都附有详细的网络结构说明和实验结果分析。通过这些注意力机制的应用,可以有效提升模型在目标检测任务上的性能。此外,作者还提供了实验数据集的基本情况及baseline模型的选择与实验结果,方便读者理解和复现。
27 0
聊一聊计算机视觉中常用的注意力机制以及Pytorch代码实现
|
1月前
|
机器学习/深度学习 算法 数据可视化
如果你的PyTorch优化器效果欠佳,试试这4种深度学习中的高级优化技术吧
在深度学习领域,优化器的选择对模型性能至关重要。尽管PyTorch中的标准优化器如SGD、Adam和AdamW被广泛应用,但在某些复杂优化问题中,这些方法未必是最优选择。本文介绍了四种高级优化技术:序列最小二乘规划(SLSQP)、粒子群优化(PSO)、协方差矩阵自适应进化策略(CMA-ES)和模拟退火(SA)。这些方法具备无梯度优化、仅需前向传播及全局优化能力等优点,尤其适合非可微操作和参数数量较少的情况。通过实验对比发现,对于特定问题,非传统优化方法可能比标准梯度下降算法表现更好。文章详细描述了这些优化技术的实现过程及结果分析,并提出了未来的研究方向。
26 1
|
4月前
|
机器学习/深度学习 数据采集 自然语言处理
注意力机制中三种掩码技术详解和Pytorch实现
**注意力机制中的掩码在深度学习中至关重要,如Transformer模型所用。掩码类型包括:填充掩码(忽略填充数据)、序列掩码(控制信息流)和前瞻掩码(自回归模型防止窥视未来信息)。通过创建不同掩码,如上三角矩阵,模型能正确处理变长序列并保持序列依赖性。在注意力计算中,掩码修改得分,确保模型学习的有效性。这些技术在现代NLP和序列任务中是核心组件。**
177 12
|
3月前
|
机器学习/深度学习 PyTorch 算法框架/工具
PyTorch代码实现神经网络
这段代码示例展示了如何在PyTorch中构建一个基础的卷积神经网络(CNN)。该网络包括两个卷积层,分别用于提取图像特征,每个卷积层后跟一个池化层以降低空间维度;之后是三个全连接层,用于分类输出。此结构适用于图像识别任务,并可根据具体应用调整参数与层数。
|
3月前
|
数据可视化 数据管理 vr&ar