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)