最近在找轻量级的语义分割模型，SalsaNext作为一个很经典的语义分割网络，在服务器的2080上面能够达到30毫秒一帧左右的推理速度，但是其网络本身提出的时间比较久远，后处理的部分使用的依然是最经典的knn，fidnet的后处理根据论文作者自己说的是更快效果更好吗，但是在服务器上的后处理速度要慢几十倍，这主要是因为其开源代码中后处理的部分写的太烂。这里记录一下看SalsaNext推理部分的代码，方便后续融合两个模型。

源代码和权重文件来自于官方开源代码：https://github.com/TiagoCortinhal/SalsaNext/

一、eval.sh

这个脚本文件用于启动几个推理相关的Python文件，这里需要修改一个地方，如果用默认的启动写法，推理的过程是在CPU上面进行的，即使在启动时设置了-g参数也没有效果，这种情况是因为启动的脚本文件中有点小问题。源代码中从命令行中加载参数的写法为：

while getopts "d:p:m:s:n:c:u:g" opt
do
   case "$opt" in
      d ) d="$OPTARG" ;;
      p ) p="$OPTARG" ;;
      m ) m="$OPTARG" ;;
      s ) s="$OPTARG" ;;
      n ) n="$OPTARG"  ;;
      g ) g="$OPTARG" ;;
      u ) u="$OPTARG" ;;
      c ) c="$OPTARG" ;;
      ? ) helpFunction ;;
   esac
done

根据shell文件的写法，getopts中第一个冒号表示忽略错误，每个字母后面的冒号表示这个选项有自己的参数，这个参数保存在内置变量OPTARG中，这样在后面的do-case里面就会赋值给具体的shell变量，只有这样g这个变量才会被赋值，这样在指定显卡的时候才能够正确读取显卡编号。

export CUDA_VISIBLE_DEVICES="$g"

getopts相关的用法可以参考：https://blog.csdn.net/a772304419/article/details/126764754

二、infer.py

根据eval.sh里面的内容，推理启动的主要是infer.py和evaluate_iou.py两个文件，这两个文件一个负责推理，另一个负责根据推理结果计算iou。

在infer.py中，主函数首先根据启动时传入的参数进行解析，并根据参数信息设置不同的参数。包括了数据集的位置、预测结果保存的位置、推理模型的位置、是否使用不确定性进行推理、采样数量以及用模型需要做的操作（训练、验证还是检验）。

读取启动参数后会根据参数中模型位置的路径打开两个配置用的yaml文件，这两个文件在官方提供的开源模型中是打包在一起的，启动指令时只要将模型位置写到所在路径的上一层就能够正确找到模型和配置文件中。

之后代码会在启动指令中给出的希望保存的路径下创建用于存储label文件的文件夹。

在检查模型文件是否存在之后，程序会初始化一个user对象，用于处理推理的过程，这个对象对应的是user,py，也是最关键的一个部分。

三、user.py

user.py文件中本质上只有一个类，这个类中包括了构造函数和推理用的两个函数

构造函数中会根据传入的参数，初始化一个parser对象，不同于fidnet的代码，在这个parser对象中十分简洁地写明了从yaml文件中读取到的数据集信息。

parser.py

这里展开说一下这个对象，parser对象对应的是parser.py这个文件，对这个类的定义个人感觉是继承了Python库文件中的一些类，除了加载参数，这个类还负责初始化pytorch里面用于加载数据的对象DataLoader。

可以看到，在初始化DataLoader对象时，第一个参数是一个数据集对象，根据参考资料的内容，如果是一些常用的数据集，可以直接用pytorch进行调用，而对于自定义的数据集，则需要实现一个数据集类，这个类中需要包含构造函数、getitem以及len三个函数。而对于语义分割这个任务而言，SalsaNext使用的是球面投影的方法，所以在这个地方直接将球面投影在getitem函数中实现了。这样子在推理的过程中，就可以直接得到球面投影后的五个通道内容。

由于在parser对象初始化的时候已经将所有的点云文件和标签文件都存储在了scan_files和label_files中，所以在getitem中就只需要根据当前帧索引来找到对应的点云。这里我们考虑的是模型推理的过程，所以自然不存在读取标签文件。在打开点云文件时，代码首先是初始化了一个LaserScan对象，这个对象对应的是laserscan.py这个文件，这个对象主要是用于打开并存储一帧点云文件，使用np.fromfile读取点云文件并将其初始化为float数据类型，点云文件中的前三列被是点的xyz坐标，第三列是回波强度。有关回波强度（remission）和反射强度（intensity）之间的差别可以参考链接。

