Skip to content

NeRF学习 —— 基于pytorch代码复现的笔记

2663字约9分钟

3DVisionNeRF3D/4D Reconstruction

2024-03-24

1 复现笔记

1 创建Dateset

核心函数包括:initgetitemlen

  • init函数负责从磁盘中load指定格式的文件,计算并存储为特定形式

  • getitem函数负责在运行时提供给网络一次训练需要的输入,以及 groundtruth 的输出

    例如对NeRF,分别是1024条 rays 以及1024个 RGB 值

  • len函数是训练或者测试的数量:getitem函数获得的index值通常是[0, len-1]

1.1 __init__

这里我们主要要做3件事(公平公平还是tmd公平

  1. 加载图像数据和相机位姿信息:

    • 从磁盘中读取指定格式的文件,这些文件包含了图像路径和相机位姿信息,位姿保存在 self.poses

    • 从图片路径保存图片,并根据参数进行处理(白色背景+缩放比例)

    • 最后,将处理后的图像添加到self.images列表中

  2. 计算相机内参:

    • 根据相机的视场角(FOV)和图像尺寸计算相机的焦距和相机矩阵
  3. 生成射线数据:

    • 首先,对于每一个位姿,使用rh.get_rays_np函数生成射线,并将所有射线堆叠起来

    • 然后,将射线和图像数据合并,并进行一系列的变换和重塑操作形成形状为 [N, H, W, ro+rd+rgb, 3] 的numpy数组 rays_rgb

      • N 是图像的数量

      • H 是图像的高度

      • W 是图像的宽度

      • ro+rd+rgb 表示每个射线的数据,包括射线的起始点(ro)、射线的方向(rd)和射线对应的像素颜色(rgb)

      • 3 表示每个数据都是一个三维向量(例如,rgb颜色是一个三维向量,包含红色、绿色和蓝色的强度)

      表示有 N*H*W 条射线,每条射线有3个数据(ro、rd、rgb),每个数据都是一个三维向量

    • 存储打乱前的射线 img_rays_rgb , 一张一张图片存储的,没有整合在一起 ,形状为 [N, H*W, ro+rd+rgb, 3],再转换为 Tensor

1.2 __getitem__

1.2.1 训练数据

主要功能是负责在运行时提供给网络一次训练需要的输入,以及 groundtruth 的输出 —— N_raysraysN_raysRGB值 ( N_rays = 1024)

  1. 将所有的光线(rays)和对应的RGB值整合到一个数组中,然后将这个数组的形状从 [N, H, W, ro+rd+rgb, 3] 改变为 [-1, 3, 3],即将所有的光线和RGB值平铺到一个一维数组中得到 [N*H*W, ro+rd+rgb, 3],并转为 Tensor

  2. 对所有的光线和RGB值进行打乱,这里使用GPU加速了一下

  3. 打乱后,根据索引 index 选择一个批次(1024条)的数据,再将批次的维度从第二位转换到第一位([1024,3,3] to [3,1024,3])

    在PyTorch中,通常希望批次维度在第一位,以便于进行批次处理

  4. 将选择的数据分为两部分:光线的数据(batch_rays)和对应的RGB值(target_s),并以字典的形式返回

    return {'H': self.H,
            'W': self.W,
            'K': self.K,
            'near': self.near,
            'far': self.far,
            'rays': batch_rays,
            'target_s': target_s}
1.2.2 测试数据

在测试时,输出一张图像的 rays 和 RGB 值,即 400 * 400 条 rays 和 400 * 400 个 RGB 值

  1. self.img_rays_rgb中根据索引index获取一张图片的光线

  2. 使用torch.transpose函数将批次的维度从第二位转换到第一位([1024,3,3] to [3,1024,3])

    在PyTorch中,通常希望批次维度在第一位,以便于进行批次处理

  3. 将数据分解为射线数据rays和对应的RGB值img_rgbrays包含了每条射线的起始点和方向img_rgb是每条射线对应的RGB值

  4. 返回一个字典,包含了射线数据和对应的RGB值,以及对应的图片索引 index

    return {'H': self.H,
            'W': self.W,
            'K': self.K,
            'near': self.near,
            'far': self.far,
            'rays': rays,
            'target_s': img_rgb,
            'index': index}

通过data_loader我们可以得到一个batch

  • train:

    image-20240428163506491
  • test:

    image-20240428163651114

1.3 __len__

返回数据集的大小,即数据集中的样本数量,即图像的数量

2 Network

核心函数包括:__init____forward__

  • init 函数负责定义网络所必需的模块

  • forward函数负责接收Dataset的输出,利用定义好的模块,计算输出

对于NeRF来说,我们需要在init中定义两个 MLP 以及 encoding 方式,在forward函数中,使用rays完成计算

定义了两个主要的神经网络模型:NeRFNetwork

2.1 NeRF 类

NeRF 类是NeRF的实现。它的结构包括了用于处理采样点和观察坐标方向的一系列MLP,以及相应的参数。主要方法是 forward,用于执行前向传播操作

  • 结构

    • __init__ 方法中定义了模型的各种参数和MLP

    • forward 方法实现了模型的前向传播逻辑,分别计算了辐射值和密度值,返回 output ——将 rgbalpha 两个张量在最后一个维度上进行拼接

      outputs = torch.cat([rgb, alpha], -1)

2.2 Network 类

Network 类是整个神经网络的集成,其中包括了对 NeRF 模型的调用以及render的实现,完成模型的调用和计算模型产生的预估值(即在这部分完成渲染)

  • 结构

    • __init__ 方法中初始化了整个网络模型,包括定义了两个 NeRF 模型(一个粗网络和一个精细网络),并根据配置参数进行设置

    • batchify 方法用于将函数应用于较小的批次数据

    • network_query_fn 方法,即网络查询函数,准备输入并将其传递给网络,并输出结果,使用 batchify 将一批一批的数据传入NeRF网络进行 forward

    • forward 方法,从 batch 中得到 dataloader 的数据,调用 render 函数:

      • 计算并创建射线,使用 render_utils 中的 batchify_rays 进行放入网络得到结果,再渲染射线得到 render_rays 返回的字典(即最后的结果):
      - 'rgb_map' (torch.Tensor): 每条光线的估计 RGB 颜色。形状:[num_rays, 3]
      - 'disp_map' (torch.Tensor): 每条光线的视差图(深度图的倒数)。形状:[num_rays]
      - 'acc_map' (torch.Tensor): 沿每条光线的权重总和。形状:[num_rays]
      - 'rgb0' (torch.Tensor): 粗模型的输出 RGB。见 rgb_map。形状:[num_rays, 3]
      - 'disp0' (torch.Tensor): 粗模型的输出视差图。见 disp_map。形状:[num_rays]
      - 'acc0' (torch.Tensor): 粗模型的输出累积不透明度。见 acc_map。形状:[num_rays]
      - 'z_std' (torch.Tensor): 每个样本沿光线的距离的标准差。形状:[num_rays]

      再对结果进行整理得到一个列表

通过dataloaderbatch传入网络,最终得到整理后的结果:

一个包含了所有渲染结果和相关信息的列表

前面 0、1、2 对应 'rgb_map''disp_map''acc_map'

image-20240331160722758

所有张量类型如下

dtype = torch.float64
device = device(type='cuda', index=0)

2.3 Render

增加 render_utils,并修改 raw2outputs()render_rays() 函数使数据移动至 cuda 进行计算

3 Loss & Evaluate

3.1 Loss

loss 模块主要为 NetworkWrapper 的类。forward 方法即是 loss 计算的过程。

  • 方法中,首先通过 self.net(batch) 调用传入的神经网络模型进行预测
  • 然后计算预测的RGB值与目标RGB值之间的MSE,并将其转换为PSNR。如果输出中包含rgb0(即包含粗网络预测的值),则还会计算粗网络的预测值RGB与目标RGB值之间的MSE和PSNR。
  • 最后,将所有的统计量(包括PSNR和损失)添加到scalar_stats字典中,并返回。

3.2 Evaluate

evaluate 模块在这个项目中主要用于评估模型的性能。并将评估结果保存为图像文件和JSON文件。

  1. 首先从模型的输出和批次数据中获取预测的RGB值和真实的RGB值
  2. 计算均方误差(MSE)和峰值信噪比(PSNR)
  3. 将预测的RGB值和真实的RGB值重塑为图像的形状。再将真实的RGB值和预测的RGB值水平拼接在一起,并保存为图像,形成对比

2 Issues

2.1 Windows bug

2.1.1 DataLoader 的多进程 pickle

image-20240317200544474

在刚开始执行时就遇到了这个错误,开始以为是环境问题,结果鼓捣半天没用,仔细分析了一下错误原因,发现主要是这句话的原因

for iteration, batch in enumerate(data_loader):

也就是序列化的问题,后面还给出了多进程的报错,就应该是多进程的pickle问题,网上一搜,还真是,可能是Windows系统的原因导致的,在 configs 中的配置文件中修改 num_workers = 0,不使用多进程,就解决了报错(速度应该没啥影响)

2.1.2 imageio 输出图片

在执行到 evaluate 时,输出图片出了下面的错误:

envs\NeRFlearning\Lib\site-packages\PIL\Image.py", line 3102, in fromarray raise TypeError(msg) from e TypeError: Cannot handle this data type: (1, 1, 3), <f4

这是因为 imageio.imwrite 函数需要接收的图像数据类型为 uint8 ,而原始的 pred_rgbgt_rgb 可能是浮点数类型的数据。因此,我们需要将它们乘以255(将范围从0-1转换为0-255),然后使用 astype 函数将它们转换为 uint8 类型

# 将数据类型转换为 uint8
pred_rgb = (pred_rgb * 255).astype(np.uint8)
gt_rgb = (gt_rgb * 255).astype(np.uint8)
# 需要添加以上两行,否则报错
imageio.imwrite(save_path, img_utils.horizon_concate(gt_rgb, pred_rgb))

这里还有一个问题,我发现在已经生成过一张图片后,再次执行到这里,输出的新图片无法覆盖之前的旧图片,所以我加上了时间和当前周期作为扩展名

now = datetime.datetime.now()
now_str = now.strftime('%Y-%m-%d_%Hh%Mm%Ss')
base_name, ext = os.path.splitext(save_path)
save_path = f"{base_name}_{now_str}_epoch_{cfg.train.epoch}{ext}"

这样就会以这样的 res_2024-03-18_09h09m40s_epoch_10.jpg 格式正确输出每次的图像了

2.1.3 I/O

此项目是由 Linux 开发的,在Windows系统上,免不了出现各种麻烦。特别是,该项目的所有 I/O 都是 Linux 的 I/O 格式,所以要进行全面修改:

  • 使用 Python 的 os 模块中的 makedirs 函数来替换

    os.system('mkdir -p ' + model_dir)
    # |
    # V
    os.makedirs(model_dir, exist_ok=True)
  • shutil 模块中的 rmtree 函数来替换

    os.system('rm -rf {}'.format(model_dir))
    # |
    # V
    import shutil
    if os.path.exists(model_dir):
    	shutil.rmtree(model_dir)
  • Python 的 os 模块中的remove函数来替换

    os.system('rm {}'.format(os.path.join(model_dir, '{}.pth'.format(min(pths)))))
    # |
    # V
    os.remove(os.path.join(model_dir, '{}.pth'.format(min(pths))))
  • 使用 exit(0) 来结束当前进程

    os.system('kill -9 {}'.format(os.getpid()))
    # |
    # V
    exit(0)