NICE-SLAM是浙大和ETH提出的一种神经隐式SLAM方法( github,论文 链接),主要创新点在于解决了基于传统nerf的方案如imap无法应用在较大场景的问题。
方法简介
NICE-SLAM主要维护一个Feature grids全局地图,该地图包含四个层级(代码中为coarse, middle, fine, color四层,论文图中为三层),每个层级是一个三维栅格地图,每个栅格存储32维特征。
渲染rgb或depth时,首先在ray上按照一定策略采样一些点,然后根据每个点的坐标找到对应Feature grids的栅格并线性插值得到每个点feature,再经过一个decoder得到该点的颜色或占据概率,最后将ray上所有采样点根据渲染公式计算出对应像素的颜色或深度。
建图的过程是tracking和mapping交替进行。Mapping的过程是以当前相机的RGB和depth作为真值,根据Feature grids渲染得到的RGB和depth作为预测值,构建几何或光度误差损失函数,同时对相机位姿和Feature grids中的feature进行优化。Tracking的过程就是基于Feature grids和估计的位姿渲染depth和rgb,通过和相机实际采集的depth和rgb对比,从而优化相机位姿。
NICE-SLAM与imap相比主要优势是可以应用在较大场景中。由于imap基于经典nerf,因此在slam的过程中需要不断用历史帧训练网络,否则网络会逐渐遗忘掉此前构建的地图;并且受限于网络大小,imap只能在一个小房间范围内建图,一旦场景变大则网络难以记忆全部场景。而NICE-SLAM相当于采用了三维栅格地图,每个栅格保存局部的特征,用decoder将特征解码即可恢复出场景,因此即使场景面积很大也不存在网络遗忘的问题。但是NICE-SLAM也有一些可以改进的点,比如:
- 场景面积较大时栅格地图尺寸也会变大,占用很多显存;
- 训练时间较长,无法做到实时;
- 很难通过回环消除累积误差。因为slam过程中对全局栅格地图的特征不断优化,即使能够回环检测,也需要将回环中的所有关联帧拿出来重新对栅格地图做优化。
论文中的框图
各模块大致流程
NICE-SLAM主要有2个核心模块,分别是Tracker和Mapper。默认情况下按照Tracker、Mapper的顺序交替处理。另外其中Mapper包括了coarse_mapping和mapping两个并行的模块,实际上是三个进程在跑,但coarse貌似没有起到什么作用,仅在最终可视化时预测了视野范围外场景,tracker中没有找到在哪里用到。
在默认配置下Mapper会每5帧执行一次建图,即Tracker会跟踪连续5帧相机的位姿,然后Mapper根据Tracker的最新位姿进一步优化地图和位姿,之后Tracker再根据Mapper的最新位姿继续跟踪5帧,不断循环。
Tracker
- 读取color和depth,等待Mapper完成上一帧
- 将最新的decoder和Feature grids拷贝
- 估计当前帧位姿的先验(直接取上一帧的位姿,或者根据匀速运动模型估计)
- 初始化optimizer_camera
- 迭代多次优化相机位姿:在optimize_cam_in_batch()中根据当前估计的位姿采样ray,计算rgb和depth的误差作为loss,从而优化位姿估计
- 将最优的位姿放入estimate_c2w_list给Mapper使用
Mapper
- 等待Tracker
- 读取color和depth,获取tracker估计的相机位姿
- 调用optimize_map()函数优化Feature grids和位姿(位姿在第5帧后才开始优化)。根据mapper的不同,流程也不同,对于coarse mapper,对coarse进行优化; 对于fine mapper,先是优化middle,再优化fine+middle,最后fine+middle+color。
- 调用mesher.get_mesh()得到mesh
一些参数
大部分默认参数在configs/nice_slam.yaml中,程序启动时通过传入自定义配置文件可以override默认参数。
grid_len: #各层级地图尺寸
- coarse: 2
- middle: 0.32
- fine: 0.16
- color: 0.16
- bound_divisible: 0.32 # 地图范围是bound_divisible的整数倍
tracking:
- seperate_LR: False # 在优化位姿时是否将平移和旋转的学习率单独设置(设置为True时效果类似于在高斯牛顿法优化位姿时,将旋转部分的信息矩阵(协方差矩阵的逆)乘了0.2的系数)
- use_color_in_tracking: True #跟踪时是否增加光度误差作为loss
- handle_dynamic: True # 是否考虑动态物体,考虑动态物体的时候会对每个像素的估计误差做筛选,估计偏差太大的不要,不确定度太小的不要
mapping:
- mapping_window_size: 5 # 采样多少关键帧用于优化,此数量包括了当前帧及上一帧。即mapping_window_size - 2为采样的历史关键帧数量
- keyframe_selection_method: 'overlap' # 历史关键帧采样方式,如果'overlap'则根据当前帧采样点落入历史帧视野范围内的点的比例来选择;如果为‘global'则从所有历史关键帧中随机选
- BA: False # 是否用历史关键帧一起优化位姿(即使为false,还是会用历史帧优化feature grids),设置为True时,当关键帧数量大于等于5时开始进行BA
- frustum_feature_selection: True # 是否根据当前帧frustum(视锥)将相关的feature grid提取出来放入优化器,如果为false则将全部范围的feature grid加到优化器中。个人理解是设为True的话只优化和当前帧有关的feature grids,如果设为False则当前帧视野外的grids也会受到历史帧梯度的影响。
- every_frame: 5 # 每隔多少帧优化一次Feature grids
rendering:
- N_samples: 32 # 每条ray上采样多少个点来render一个像素
- N_surface: 16 # 每条ray上在depth真值附近(即物体表面附近)采样多少个点
- lindisp: False # 是否线性间隔采样
model:
- c_dim: 32 # Feature grids的特征长度
比较重要的变量
这些变量主要在NICE-SLAM.py中,由Tracker和Mapper共享。
- self.shared_c 构建的地图,即Feature grids。包含corase、middle、fine、color四个层级。在grid_init()函数初始化各个level的grid,尺寸是[B, c_dim, z, y, x],用dict组织起来。
- self.estimate_c2w_list 存储了相机位姿的列表,位姿为4x4维变换矩阵,列表长度为数据集中的图像总数
- self.idx 记录当前处理到数据集第几帧了,由Tracker进行更新
- self.mapping_idx 记录mapping处理完数据集第几帧了,由fine mapper在地图更新结束后更新
比较重要的函数
optimize_map (Mapper.py)
输入:当前帧相机rgb、深度图;历史关键帧list;当前帧的位姿估计值等
输出:相机位姿
大致流程是选择关键帧->选择待优化的Feature grid->构建optimizer->循环对位姿和地图优化若干次。
关键帧选择
首先选择关键帧,对应论文3.4(Keyframe Selection)。可选两种策略,分别是随机选择和根据overlap,最后再加上最近的关键帧和当前帧。
if len(keyframe_dict) == 0:
optimize_frame = []
else:
if self.keyframe_selection_method == 'global':
num = self.mapping_window_size-2
optimize_frame = random_select(len(self.keyframe_dict)-1, num)
elif self.keyframe_selection_method == 'overlap':
num = self.mapping_window_size-2
optimize_frame = self.keyframe_selection_overlap(
cur_gt_color, cur_gt_depth, cur_c2w, keyframe_dict[:-1], num)
# add the last keyframe and the current frame(use -1 to denote)
oldest_frame = None
if len(keyframe_list) > 0:
optimize_frame = optimize_frame + [len(keyframe_list)-1]
oldest_frame = min(optimize_frame)
optimize_frame += [-1]
其中self.keyframe_selection_overlap(...)函数的主要流程是:
- 在当前帧上采样一些ray(默认100条),每个ray采样16个点,然后投影到每个历史关键帧中。
- 计算落入历史帧视野范围内的点的比例,作为overlap的依据
- 将完全没有overlap的关键帧去掉后,剩下的关键帧随机选k=K-2个
选择待优化Feature grid
选择需要优化的Feature grids并将requires_grad设为True,然后将各个level待优化的grid放入各自的list中。
decoders_para_list = []
coarse_grid_para = []
middle_grid_para = []
fine_grid_para = []
color_grid_para = []
gt_depth_np = cur_gt_depth.cpu().numpy()
if self.nice:
if self.frustum_feature_selection:
masked_c_grad = {}
mask_c2w = cur_c2w
for key, val in c.items():
if not self.frustum_feature_selection:
val = Variable(val.to(device), requires_grad=True)
c[key] = val
if key == 'grid_coarse':
coarse_grid_para.append(val)
elif key == 'grid_middle':
middle_grid_para.append(val)
elif key == 'grid_fine':
fine_grid_para.append(val)
elif key == 'grid_color':
color_grid_para.append(val)
else:
mask = self.get_mask_from_c2w(
mask_c2w, key, val.shape[2:], gt_depth_np)
mask = torch.from_numpy(mask).permute(2, 1, 0).unsqueeze(
0).unsqueeze(0).repeat(1, val.shape[1], 1, 1, 1)
val = val.to(device)
# val_grad is the optimizable part, other parameters will be fixed
val_grad = val[mask].clone()
val_grad = Variable(val_grad.to(
device), requires_grad=True)
masked_c_grad[key] = val_grad
masked_c_grad[key+'mask'] = mask
if key == 'grid_coarse':
coarse_grid_para.append(val_grad)
elif key == 'grid_middle':
middle_grid_para.append(val_grad)
elif key == 'grid_fine':
fine_grid_para.append(val_grad)
elif key == 'grid_color':
color_grid_para.append(val_grad)
if self.nice:
if not self.fix_fine:
decoders_para_list += list(
self.decoders.fine_decoder.parameters())
if not self.fix_color:
decoders_para_list += list(
self.decoders.color_decoder.parameters())
else:
# imap*, single MLP
decoders_para_list += list(self.decoders.parameters())
其中get_mask_from_c2w(...)函数的是根据当前相机的位姿和深度图选择需要被优化的grid,输出是和Feature grids空间尺寸相同的0-1mask。具体流程是:
- 首先根据grid的尺寸创建一个对应的尺寸(长*宽*高,3)的三维点云,假如是coarse层的直接返回全1的mask即可
- 根据相机内参和相机位姿,将点云投影到图像平面,得到尺寸(长*宽*高, 2, 1)的矩阵,保存了点云在图像平面的坐标
- 利用cv2.remap()函数从深度图中取对应的深度
- 根据图像的长宽范围、深度范围得到mask,最后reshape到与grid的尺寸相同
构建optimizer
构建optimizer,配置待优化变量,包括上一步获取的待优化的grid,以及如果BA的话就将此前选出的optimize_frame的相机位姿也一起优化。注意优化位姿时需要将变换矩阵转换成四元数,优化完成后再转回变换矩阵,这在Tracker中也是一样的。
if self.BA:
camera_tensor_list = []
gt_camera_tensor_list = []
for frame in optimize_frame:
# the oldest frame should be fixed to avoid drifting
if frame != oldest_frame:
if frame != -1:
c2w = keyframe_dict[frame]['est_c2w']
gt_c2w = keyframe_dict[frame]['gt_c2w']
else:
c2w = cur_c2w
gt_c2w = gt_cur_c2w
camera_tensor = get_tensor_from_camera(c2w)
camera_tensor = Variable(
camera_tensor.to(device), requires_grad=True)
camera_tensor_list.append(camera_tensor)
gt_camera_tensor = get_tensor_from_camera(gt_c2w)
gt_camera_tensor_list.append(gt_camera_tensor)
if self.nice:
if self.BA:
# The corresponding lr will be set according to which stage the optimization is in
optimizer = torch.optim.Adam([{'params': decoders_para_list, 'lr': 0},
{'params': coarse_grid_para, 'lr': 0},
{'params': middle_grid_para, 'lr': 0},
{'params': fine_grid_para, 'lr': 0},
{'params': color_grid_para, 'lr': 0},
{'params': camera_tensor_list, 'lr': 0}])
else:
optimizer = torch.optim.Adam([{'params': decoders_para_list, 'lr': 0},
{'params': coarse_grid_para, 'lr': 0},
{'params': middle_grid_para, 'lr': 0},
{'params': fine_grid_para, 'lr': 0},
{'params': color_grid_para, 'lr': 0}])
else:
# imap*, single MLP
......
优化
循环中的主要流程是:
- 根据mapper的不同设置stage,对于coarse mapper,stage为 'coarse'; 对于fine mapper,先是middle stage,再fine stage,最后color stage,每个stage的iter数量根据middle_iter_ratio和fine_iter_ratio来设置。各种stage中各种待优化变量的学习率在configs/nice_slam.yaml中配置。对于相机位姿的优化仅在color stage进行。
- 遍历optimize_frame中的每个关键帧,调用get_samples(...)函数根据相机位姿采样一定数量的ray,加入到list中,总计采样5000条。
- 调用renderer.render_batch_ray(...)函数,得到render后的rgb和depth
- render后的rgb和depth和真值计算loss,并backward
- 将优化后的grid更新回全局Feature grids中
- 将相机位姿更新回keyframe_dict中
for joint_iter in range(num_joint_iters):
#
if self.nice:
if self.frustum_feature_selection:
for key, val in c.items():
if (self.coarse_mapper and 'coarse' in key) or \
((not self.coarse_mapper) and ('coarse' not in key)):
val_grad = masked_c_grad[key]
mask = masked_c_grad[key+'mask']
val = val.to(device)
val[mask] = val_grad
c[key] = val
if self.coarse_mapper:
self.stage = 'coarse'
elif joint_iter <= int(num_joint_iters*self.middle_iter_ratio):
self.stage = 'middle'
elif joint_iter <= int(num_joint_iters*self.fine_iter_ratio):
self.stage = 'fine'
else:
self.stage = 'color'
optimizer.param_groups[0]['lr'] = cfg['mapping']['stage'][self.stage]['decoders_lr']*lr_factor
optimizer.param_groups[1]['lr'] = cfg['mapping']['stage'][self.stage]['coarse_lr']*lr_factor
optimizer.param_groups[2]['lr'] = cfg['mapping']['stage'][self.stage]['middle_lr']*lr_factor
optimizer.param_groups[3]['lr'] = cfg['mapping']['stage'][self.stage]['fine_lr']*lr_factor
optimizer.param_groups[4]['lr'] = cfg['mapping']['stage'][self.stage]['color_lr']*lr_factor
if self.BA:
if self.stage == 'color':
optimizer.param_groups[5]['lr'] = self.BA_cam_lr
else:
self.stage = 'color'
optimizer.param_groups[0]['lr'] = cfg['mapping']['imap_decoders_lr']
if self.BA:
optimizer.param_groups[1]['lr'] = self.BA_cam_lr
if (not (idx == 0 and self.no_vis_on_first_frame)) and ('Demo' not in self.output):
self.visualizer.vis(
idx, joint_iter, cur_gt_depth, cur_gt_color, cur_c2w, self.c, self.decoders)
optimizer.zero_grad()
batch_rays_d_list = []
batch_rays_o_list = []
batch_gt_depth_list = []
batch_gt_color_list = []
camera_tensor_id = 0
for frame in optimize_frame:
if frame != -1:
gt_depth = keyframe_dict[frame]['depth'].to(device)
gt_color = keyframe_dict[frame]['color'].to(device)
if self.BA and frame != oldest_frame:
camera_tensor = camera_tensor_list[camera_tensor_id]
camera_tensor_id += 1
c2w = get_camera_from_tensor(camera_tensor)
else:
c2w = keyframe_dict[frame]['est_c2w']
else:
gt_depth = cur_gt_depth.to(device)
gt_color = cur_gt_color.to(device)
if self.BA:
camera_tensor = camera_tensor_list[camera_tensor_id]
c2w = get_camera_from_tensor(camera_tensor)
else:
c2w = cur_c2w
batch_rays_o, batch_rays_d, batch_gt_depth, batch_gt_color = get_samples(
0, H, 0, W, pixs_per_image, H, W, fx, fy, cx, cy, c2w, gt_depth, gt_color, self.device)
batch_rays_o_list.append(batch_rays_o.float())
batch_rays_d_list.append(batch_rays_d.float())
batch_gt_depth_list.append(batch_gt_depth.float())
batch_gt_color_list.append(batch_gt_color.float())
batch_rays_d = torch.cat(batch_rays_d_list)
batch_rays_o = torch.cat(batch_rays_o_list)
batch_gt_depth = torch.cat(batch_gt_depth_list)
batch_gt_color = torch.cat(batch_gt_color_list)
if self.nice:
# should pre-filter those out of bounding box depth value
with torch.no_grad():
det_rays_o = batch_rays_o.clone().detach().unsqueeze(-1) # (N, 3, 1)
det_rays_d = batch_rays_d.clone().detach().unsqueeze(-1) # (N, 3, 1)
t = (self.bound.unsqueeze(0).to(
device)-det_rays_o)/det_rays_d
t, _ = torch.min(torch.max(t, dim=2)[0], dim=1)
inside_mask = t >= batch_gt_depth
batch_rays_d = batch_rays_d[inside_mask]
batch_rays_o = batch_rays_o[inside_mask]
batch_gt_depth = batch_gt_depth[inside_mask]
batch_gt_color = batch_gt_color[inside_mask]
ret = self.renderer.render_batch_ray(c, self.decoders, batch_rays_d,
batch_rays_o, device, self.stage,
gt_depth=None if self.coarse_mapper else batch_gt_depth)
depth, uncertainty, color = ret
depth_mask = (batch_gt_depth > 0)
loss = torch.abs(
batch_gt_depth[depth_mask]-depth[depth_mask]).sum()
if ((not self.nice) or (self.stage == 'color')):
color_loss = torch.abs(batch_gt_color - color).sum()
weighted_color_loss = self.w_color_loss*color_loss
loss += weighted_color_loss
# for imap*, it uses volume density
regulation = (not self.occupancy)
if regulation:
point_sigma = self.renderer.regulation(
c, self.decoders, batch_rays_d, batch_rays_o, batch_gt_depth, device, self.stage)
regulation_loss = torch.abs(point_sigma).sum()
loss += 0.0005*regulation_loss
loss.backward(retain_graph=False)
optimizer.step()
if not self.nice:
# for imap*
scheduler.step()
optimizer.zero_grad()
# put selected and updated features back to the grid
if self.nice and self.frustum_feature_selection:
for key, val in c.items():
if (self.coarse_mapper and 'coarse' in key) or \
((not self.coarse_mapper) and ('coarse' not in key)):
val_grad = masked_c_grad[key]
mask = masked_c_grad[key+'mask']
val = val.detach()
val[mask] = val_grad.clone().detach()
c[key] = val
get_samples (common.py)
函数输入:待采样的ROI范围,相机内参,depth,rgb,以及待采样ray的数量 n,
输出:ray的原点o、方向dir以及rgb(尺寸均为[ n,3])和depth(尺寸[ n])
- 首先get_sample_uv()在图像ROI范围内采样n个点的图像坐标
- 然后get_rays_from_uv()根据采样的坐标和相机内参,计算ray的在world坐标系的方向向量即rays_d,而rays_o则是相机在world坐标系的坐标。
render_batch_ray (Render.py)
函数输入: Feature grids,decoder,ray的dir和原点,stage,以及深度图(可选)
函数输出: 渲染后的深度、颜色、uncertainty图像
- 计算每个ray的范围[near,far]。其中far的计算比较难懂,我的理解是通过用地图范围(bound)减去ray的原点,再除以dir向量后,可以得到bound在dir向量上的投影,内层的torch.max(t, dim=2)[0]可以拿到dir和bound同向的部分,外层的torch.min拿到三个坐标轴上最小的。这样可以保证ray上距离小于far_bb的一定在bound范围内。
det_rays_o = rays_o.clone().detach().unsqueeze(-1) # (N, 3, 1)
det_rays_d = rays_d.clone().detach().unsqueeze(-1) # (N, 3, 1)
t = (self.bound.unsqueeze(0).to(device) -
det_rays_o)/det_rays_d # (N, 3, 2)
far_bb, _ = torch.min(torch.max(t, dim=2)[0], dim=1)
2. 对于每个ray,在[near,far]范围中采样N_samples个点,在深度真值附近采样N_surface个点。得到尺寸[N_rays, N_samples+N_surface, 3]的tensor,再reshape成[N_rays*(N_samples+N_surface), 3]的点云。
3. raw = self.eval_points()渲染,得到[N_rays*(N_samples+N_surface), 4]的tensor,再将尺寸reshape成[N_rays, N_samples+N_surface, 4]
4. 调用raw2outputs_nerf_color()函数,将每条ray按照渲染公式计算对应像素最终的color、depth以及depth的variance
eval_points(Render.py)
输入:若干ray上采样后的的所有点(尺寸[N,3])、decoder、Feature grids、stage等
输出: 渲染结果(尺寸[N,4])
首先用torch.split()按照batch_size拆分points,然后对每个batch中每个点先mask掉范围外的,然后送入decoder,得到输出,最后再拼接起来返回。
其中decoder中的大致流程是:
- 首先调用sample_grid_feature()从Feature grids线性插值获取每个点的特征 (尺寸[Batch=1,32,N],如果是fine层则将fine和middle两个层的特征 concatenate在一起。
- positional_embedding对每个点的坐标进行编码,如果是coarse层则没有这一步。
- 过一个5block的mlp,得到尺寸[N]的tensor,即占据概率occ。
- 如果是color层,还会过一个color_decoder,得到[N, 4]的tensor。
- 最后的输出根据stage的不同而不同:coarse和middle最终输出尺寸是[N,4],只不过仅最后一个通道为occ,前三个通道rgb均为0。fine最终输出尺寸是[N,4],前三个通道rgb均为0,最后一个通道为fine_occ+middle_occ,fine起到残差修正的作用;color的最终输出尺寸是[N,4],最后一个通道用fine_occ+middle_occ替代。