在LaserScan对象中，open_scan函数执行后保存的是原始的点云数据，而在模型推理过程使用的输入是球面投影后64×2048分辨率的图像，所以在后续的do_range_projection函数中还需要同时完成投影的部分。

点云球面投影的实现过程

函数首先会计算一个可视范围，同时利用np.linalg.norm函数计算所有点的深度，np.linalg.norm本身是用来计算范数的，在这里也可以直接当做计算距离的函数去使用。

之后根据点的坐标计算点的俯仰角pitch和偏航角yaw，这两个角用于描述点与y轴和x轴之间的夹角，在投影过程会用到这两个夹角。根据欧拉角的定义，位于xoy面上方时pitch为正，位于yoz面右侧时yaw为正。在投影过程中，首先要将这两个角调整到0-1的区间范围内。对于偏航角yaw，在一帧点云中其分布范围为[-π, π]，除以π后为[-1.0, 1.0]的范围，最后通过加一除以二调整到 [0.0, 1.0]的范围，在这个新的范围中，坐标轴的原点也发生了变化，调整前yaw的原点为正前方，调整后原点为原始点云角度为-π位置的点，0.0表示点云点在图像的最左侧，1.0表示点云点在图像的最右侧。而对于俯仰角pitch，先加上视场角的下界的绝对值将俯仰角的范围调整到正值的范围，之后除以总的视场角进行归一化，这样调整后0.0表示点云点在图像的最上方，1.0表示点云点在图像的最下方。现在得到的投影坐标是一个比例范围，用这个比例范围乘以投影图像的尺寸，就可以计算出点云所有点在投影图中的距离。

由于坐标是不存在小数部分的，所以现在计算出的投影坐标还需要做修正，分别将proj_x和proj_y限制在[0, W-1]和[0, H-1]的范围内。最后转换为整形数，就可以完成整个的点云投影。这里补充一点，正是因为这一步修正，导致一些点投影在深度图上的时候产生了重叠，这进一步导致了反投影回点云的时候出现的物体边缘语义不准确的问题，这也是很多基于球面投影的语义分割方法需要修正的一个问题。

最后LaserScan对象会保存投影图上每个像素对应的深度、对应点云中的原始点、对应点的回波强度。order数组表示了按照深度从大到小重新排序后的索引顺序。接下来，代码根据order数组对各个数组（depth、indices、points、remission、proj_y 和 proj_x）进行重新排序，确保它们的顺序与深度数据的顺序一致。这样，depth 数组中的深度值将按照从大到小的顺序排列，而其他对应的数组中的元素也将按照相同的顺序进行重新排列。然后，代码将重新排序后的点云数据投影到图像中的对应位置。具体来说，代码使用 proj_y 和 proj_x 数组作为投影位置的索引，将深度值、点坐标、反射强度、索引等信息分别赋值给对应的图像数组。这样，图像中的每个像素位置 (proj_y, proj_x) 将与点云数据中的某个点对应，并记录了该点的深度、坐标、反射强度和索引信息。最后，代码根据投影的索引信息生成了一个投影掩码 proj_mask，其中大于0的元素表示有效的投影点，用整数类型（np.int32）表示。

回到parser.py文件，现在我们已经在scan对象中存储了点云处理后的信息，但是这些信息都是numpy形式的，因此在parser中还会进行一些数据类型上的转换。

除了投影前原始点云的信息，这里同样是投影后的信息进行了格式转换。

这样子整个parser对象也就基本结束了，这个对象的关键在于利用自定义的数据集对象SemanticKitti，重写了getitem函数，并在这个函数中利用LaserScan对象对点云进行球面投影处理，这样在使用DataLoader加载数据时，就可以直接得到处理好的点云。所以严格地说，parser对象不只是一个存储参数用的对象，更是一个加载数据集的对象。

回到user.py中，使用torch.load加载模型权重，再使用load_state_dict将提供的模型文件加载到显卡上。相关操作可参考链接。

模型推理 infer_subset()

infer.py的最后调用了user对象的infer函数，infer函数会根据参数中的split选择用验证集还是测试集。不管是哪种操作，程序都会进入infer_subset函数。

