NeRF Pytorch 代码笔记

威尔士矮脚狗

已于 2023-05-28 09:40:32 修改

阅读量4.1k

点赞数 12

分类专栏： NeRF 文章标签： pytorch 深度学习

于 2022-10-27 13:12:39 首次发布

本文链接：https://blog.csdn.net/weixin_46117498/article/details/127419521

版权

NeRF 专栏收录该内容

4 篇文章 1 订阅

订阅专栏

NeRF代码解析

Youtube 代码讲解链接
NeRF的代码分为两部分，训练和推论。
训练阶段分为三部分，准备数据，构建全连接神经网络，计算loss。

准备数据

有三个资料集：

第一个是 nerf_synthetic 其中包含
lego小车不同角度的图片和Camera 的位姿。
在这里插入图片描述
第二个是 nerf_llff_data
这个数据集是由真实照片制作的，Camera的位姿是由colmap生成的，位姿保存在poses_bound.npy的文件中。

代码的解读

main()

if __name__=='__main__':
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
    train()

首先是设置浮点数类型，设置为cuda类型可以直接使用GPU进行加速训练。
在这里插入图片描述
在设置好浮点数类型后调用train函数。

train ()

config_parser 指定训练或渲染使用的参数

其中有如神经网络的层数和data存放的位置等信息。

	args = parser.parse_args()

将参数传到 args 中，args 在使用时通过 args.name 调用相应属性的值。

if args.dataset_type == 'llff':...
elif args.dataset_type == 'blender':...
elif args.dataset_type == 'LINEMOD':...
elif args.dataset_type == 'deepvoxels':...
    else:
        print('Unknown dataset type', args.dataset_type, 'exiting')
        return

load_blender 获取图片、位姿、图片的长宽与焦距、训练集与测试集的划分

images：物体所有图片
poses：所有图片的位姿
render_poses:
hwf:长宽焦距
i_split:将下标划分为train val test
hwf 为list 里面有图片的长宽和焦距

images, poses, render_poses, hwf, i_split = load_blender_data(args.datadir, args.half_res, args.testskip)

将所有图片和相机的位姿读入imgs和pos数组中

json文件中记录位姿部分的内容

    "camera_angle_x": 0.6911112070083618,
    "frames": [
        {
            "file_path": "./test/r_0",
            "rotation": 0.031415926535897934,
            "transform_matrix": [
                [
                    -0.9999999403953552,
                    0.0,
                    0.0,
                    0.0
                ],
                [
                    0.0,
                    -0.7341099977493286,
                    0.6790305972099304,
                    2.737260103225708
                ],
                [
                    0.0,
                    0.6790306568145752,
                    0.7341098785400391,
                    2.959291696548462
                ],
                [
                    0.0,
                    0.0,
                    0.0,
                    1.0
                ]
            ]
        },

for frame in meta['frames'][::skip]:
    fname = os.path.join(basedir, frame['file_path'] + '.png')
    imgs.append(imageio.imread(fname))
    poses.append(np.array(frame['transform_matrix']))

将每张图正规化到0~1之间
一张图片有四个通道，RGB和A透明度

imgs = (np.array(imgs) / 255.).astype(np.float32)
poses = np.array(poses).astype(np.float32)
counts.append(counts[-1] + imgs.shape[0])
all_imgs.append(imgs)
all_poses.append(poses)

计算 hwf 中的 f (focal:焦距)，

focal = .5 * W / np.tan(.5 * camera_angle_x)

在这里插入图片描述

**可选部分 half_res

NeRF中训练是每张图片的每个像素点都进行训练，数据量比较大，可以选择设置 half_res 将图片变为之前的一半，设备性能不好时使用。

    if half_res:
        H = H//2
        W = W//2
        focal = focal/2.
        imgs_half_res = np.zeros((imgs.shape[0], H, W, 4))
        for i, img in enumerate(imgs):
            # INTER_AREA：基于局部像素的重采样
            imgs_half_res[i] = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)
        imgs = imgs_half_res

返回所有图片的RGBA，位姿，render_pose未知，每张图的长宽和焦距，划分train val test的下标

    return imgs, poses, render_poses, [H, W, focal], i_split

设定距离远近的平面

由NeRF的论文可知，在渲染过程中需要在一条 Camera Ray 上采集采样点，并且根据采样点获得像素点的RGB值，采样时要确定图片的远近位置，near和far的值为距离相机中心的距离。

由于blender是合成资料集，near和far是在制作资料集时设定好的。至于在真实拍摄的图片中如何确定near和far的值还没有看懂相应的代码。
在这里插入图片描述

