NICE-SLAM 代码简析

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

  1. 读取color和depth,等待Mapper完成上一帧
  2. 将最新的decoder和Feature grids拷贝
  3. 估计当前帧位姿的先验(直接取上一帧的位姿,或者根据匀速运动模型估计)
  4. 初始化optimizer_camera
  5. 迭代多次优化相机位姿:在optimize_cam_in_batch()中根据当前估计的位姿采样ray,计算rgb和depth的误差作为loss,从而优化位姿估计
  6. 将最优的位姿放入estimate_c2w_list给Mapper使用

Mapper

  1. 等待Tracker
  2. 读取color和depth,获取tracker估计的相机位姿
  3. 调用optimize_map()函数优化Feature grids和位姿(位姿在第5帧后才开始优化)。根据mapper的不同,流程也不同,对于coarse mapper,对coarse进行优化; 对于fine mapper,先是优化middle,再优化fine+middle,最后fine+middle+color。
  4. 调用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(...)函数的主要流程是:

  1. 在当前帧上采样一些ray(默认100条),每个ray采样16个点,然后投影到每个历史关键帧中。
  2. 计算落入历史帧视野范围内的点的比例,作为overlap的依据
  3. 将完全没有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。具体流程是:

  1. 首先根据grid的尺寸创建一个对应的尺寸(长*宽*高,3)的三维点云,假如是coarse层的直接返回全1的mask即可
  2. 根据相机内参和相机位姿,将点云投影到图像平面,得到尺寸(长*宽*高, 2, 1)的矩阵,保存了点云在图像平面的坐标
  3. 利用cv2.remap()函数从深度图中取对应的深度
  4. 根据图像的长宽范围、深度范围得到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
            ......

优化

循环中的主要流程是:

  1. 根据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进行。
  2. 遍历optimize_frame中的每个关键帧,调用get_samples(...)函数根据相机位姿采样一定数量的ray,加入到list中,总计采样5000条。
  3. 调用renderer.render_batch_ray(...)函数,得到render后的rgb和depth
  4. render后的rgb和depth和真值计算loss,并backward
  5. 将优化后的grid更新回全局Feature grids中
  6. 将相机位姿更新回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])
  1. 首先get_sample_uv()在图像ROI范围内采样n个点的图像坐标
  2. 然后get_rays_from_uv()根据采样的坐标和相机内参,计算ray的在world坐标系的方向向量即rays_d,而rays_o则是相机在world坐标系的坐标。

render_batch_ray (Render.py)

函数输入: Feature grids,decoder,ray的dir和原点,stage,以及深度图(可选)
函数输出: 渲染后的深度、颜色、uncertainty图像
  1. 计算每个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中的大致流程是:

  1. 首先调用sample_grid_feature()从Feature grids线性插值获取每个点的特征 (尺寸[Batch=1,32,N],如果是fine层则将fine和middle两个层的特征 concatenate在一起。
  2. positional_embedding对每个点的坐标进行编码,如果是coarse层则没有这一步。
  3. 过一个5block的mlp,得到尺寸[N]的tensor,即占据概率occ。
  4. 如果是color层,还会过一个color_decoder,得到[N, 4]的tensor。
  5. 最后的输出根据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替代。
  • 1
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值