进入函数后首先初始化一些记录用的参数，之后使用torch.no_grad()函数禁用梯度计算，这样子可以显著减少内存使用并加速计算。之后使用前面初始化的loader进行数据的遍历，由于进入infer_subset函数时会根据参数传入不同的loader，所以在推理时这里传入的就只是parser对象里面我们已经读好数据的SemanticKitti对象，这里也看到，使用枚举从loader里面加载数据时，读到的内容和之前重写的getitem是相互对应的。

在得到当前帧点云的数据后，程序调用cuda()函数将训练所需的参数加载到gpu上，这其中proj_in是模型推理的内容，而剩下的都是后处理需要的内容。

之后利用self.model函数将输入送入模型并得到proj_output输出，同时还是用argmax得到了一个最大值矩阵，从代码写法来看，这个矩阵是将proj_output中第0个通道里面最大值所在的维度取了出来，模型的输入尺度是[1, 5, 64, 2048]，模型的输出尺度是[1, 20, 64, 2048]，输入中5对应的是输入的五个通道，64和2048对应的是每个通道投影图的大小；输出尺度的区别在第二个维度，因为salsanext原始模型是20分类的，所以相当于输出是20个大小为64×2048矩阵，这个矩阵表示投影图中每个像素可能是某个语义类别的概率。这样的话，使用argmax就将这个输出的四维矩阵转换为了一个矩阵，这个矩阵代表了每个像素最可能的语义类别。

如果使用后处理，则会跳转到knn.py中进行计算，如果不用则直接根据投影位置将点云内点的语义信息进行赋值。

最后将所有的语义信息进行保存，展开成一列并保存为label文件，文件中每个位置为原点云中每个点的语义类别编号。

后处理 knn.py

后处理的部分则是在knn.py文件中，在user对象初始化的时候会根据yaml文件进行后处理的初始化，对应的yaml文件和提供的模型参数文件放在一起，有关的配置参数为：

后处理的实现则是在forward函数里面，一进函数就可以看到，代码会根据设备选择使用cpu还是gpu进行后处理，这显然就要比fidnet的速度快得多。函数输入一共包括五个参数：proj_range, unproj_range, proj_argmax, px, py；对应的是投影图，所有点云的深度，投影图每个位置的语义类别，点云中每个点对应到深度图的xy坐标，这里需要注意，投影图尺寸是64×2048，相当于一共有131072个点，但是一帧点云中并没有这么多点，一般都是要小于131072的，这稍微有点反常识，个人理解是点云并不是每个点都合法，所以丢失了一部分，所以总数目要小于投影图的像素点，但是这并不妨碍投影过程产生的重叠。

knn后处理，本质上就是让点云里面的所有点，根据在投影图上的位置，选择一定范围内深度距离最近的固定数目的点，根据这些点的语义信息来确定中心点的语义信息。这个过程涉及筛选邻近范围内的点，salsanext使用的是torch.nn.functional.unfold这个函数，函数说明也没很看懂，反正大体意思就是进行一个类似卷积的操作，函数输入是64×2048大小的深度图，输出变成了1×25×131072，其中25表示的是投影图中每个像素邻近的5×5范围共计25个像素的深度，131072则表示深度图中的每个像素。展开后为了找到点云中每个点对应的是哪一个位置，同时还需要用py和px算出来一个一维的对应的位置idx_list。这样根据idx_list就可以筛选出点云中每个点在深度图上一定范围内的临近点。

在进行knn计算的过程中，salsanext采用的是加权的距离，计算时首先用点云中每个点的深度替换中心点的深度，计算中心点到所有临近点的深度距离，这个距离加权后会作为knn中衡量距离的指标。

计算距离后根据这个距离挑选出五个最近的点，之后根据这五个点的语义信息投票决定中心点的语义信息。

总的来说，knn后处理是用语义分割后的投影图来对点云的语义做了一次再处理，点云的每个点根据其投影位置确定在投影图上邻近范围，在这个邻近范围内，根据深度信息加权计算距离，利用加权距离进行knn的计算，得到的k个点的语义信息会被进行统计以此确定点云中点的语义信息。