**可选 white_bkgd (是否将背景设为白色)

imges 的形状是(n * h * w * 4)
其中 n 为照片的数目
h 为图像的高度
w 为图像的宽度
4 为图像的 RGBA 的值，A只有两种值，0代表没有物体，1代表有物体
下面是参数设置中对white_bkgd的解释和train中对应部分的代码。

white_bkgd 参数的含义

parser.add_argument("--white_bkgd", action='store_true',
                    help='set to render synthetic data on a white bkgd (always use for dvoxels)')

将透明背景换成白色的算法

if args.white_bkgd:
    images = images[...,:3]*images[...,-1:] + (1.-images[...,-1:])

至此，已将数据集处理完成，可以传入DNN进行训练。
下面的工作主要是设置网络与训练的细节

创建保存模型的目录和文件

basedir = args.basedir
expname = args.expname
os.makedirs(os.path.join(basedir, expname), exist_ok=True)
f = os.path.join(basedir, expname, 'args.txt')
with open(f, 'w') as file:
    for arg in sorted(vars(args)):
        attr = getattr(args, arg)
        file.write('{} = {}\n'.format(arg, attr))
if args.config is not None:
    f = os.path.join(basedir, expname, 'config.txt')
    with open(f, 'w') as file:
        file.write(open(args.config, 'r').read())

create_nerf() :初始化训练 nerf 的神经网络

get_embedder:获取位置编码(positional encoding)

在论文中有证明过PE的重要性，如果直接输入5D坐标的话训练的效果会不太好，通过 γ 函数将数据升至高维后再输入网络，可以提升对细节的展示效果。简单的说就是将输入升至高维可以增大两个点的距离，两个函数输入很近输出也会很近，所以增大输入的距离可以使得神经网络更好地理解高频信息。
论文中对于3D位置坐标和2D视角坐标的编码方式是不同的，对于3D位置坐标，论文中使用频率为10的PE，对于2D坐标，论文中使用频率为4的PE。

下面是参数配置的代码：

parser.add_argument("--multires", type=int, default=10,
                    help='log2 of max freq for positional encoding (3D location)')
parser.add_argument("--multires_views", type=int, default=4,
                    help='log2 of max freq for positional encoding (2D direction)')

get_embedder 的返回值是编码后的高维数据以及它的维度。
下面的图介绍了网络的结构和在什么位置使用这些PE。

对于这方面的知识如果想要深入了解可以去看一下这篇论文：
https://arxiv.org/abs/2006.10739
这篇文章介绍了一种改进PE的新方法，同时也介绍了之前的方法，以及为什么采用PE，以及PE的效果。
在这里插入图片描述

use_viewdirs 解决光线反射问题

有无将视角信息加到网络的输入中。
在这里插入图片描述
角度小幅度变动会导致图像的RGB有大幅度变动，这也属于高频信息，所以在这里也考虑使用PE进行改进。这部分论文中也有提到，最好的频率为4。

if args.use_viewdirs:
    embeddirs_fn, input_ch_views = get_embedder(args.multires_views, args.i_embed)

构造 NeRF 神经网络模型

model = NeRF(D=args.netdepth, W=args.netwidth,
             input_ch=input_ch, output_ch=output_ch, skips=skips,
             input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(device)

转自知乎

神经网络的第8层会输出密度值，并且加上direct相关的信息，最后一层输出5D坐标产生的RGB值。
网络比较简单，就是一些Dense的拼接，有几个地方拼接一下新输入的信息。

训练与保存模型参数

NeRF网络的训练分为两个阶段，第一阶段粗略地进行采样，找到物体的大概位置，第二阶段为精细采样，在密度高的地方放更多的采样点。
course 和 fine 两个网络参数都要放在grad_vars中，

grad_vars = list(model.parameters())

network_query_fn = lambda inputs, viewdirs, network_fn : run_network(inputs, viewdirs, network_fn,embed_fn=embed_fn,embeddirs_fn=embeddirs_fn,netchunk=args.netchunk)

随着训练的进行学习率不断减小

optimizer = torch.optim.Adam(params=grad_vars, lr=args.lrate, betas=(0.9, 0.999))

计算loss

对于每个像素使用相机原点和像素点建立一条 Camera Ray 并在 Camera Ray 上采样进行的训练。最终获得RGB值和密度。
将获得的颜色和实际的颜色进行对比，结果作为loss，使差值最小化。
在这里插入图片描述

use_batching

平等对待每个像素点，随机地取batchsize个像素点进行训练。
use_batching = True:每次训练都会从不同照片获取像素的信息。
use_batching = False:每次训练都会从相同的照片上获取不同的Batchsize个像素。

get_rays_np

将 camera 中的坐标转换成真实世界的坐标。
在这里插入图片描述
需要一个 camera to world 的矩阵，将 camera 坐标和 c2w 矩阵相乘就能得到真实世界下的坐标。

rays_d = np.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1)
rays_o = np.broadcast_to(c2w[:3,-1], np.shape(rays_d))