带有备注的代码如下：

    def forward(self, proj_range, unproj_range, proj_argmax, px, py):
        # 投影图 所有点云的深度 投影图每个位置的语义类别 点云中每个点对应到深度图的xy坐标
        ''' Warning! Only works for un-batched pointclouds.
            If they come batched we need to iterate over the batch dimension or do
            something REALLY smart to handle unaligned number of points in memory
        '''
        # get device
        if proj_range.is_cuda:
            device = torch.device("cuda")
        else:
            device = torch.device("cpu")

        # sizes of projection scan
        # 深度图大小
        H, W = proj_range.shape

        # number of points
        # 点云中所有点的数目
        P = unproj_range.shape
        # check if size of kernel is odd and complain
        if (self.search % 2 == 0):
            raise ValueError("Nearest neighbor kernel must be odd number")

        # calculate padding
        pad = int((self.search - 1) / 2)

        # unfold neighborhood to get nearest neighbors for each pixel (range image)
        # 将深度图展开 展开前为64×2048 展开后为1×25×131072 展开时以每个像素为中心 选择kernel_size范围内的点 边缘位置补充2宽度的0
        # 25表示的是投影图中每个像素邻近的5×5范围共计25个像素的深度
        proj_unfold_k_rang = F.unfold(proj_range[None, None, ...],
                                      kernel_size=(self.search, self.search),
                                      padding=(pad, pad))
        # index with px, py to get ALL the pcld points
        idx_list = py * W + px
        # py和px都是原始点云对应的投影点的坐标 利用这一步操作 将二维坐标转换为一维坐标

        unproj_unfold_k_rang = proj_unfold_k_rang[:, :, idx_list]
        # 筛选出点云对应的投影点 从而得到每个投影点在深度图上对应点通过knn计算出的25个临近点的深度

        # WARNING, THIS IS A HACK
        # Make non valid (<0) range points extremely big so that there is no screwing
        # up the nn self.search
        unproj_unfold_k_rang[unproj_unfold_k_rang < 0] = float("inf")
        # 深度不能小于0 去掉非法值

        # now the matrix is unfolded TOTALLY, replace the middle points with the actual range points
        center = int(((self.search * self.search) - 1) / 2)
        unproj_unfold_k_rang[:, center, :] = unproj_range
        # 将中心位置的深度 替换为投影之前的深度值

        # now compare range
        k2_distances = torch.abs(unproj_unfold_k_rang - unproj_range)
        # 计算中心点到邻近范围内其他点的距离

        # make a kernel to weigh the ranges according to distance in (x,y)
        # I make this 1 - kernel because I want distances that are close in (x,y)
        # to matter more
        inv_gauss_k = (
                1 - get_gaussian_kernel(self.search, self.sigma, 1)).view(1, -1, 1)
        inv_gauss_k = inv_gauss_k.to(device).type(proj_range.type())
        # 计算高斯核权重

        # apply weighing
        k2_distances = k2_distances * inv_gauss_k
        # 加权计算距离值

        # find nearest neighbors
        _, knn_idx = k2_distances.topk(
            self.knn, dim=1, largest=False, sorted=False)
        # 根据加权距离值筛选最近邻的五个点 knn_idx为点云中每个点最近邻的索引

        # do the same unfolding with the argmax
        proj_unfold_1_argmax = F.unfold(proj_argmax[None, None, ...].float(),
                                        kernel_size=(self.search, self.search),
                                        padding=(pad, pad)).long()
        unproj_unfold_1_argmax = proj_unfold_1_argmax[:, :, idx_list]
        # 对语义信息做同样的展开操作 获取每个点邻近范围内的25个位置的语义信息
        print(knn_idx.shape)
        # get the top k predictions from the knn at each pixel
        knn_argmax = torch.gather(
            input=unproj_unfold_1_argmax, dim=1, index=knn_idx)
        # 获得最近邻的语义信息

        # fake an invalid argmax of classes + 1 for all cutoff items
        # 根据knn距离进行再次筛选
        if self.cutoff > 0:
            knn_distances = torch.gather(input=k2_distances, dim=1, index=knn_idx)
            knn_invalid_idx = knn_distances > self.cutoff
            knn_argmax[knn_invalid_idx] = self.nclasses

        # now vote
        # argmax onehot has an extra class for objects after cutoff
        # 根据最近邻的点的语义信息 恢复点的语义信息
        knn_argmax_onehot = torch.zeros(
            (1, self.nclasses + 1, P[0]), device=device).type(proj_range.type())
        ones = torch.ones_like(knn_argmax).type(proj_range.type())
        knn_argmax_onehot = knn_argmax_onehot.scatter_add_(1, knn_argmax, ones)

        # now vote (as a sum over the onehot shit)  (don't let it choose unlabeled OR invalid)
        knn_argmax_out = knn_argmax_onehot[:, 1:-1].argmax(dim=1) + 1

        # reshape again
        knn_argmax_out = knn_argmax_out.view(P)

        return knn_argmax_out