将 RGB 的值加入到数组中

rays_d = np.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1)  
rays_o = np.broadcast_to(c2w[:3,-1], np.shape(rays_d))

如果将 use_batching = True 时，就会按相应的方法按 batch 训练。
从 get_rays_np 至此，代码完成了将 Camera 中各个像素点的视角转换为真实世界中的坐标，并且将其像素值和坐标合并到一起。
后续的一些收尾工作

将RGB加入其中并重新排布各列位置，方便训练网络：

rays_rgb = np.concatenate([rays, images[:,None]], 1)
rays_rgb = np.transpose(rays_rgb, [0,2,3,1,4])
rays_rgb = np.stack([rays_rgb[i] for i in i_train], 0)
rays_rgb = np.reshape(rays_rgb, [-1,3,3])
rays_rgb = rays_rgb.astype(np.float32)
print('shuffle rays')
np.random.shuffle(rays_rgb)

如果之前设置了use_batching的话，每次训练都会按照设定好的batch取数据。
具体方法是设定N_rand 每个batch 取 N_rand 个数据后令i_batch += N_rand。

if use_batching:
    batch = rays_rgb[i_batch:i_batch + N_rand]
    batch = torch.transpose(batch, 0, 1)
    batch_rays, target_s = batch[:2], batch[2]

    i_batch += N_rand
    if i_batch >= rays_rgb.shape[0]:
        print("Shuffle data after an epoch!")
        rand_idx = torch.randperm(rays_rgb.shape[0])
        rays_rgb = rays_rgb[rand_idx]
        i_batch = 0

render

原点和视角通过 MLP 预测出颜色，再将颜色和真实颜色进行对比，计算 loss。

rgb, disp, acc, extras = render(H, W, K, chunk=args.chunk, rays=batch_rays,
                                                verbose=i < 10, retraw=True,
                                                **render_kwargs_train)

viewdirs :设置这个参数主要为了解决光线反射问题，位置变化很小的时候由于存在光线的反射，会导致RGB变化很大，之前也提到过解决方法，就是获取 viewdir 的位置编码并投入网络进行训练。
这里的做法是将其正规化，并改变形状。

if use_viewdirs:
    viewdirs = rays_d
    if c2w_staticcam is not None:
        rays_o, rays_d = get_rays(H, W, K, c2w_staticcam)
    viewdirs = viewdirs / torch.norm(viewdirs, dim=-1, keepdim=True)
    viewdirs = torch.reshape(viewdirs, [-1,3]).float()

创建用于渲染的 batch
形状为：[N_rand,3,3,1,1]

rays_o = torch.reshape(rays_o, [-1,3]).float()
rays_d = torch.reshape(rays_d, [-1,3]).float()

near, far = near * torch.ones_like(rays_d[...,:1]), far * torch.ones_like(rays_d[...,:1])
rays = torch.cat([rays_o, rays_d, near, far], -1)
if use_viewdirs:
    rays = torch.cat([rays, viewdirs], -1)

batchify_rays:
按照Batchsize的大小进行训练需要可能会out of memory。
每次训练1024个像素点，训练时要投入1024*256（192个fine采样点，64个coarse采样点）显存不足以支持这些，所以要切成小块进行。

all_ret = batchify_rays(rays, chunk, **kwargs)

将1024 * 256的数据切成以1024 * 32为单位的小块进行训练，返回值也是按小块返回的。

all_ret = {}
for i in range(0, rays_flat.shape[0], chunk):
    ret = render_rays(rays_flat[i:i+chunk], **kwargs)
    for k in ret:
        if k not in all_ret:
            all_ret[k] = []
        all_ret[k].append(ret[k])

all_ret = {k : torch.cat(all_ret[k], 0) for k in all_ret}

下面简单介绍下采样的过程
采样分为两部分组成，对应两个网络，coarse 网络负责粗采样，fine网络负责细采样，沿着 camera ray 先进行均匀采样找出体素高的地方，在进行细采样。
在这里插入图片描述
下面是采样公式的积分形式，上下限中的 tf 和 tn 为之前求出的 near 和 far，分别为物体到 camera 的最近距离和最远距离。
T(t) 为该点体素值的权重系数，从式子中可以看出，系数的值为体素从tn到该点的积分取负指数，从起始点到该点的体素值的积分越大，表示光线传到该点时损失越大，即该点的RGB值对 Camera ray 上采样出的 RGB 值的贡献度小。σ(t) 为该点的体素值，c(r,d)为该点的 rgb 值，被积函数的意义为一个点对Camera ray上得到的rgb值的贡献，从起点积分到终点的含义就是将起点到终点上的每个点对最终RGB的贡献求和，最终得到像素点的颜色。
在这里插入图片描述
将采样离散化之后得到了如下公式，将RGB的值转换为 Camera ray 上每一点的 RGB 值加权求和得到的结果。其中δi为上一个采样点到下一个采样点的距离。

这个公式将前面一点对颜色的贡献作为两点间一段距离上所有点对颜色的贡献。就像下面这张图一样，t1 与 t2 之间的所有点对颜色的贡献被看作t1对颜色的贡献，所以公式中使用 ti 的体素乘上 ti 到 ti+1 的距离。
在这里插入图片描述
Ti 如下：ti 与积分形式下的 ti 表示含义相同，都体现了光线传播到该点后的损失程度，也就是体现遮挡这一特性，体素高的地方势必会遮挡后面物体，从而减弱后面物体对颜色的贡献。

综上，c = α1*c1 + (1-α1)α2c2 + … + αn(1-α1)(1-α2)…(1-αn-1)

下面为计算c的代码

计算α的公式

raw2alpha = lambda raw, dists, act_fn=F.relu: 1.-torch.exp(-act_fn(raw)*dists)

计算相邻采样点的距离
在这里插入图片描述

dists = z_vals[...,1:] - z_vals[...,:-1]

取出所有射线和采样点的RGB值，[N_rays, N_samples, 3]

rgb = torch.sigmoid(raw[...,:3])

计算 αn(1-α1)(1-α2)…(1-αn-1)

# 计算 c = α1*c1 + (1-α1)*α2*c2 + ... + αn(1-α1)(1-α2)...(1-αn-1)
# α
alpha = raw2alpha(raw[...,3] + noise, dists)  # [N_rays, N_samples]
# αn(1-α1)(1-α2)...(1-αn-1)
weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1)), 1.-alpha + 1e-10], -1), -1)[:, :-1]

根据采样公式将所有的采样点加权求和，得出RGB的值。

rgb_map = torch.sum(weights[...,None] * rgb, -2)

sample_pdf:在weight较高的地方多做采样。

最终获得[n_rays,n_samples+n_importance,3]的张量。

z_samples = sample_pdf(z_vals_mid, weights[...,1:-1], N_importance, det=(perturb==0.), pytest=pytest)
z_samples = z_samples.detach()
z_vals, _ = torch.sort(torch.cat([z_vals, z_samples], -1), -1)
pts = rays_o[...,None,:] + rays_d[...,None,:] * z_vals[...,:,None]

经过fine网络返回rgb值

run_fn = network_fn if network_fine is None else network_fine
raw = network_query_fn(pts, viewdirs, run_fn)
rgb_map, disp_map, acc_map, weights, depth_map = raw2outputs(raw, z_vals, rays_d, raw_noise_std, white_bkgd, pytest=pytest)

将结果以字典的形式返回

ret = {'rgb_map' : rgb_map, 'disp_map' : disp_map, 'acc_map' : acc_map}

小结：
取得采样点的位置，经过第一个 coarse 网络计算，但是不使用第一个神经网络计算的结果，只为了获取各个位置的权重，在权重较高的地方多设置几个采样点，再使用fine网络进行精细采样，最终将粗细网络采样得到的结果与grand truth 进行对比，最小化它们的差值。
下面是计算loss的代码。

计算 loss

梯度置零，计算渲染图片和真实图片的均方误差。以及psnr。

optimizer.zero_grad()
img_loss = img2mse(rgb, target_s)
trans = extras['raw'][...,-1]
loss = img_loss
psnr = mse2psnr(img_loss)

DONE

威尔士矮脚狗

关注

12
点赞
踩
74

收藏

觉得还不错? 一键收藏
6
评论
NeRF Pytorch 代码笔记

https://www.youtube.com/watch?v=kh_hxFnuQNI&t=8s
复制链接

扫一扫

专栏目录