点云深度学习开山之作——PointNet、PointNet++。使用Python（PyTorch）实现物体分类、部件分割、场景分割。
原文：
1 原文解读
以《PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space》为主，适当结合其他论文。
1.1 引言
主要介绍了研究背景、PointNet的局限性、PointNet++的设计思路以及研究的主要贡献，具体内容如下：
研究背景：本文聚焦于分析欧几里得空间中的几何点集，其中3D扫描仪获取的点云数据是一种重要类型。点集数据具有成员排列不变性，并且距离度量定义的局部邻域可能呈现不同特性，例如点的密度和其他属性在不同位置可能不均匀，这在实际3D扫描中普遍存在，如受透视效果、径向密度变化、运动等因素影响。
PointNet的局限性：在点集深度学习领域，PointNet是开创性工作，它直接处理点集，基本思路是学习每个点的空间编码，然后聚合所有点的特征得到全局点云签名。然而，PointNet设计上没有捕捉由度量诱导的局部结构。在卷积架构中，利用局部结构对其成功至关重要，如CNN能够在多分辨率层次结构中逐步捕获越来越大尺度的特征，较低层神经元感受野小，较高层感受野大，这种沿层次结构抽象局部模式的能力使其对未见情况具有更好的泛化性，而PointNet缺乏这种能力。
PointNet++的设计思路：为解决上述问题，作者提出一种名为PointNet++的分层神经网络，以分层方式处理度量空间中采样的点集。其基本思想是通过底层空间的距离度量将点集划分为重叠的局部区域，类似CNN，从小邻域中提取捕获精细几何结构的局部特征，再将这些局部特征组合成更大的单元并处理以生成更高层次的特征，不断重复该过程直至获得整个点集的特征。在设计PointNet++时，需要解决两个关键问题：一是如何生成点集的划分；二是如何通过局部特征学习器抽象点集或局部特征。作者选择PointNet作为局部特征学习器，因为它能有效处理无序点集进行语义特征提取，且对输入数据损坏具有鲁棒性。PointNet++通过在输入集的嵌套划分上递归应用PointNet来实现分层特征学习。对于生成点集的重叠划分，每个划分被定义为底层欧氏空间中的邻域球，其参数包括质心位置和尺度，质心通过最远点采样（FPS）算法从输入点集中选择，与固定步长扫描空间的体积CNN相比，PointNet++的局部感受野依赖于输入数据和度量，更高效有效。
研究的主要贡献：确定局部邻域球的合适尺度是一个具有挑战性的问题，因为特征尺度和输入点集的非均匀性相互纠缠。作者假设输入点集在不同区域可能具有可变密度，这与CNN输入不同（CNN输入可视为定义在具有均匀恒定密度的规则网格上的数据）。在CNN中，较小的内核有助于提高性能，但在点集数据实验中发现，小邻域可能因采样不足而包含太少点，不足以让PointNets稳健地捕获模式。PointNet++的重要贡献在于利用多个尺度的邻域来实现稳健性和细节捕获。在训练过程中通过随机输入丢弃，网络学习根据输入数据自适应地权衡不同尺度检测到的模式并组合多尺度特征。实验表明，PointNet++能够高效且稳健地处理点集，在具有挑战性的3D点云基准测试中取得了显著优于现有方法的结果。
1.2 问题描述
主要对要解决的问题进行了数学描述，明确了研究的输入和输出，为后续提出的方法奠定了基础，具体内容如下（了解即可）：
假设X=(M, d)是一个离散度量空间，其度量继承自欧几里得空间\mathbb{R}^{n}。其中，M \subseteq \mathbb{R}^{n}是点的集合，d是距离度量。并且，M在环境欧几里得空间中的密度并非处处均匀。
研究的目标是学习集函数f，该函数以这样的X（以及每个点的额外特征）作为输入，并产生与X相关的语义信息。在实际应用中，f可以是分类函数，为X分配一个标签；也可以是分割函数，为M中的每个成员分配一个逐点标签 。通过这样的设定，将点集处理任务转化为对特定集函数的学习问题，后续的PointNet++方法就是围绕如何有效学习这个集函数展开的。
1.3 方法
本研究工作可看作是对PointNet的拓展，增添了分层结构。首先回顾PointNet（论文中3.1节），接着介绍带有分层结构的PointNet基本扩展（论文中3.2节）。最后，提出即便在非均匀采样的点集中也能够稳健学习特征的PointNet++（论文中3.3节）。
1.3.1 PointNet回顾：通用连续集函数逼近器
给定一个无序点集\{x_{1}, x_{2}, \ldots, x_{n}\}，其中x_{i} \in \mathbb{R}^{d}，可以定义一个集函数f: X \to \mathbb{R}，它将一组点映射到一个向量：
f\left(x_{1}, x_{2}, \ldots, x_{n}\right)=\gamma\left(\underset{i = 1,\ldots, n}{\text{MAX}}\left\{h\left(x_{i}\right)\right\}\right)其中\gamma和h通常是多层感知器（MLP）网络。公式中的集函数f对输入点的排列具有不变性（置换不变性，Permutation Invariance，例如\max(x_1,x_2,\cdots,x_n)、x_1+x_2+\cdots+x_n等），并且可以任意逼近任何连续集函数。h的响应可解释为一个点的空间编码。
PointNet Classification Network：
推广——PointNet Segmentation Network：
该网络即为PointNet (vanilla)，对于数据缺失的鲁棒性显著优于同期其他网络（如VoxNet），因为点云中存在一些关键点（critical points）对全局特征（global feature）贡献较大。
PointNet的特点为Light-weight & Fast，在一些基准测试中取得了不错的成绩。然而，它缺少在不同尺度下捕捉局部上下文的能力。后续将介绍一种分层特征学习框架来解决这一局限。
1.3.2 分层点集特征学习
虽然PointNet使用单一的最大池化操作来聚合整个点集，但新架构构建了点的分层分组，并在层级结构中逐步抽象出越来越大的局部区域。
基本改进思想：在局部迭代使用PointNet（分层特征学习，Hierarchical Feature Learning）
PointNet++总体思路：首先通过将点集划分为重叠的局部区域。与CNN类似，提取局部特征以捕获来自小邻域的精细几何结构。这些局部特征将进一步分组为较大的单元，并进行处理以生成更高级别的特征。重复此过程，直至获得整个点集的特征。
图2展示了分层特征学习架构及其在2D欧几里得空间中用于点集分割和分类的应用示例。如图所示，该分层结构由多个集合抽象层（set abstraction level）组成。在每个层级，一组点被处理和抽象，生成一个元素更少的新集合。集合抽象层由采样层、分组层和PointNet层这三个关键层构成：
- 采样层（Sampling layer）：给定输入点
\{x_{1}, x_{2}, \ldots, x_{n}\}，采用迭代最远点采样（Farthest Point Sampling，FPS）选择一个点子集\{x_{i_{1}}, x_{i_{2}}, \ldots, x_{i_{m}}\}，使得x_{i_{j}}是相对于集合\{x_{i_{1}}, x_{i_{2}}, \ldots, x_{i_{j - 1}}\}中其余点距离最远的点（按度量距离）。与随机采样相比，在相同数量质心的情况下，FPS能更好地覆盖整个点集。与不考虑数据分布扫描向量空间的CNN不同，该采样策略以依赖数据的方式生成感受野。 - 分组层（Grouping layer）：该层输入为大小
N \times(d + C)的点集和大小N' \times d的一组质心坐标，输出为大小N' \times K \times(d + C)的点集组，其中每个组对应一个局部区域，K是质心点邻域内的点数。需要注意的是，K在不同组之间有所变化，但后续的PointNet层能够将数量可变的点转换为固定长度的局部区域特征向量。在卷积神经网络中，像素的局部区域由与该像素具有特定曼哈顿距离（内核大小）内的像素组成。而在从度量空间采样的点集中，点的邻域由度量距离定义。球查询（Ball query）用于找到所有距离查询点在一定半径内的点（在实现中设置查询点数K的上限），另一种范围查询是K-最近邻（kNN）搜索，其仅查询固定数量的相邻点。与kNN相比，球查询的局部邻域保证了固定的区域尺度，使局部区域特征在空间上更具泛化性，更适合需要局部模式识别的任务（如语义点标记） 。 - PointNet层：在这一层，输入是
N'个局部区域的点，数据大小为N' \times K \times(d + C)。每个局部区域在输出中由其质心和编码质心邻域的局部特征进行抽象，输出数据大小为N' \times(d + C')。首先将局部区域中点的坐标相对于质心点转换到局部坐标系：x_{i}^{(j)} = x_{i}^{(j)} - \hat{x}^{(j)}，其中i = 1, 2, \ldots, K，j = 1, 2, \ldots, d，\hat{x}是质心的坐标。利用3.1节中描述的PointNet作为局部模式学习的基本构建块，通过使用相对坐标和点特征，可以捕获局部区域内点与点之间的关系。
一个集合抽象层以N \times(d + C)矩阵作为输入（来自具有d维坐标和c维点特征的N个点），输出一个N' \times(d + C')矩阵，包含N'个下采样点的d维坐标和新的C'维特征向量，这些特征向量概括了局部上下文。通过这三个关键层的协同工作，实现了点集的分层特征学习，逐步从局部区域提取特征并进行抽象，为后续的分类、分割等任务提供更具代表性的特征表示。
1.3.3 非均匀采样密度下的稳健特征学习
如前文所述，点集在不同区域具有非均匀密度的情况很常见。这种非均匀性给点集特征学习带来了重大挑战。在密集数据中学习到的特征可能无法推广到稀疏采样区域。因此，为稀疏点云训练的模型可能无法识别精细的局部结构。
理想情况下，人们希望尽可能仔细地检查点集，以捕捉密集采样区域中的最精细细节。然而，在低密度区域，由于采样不足可能会破坏局部模式，这种仔细检查是不可行的。在这种情况下，应该在更大的邻域中寻找更大尺度的模式。为了实现这一目标，文章提出了密度自适应PointNet层（图3），当输入采样密度变化时，该层学习组合不同尺度区域的特征。将具有密度自适应PointNet层的分层网络称为PointNet++。
- 多尺度分组（Multi-Scale Grouping，MSG）：如图3(a)所示，一种简单有效的捕捉多尺度模式的方法是应用不同尺度的分组层，然后通过相应的PointNets提取每个尺度的特征。与单尺度分组（Single-Scale Grouping，SSG）相反，不同尺度的特征被连接起来形成多尺度特征。通过随机输入丢弃（random input dropout）训练网络学习优化的多尺度特征组合策略。具体来说，对于每个训练点集，从
[0, p]（通常设p = 0.95以避免生成空点集）中均匀采样一个丢弃率\theta，每个点以概率\theta随机丢弃。这样在训练过程中，网络会接触到不同稀疏度（由\theta引起）和均匀度（由丢弃随机性引起）的训练集，而在测试时保留所有可用点。 - 多分辨率分组（Multi-Resolution Grouping，MRG）：上述MSG方法计算成本较高，因为它需要为每个质心点在大尺度邻域运行局部PointNet。特别是在最低层级，质心点数量通常很大，时间成本显著。这里提出一种替代方法，避免了这种昂贵的计算，同时仍保留根据点的分布属性自适应聚合信息的能力。如图3(b)所示，在某一层
L_{i}，一个区域的特征是两个向量的连接。一个向量（图中左侧）通过使用集合抽象层汇总下一层L_{i - 1}每个子区域的特征得到；另一个向量（图中右侧）是通过单个PointNet直接处理局部区域内所有原始点得到的特征。当局部区域密度较低时，第一个向量可能不如第二个向量可靠，因为计算第一个向量时涉及的子区域包含更稀疏的点，受采样不足影响更大，此时应给第二个向量更高的权重。反之，当局部区域密度较高时，第一个向量能提供更精细的细节信息，因为它能够在较低层级递归地以更高分辨率进行检查。
与MSG相比，MRG计算效率更高，因为避免了在最低层级大尺度邻域的特征提取。
通过这两种密度自适应层，PointNet++能够根据点集的局部密度自适应地组合多尺度特征，从而在非均匀采样的点集中稳健地学习特征，有效提高了模型对不同密度点集的适应性和特征学习能力，增强了模型在复杂场景下的性能。
1.3.4 用于集合分割的点特征传播
在集合抽象层中，原始点集会进行下采样。然而在集合分割任务，如语义点标记中，需要获取所有原始点的点特征。一种解决方案是在所有集合抽象层中始终将所有点作为质心进行采样，但这会导致计算成本过高。另一种方法是将特征从下采样后的点传播到原始点（特征传播，Feature Propagation），即上采样。
本文采用一种基于距离插值（distance based interpolation，属于线性插值）和跨层跳跃连接（accross level skip links）的分层传播策略（如图2所示）。在特征传播层，将点特征从N_{l} \times(d + C)个点传播到N_{l - 1}个点，其中N_{l - 1}和N_{l}（N_{l} \leq N_{l - 1}）分别是集合抽象层l的输入和输出点集大小。通过在N_{l - 1}个点的坐标处插值N_{l}个点的特征值f来实现特征传播。在众多插值方法中，使用基于k最近邻的逆距离加权平均（inverse distance weighted average）（默认p = 2，k = 3），公式为
f^{(j)}(x)=\frac{\sum_{i = 1}^{k} w_{i}(x) f_{i}^{(j)}}{\sum_{i = 1}^{k} w_{i}(x)} \text{ where } w_{i}(x)=\frac{1}{d\left(x, x_{i}\right)^{p}},\ j = 1, \ldots, C在N_{l - 1}个点上插值得到的特征，随后与来自集合抽象层的跳跃连接点特征（skip linked point features）进行连接。然后，连接后的特征会经过一个“单位点网”（unit pointnet），它类似于CNN中的1\times 1卷积。通过几个共享的全连接层和ReLU层来更新每个点的特征向量。这个过程不断重复，直到将特征传播到原始点集。
通过这种点特征传播方式，PointNet++能够在集合分割任务中为原始点集中的每个点都获取到有意义的特征表示，从而为准确的语义点标记等任务提供了有力支持，使模型在复杂场景下的点集分割任务中能够更好地识别和区分不同的物体和区域。
1.4 实验
数据集：文章在四个数据集上进行评估，这些数据集涵盖了从2D对象到3D对象，再到真实3D场景的多种类型，具体如下——
- MNIST：包含手写数字图像，有60,000个训练样本和10,000个测试样本。该数据集用于评估模型在2D点集分类任务上的性能，通过将MNIST图像转换为数字像素位置的2D点云来进行实验。
- ModelNet40：由40个类别的CAD模型（大多为人造模型）组成。使用官方划分，其中9,843个形状用于训练，2,468个用于测试。在实验中，从ModelNet40形状的网格表面采样3D点云，用于3D点云分类任务的评估。
- SHREC15：包含来自50个类别的1200个形状，每个类别有24个形状，大多为有机物体的不同姿态，如马、猫等。采用五折交叉验证来获取该数据集上的分类准确率，用于评估模型在非刚性形状分类任务中的表现。
- ScanNet：包含1513个扫描并重建的室内场景。遵循文献[5]中的实验设置，使用1201个场景进行训练，312个场景进行测试。该数据集用于语义场景标记任务的评估，通过将点云标签预测转换为体素标记，与基于体素化扫描的方法进行公平比较。
选择这些数据集进行实验，能够全面评估PointNet++在不同类型点集数据上的性能，涵盖了从简单的2D数字点集到复杂的3D室内场景点集，以及刚性和非刚性3D物体的分类任务，充分验证模型在多种实际应用场景中的有效性和泛化能力。
1.4.1 欧几里得度量空间中的点集分类
文章在从2D（MNIST）和3D（ModelNet40）欧几里得空间采样的点云分类任务上对网络进行评估。MNIST图像被转换为数字像素位置的2D点云，3D点云则从ModelNet40形状的网格表面采样获得。默认情况下，MNIST使用512个点，ModelNet40使用1024个点。在表格最后一行（“ours normal”），将面法线作为额外的点特征，并且使用更多的点（N = 5000）来进一步提升性能。所有点集都被归一化，使其均值为零且在单位球内。使用一个具有三个全连接层的三级分层网络进行实验。
- 结果对比：在表1和表2中，文章将所提方法与一系列先前的代表性方法进行了比较。从结果中可以明显看出：
- 文章提出的分层学习架构比非分层的PointNet性能有显著提升。在MNIST数据集上，从PointNet (vanilla)和PointNet到文章所提方法，错误率分别相对降低了60.8%和34.6%；在ModelNet40分类任务中，使用相同的输入数据大小（1024个点）和特征（仅坐标），文章所提方法也明显强于PointNet。
- 基于点集的方法甚至可以达到与成熟的图像CNN相当或更好的性能。在MNIST数据集上，文章基于2D点集的方法达到的准确率接近Network in Network CNN；在ModelNet40数据集上，使用法线信息的文章所提方法显著优于之前的最先进方法MVCNN。
- 对采样密度变化的鲁棒性：从现实世界直接捕获的传感器数据通常存在严重的不规则采样问题。文章通过在测试时随机丢弃点（见图4左）来验证网络对非均匀和稀疏数据的鲁棒性。从图4右可以看出，采用随机输入丢弃训练的多尺度分组（MSG+DP）和多分辨率分组（MRG+DP）在处理采样密度变化时表现出很强的鲁棒性。MSG+DP在测试点数量从1024减少到256时，性能下降不到1%，并且在几乎所有采样密度下都比其他方法表现更好。PointNet (vanilla)由于专注于全局抽象而非细节，在密度变化下也有一定的鲁棒性，但相比之下，其细节信息的缺失使其性能不如文章所提方法。单尺度分组的消融模型SSG在处理稀疏采样密度时无法很好地泛化，而SSG+DP在训练时通过随机丢弃点改进了这个问题。
这些实验结果充分验证了PointNet++在欧几里得度量空间中点集分类任务上的有效性和优越性，以及对采样密度变化的鲁棒性，展示了该模型在复杂实际数据情况下的良好性能。
1.4.2 用于语义场景标记的点集分割
为了验证所提方法适用于大规模点云分析，文章在语义场景标记任务上进行评估。该任务的目标是预测室内扫描中点的语义对象标签。文献[5]（A. Dai, A. X. Chang, M. Savva, M. Halber, T. Funkhouser, and M. Nießner. Scannet: Richly-annotated 3d reconstructions of indoor scenes. arXiv preprint arXiv:1702.04405, 2017.）提供了一个基于体素化扫描的全卷积神经网络的基线方法，该方法仅依赖扫描几何信息而非RGB信息，并以每个体素为基础报告准确率。为了进行公平比较，文章在所有实验中去除RGB信息，并按照文献[5]将点云标签预测转换为体素标签。同时，文章也与文献[20]（C. R. Qi, H. Su, K. Mo, and L. J. Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. arXiv preprint arXiv:1612.00593, 2016.）进行了比较，准确率以每个体素为基础，如图5（蓝色条形）所示：
- 性能优势：文章所提方法在语义场景标记任务上大幅超越了所有基线方法。与在体素化扫描上学习的文献[5]相比，文章方法直接在点云上学习，避免了额外的量化误差，并且通过数据相关采样实现了更有效的学习。与文献[20]相比，文章方法引入了分层特征学习，能够捕获不同尺度的几何特征，这对于多层次理解场景和标记不同大小的物体非常重要。图6展示了示例场景标记结果，PointNet能够正确捕获房间的整体布局，但无法识别家具；而文章所提方法在分割除房间布局之外的物体方面表现得更好。
- 对采样密度变化的鲁棒性：为了测试训练好的模型在具有非均匀采样密度的扫描上的性能，文章合成了类似于图1的ScanNet场景的虚拟扫描，并在这些数据上评估网络。具体的虚拟扫描生成方法在补充材料中给出。文章在三种设置（SSG、MSG+DP、MRG+DP）下评估框架，并与基线方法文献[20]进行比较。性能比较结果如图5（黄色条形）所示，SSG的性能由于从均匀点云到虚拟扫描场景的采样密度变化而大幅下降；MRG网络对采样密度变化更具鲁棒性，因为在采样稀疏时它能够自动切换到描述更粗粒度的特征；尽管训练数据（具有随机丢弃的均匀点）和具有非均匀密度的扫描数据之间存在域差距，但MSG网络仅受到轻微影响，并且在比较的方法中获得了最高的准确率。这些结果证明了文章所设计的密度自适应层的有效性。
通过语义场景标记任务的实验，进一步验证了PointNet++在处理大规模点云数据进行语义分割的能力，以及在面对非均匀采样密度时的鲁棒性，展示了该模型在实际场景应用中的潜力。
1.4.3 非欧几里得度量空间中的点集分类
在本节中，文章展示了所提方法在非欧几里得空间中的泛化能力。在非刚性形状分类（见图7）任务中，一个优秀的分类器应能正确地将图7中(a)和(c)分类为同一类别，尽管它们的姿态不同，这需要对内在结构的理解。SHREC15中的形状是嵌入在3D空间中的2D表面，沿着表面的测地距离自然地诱导出一个度量空间。文章通过实验表明，在这个度量空间中采用PointNet++是捕获底层点集内在结构的有效方法。
- 实验设置：对于SHREC15中的每个形状，首先构建由成对测地距离诱导的度量空间。文章遵循文献[23]（R. M. Rustamov, Y. Lipman, and T. Funkhouser. Interior distance using barycentric coordinates. In Computer Graphics Forum, volume 28, pages 1279–1288. Wiley Online Library, 2009.）获取模拟测地距离的嵌入度量。接着，在这个度量空间中提取内在点特征，包括波核签名（WKS）、热核签名（HKS）和多尺度高斯曲率，将这些特征作为输入，然后根据底层度量空间对进行采样和分组。这样，网络就能学习捕获不受形状特定姿态影响的多尺度内在结构。对比实验设置了不同的选择，包括使用XYZ坐标作为点特征，或使用欧几里得空间
R^{3}作为底层度量空间。 - 结果分析：文章将所提方法与之前的最先进方法文献[14]（L. Luciano and A. B. Hamza. Deep learning with geodesic moments for 3d shape classification. Pattern Recognition Letters, 2017.）进行比较，结果如表3所示。文献[14]提取测地矩作为形状特征，并使用堆叠稀疏自动编码器处理这些特征以预测形状类别。文章使用非欧几里得度量空间和内在特征的方法在所有设置下都取得了最佳性能，大幅超越了文献[14]。比较文章所提方法的不同设置可以发现：内在特征对于非刚性形状分类非常重要，使用XYZ特征无法揭示内在结构，并且受姿态变化影响很大；使用测地邻域比欧几里得邻域更有益，因为欧几里得邻域可能会包含表面上距离较远的点，当形状发生非刚性变形时，邻域会发生显著变化，这给有效的权重共享带来困难，导致局部结构变得复杂，而表面上的测地邻域避免了这个问题，提高了学习效果。
通过在非欧几里得度量空间中的实验，充分验证了PointNet++能够有效捕获非刚性形状的内在结构，展示了该方法在不同类型度量空间中的良好适应性和泛化能力，拓宽了其在复杂形状分析任务中的应用范围。
1.4.4 特征可视化
在图8中，文章对分层网络第一层内核所学习到的内容进行了可视化。具体做法是在空间中创建一个体素网格，聚合那些在网格单元中激活特定神经元程度最高的局部点集（使用最高的100个示例）。保留那些获得高投票数的网格单元，并将其转换回3D点云，这些点云就代表了神经元所识别的模式。由于模型是在主要包含家具的ModelNet40数据集上进行训练的，所以在可视化结果中可以看到平面、双平面、线条、角落等结构。
这种特征可视化方式，直观地展示了网络在学习过程中所捕捉到的局部模式，帮助读者更好地理解网络是如何从点集中提取特征的。它为研究网络的行为和性能提供了一个可视化的视角，进一步验证了网络能够学习到有意义的几何特征，有助于解释为什么PointNet++在点集分类和分割等任务中能够取得良好的效果。
1.5 相关工作
分层特征学习的理念取得了巨大成功，在众多学习模型中，卷积神经网络（Convolutional Neural Network, CNN）[10, 25, 8]是其中最为突出的代表之一。然而，卷积操作并不适用于具有距离度量的无序点集，而这正是本文的研究重点。
近期有一些工作[20, 28]探索了如何将深度学习应用于无序集合，但它们忽略了点集所具有的底层距离度量。这就导致这些方法无法捕捉点的局部上下文信息，并且对全局集合的平移和归一化较为敏感。而在本文中，研究目标是处理从度量空间中采样得到的点，通过在设计中明确考虑底层距离度量来解决上述问题。
从度量空间采样得到的点通常存在噪声，且采样密度不均匀，这影响了有效的点特征提取，给学习过程带来了困难。其中一个关键问题是为点特征设计选择合适的尺度。此前，在几何处理领域以及摄影测量和遥感领域已经开发了多种方法[19, 17, 2, 6, 7, 30]来解决该问题。与这些工作不同，本文的方法以端到端的方式学习提取点特征，并平衡多个特征尺度。
在3D度量空间中，除了点集，还有一些常用于深度学习的表示形式，包括体素网格[21, 22, 29]和几何图[3, 15, 33]。然而，这些工作都没有明确考虑非均匀采样密度的问题。
这部分内容通过与CNN、其他处理无序集合的深度学习方法、传统特征尺度选择方法以及3D度量空间中其他数据表示形式的深度学习方法进行对比，突出了本文方法在处理点集数据时，针对距离度量利用、特征提取和尺度平衡以及对非均匀采样密度处理等方面的独特性和优势，从而体现了研究的创新性和价值。
1.6 结论
在本文中，研究人员提出了PointNet++，这是一种强大的神经网络架构，可用于处理在度量空间中采样的点集。PointNet++通过对输入点集进行嵌套划分递归地发挥作用，能够有效地学习与距离度量相关的分层特征。
为解决点采样不均匀的问题，研究人员提出了两种新颖的集合抽象层，它们可以根据局部点密度智能地聚合多尺度信息。这些改进使得模型在具有挑战性的3D点云基准测试中达到了最先进的性能。
在未来研究方向上，研究人员认为值得思考如何通过在每个局部区域共享更多计算来加速所提网络（尤其是MSG和MRG层）的推理速度。另外，在高维度量空间中寻找应用也很有意义，因为基于CNN的方法在高维度量空间中计算上可能不可行，而PointNet++方法有望很好地扩展应用。
参考文献见原文
2 PointNet模型
models
2.1 基础算法实现
pointnet_utils.py：PointNet的T-Net和编码器，以及特征转换矩阵的正则化函数
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
from torch.autograd import Variable
import numpy as np
import torch.nn.functional as F
# STN3d: T-Net 3 * 3 transform
class STN3d(nn.Module):
def __init__(self, channel):
super(STN3d, self).__init__()
# torch.nn.Conv1d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True)
self.conv1 = torch.nn.Conv1d(channel, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 9) # 9 = 3 * 3
self.relu = nn.ReLU()
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(256)
def forward(self, x):
batchsize = x.size()[0]
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.relu(self.bn3(self.conv3(x)))
x = torch.max(x, 2, keepdim=True)[0] # Symmetric function: max pooling
x = x.view(-1, 1024) # 参数拉直（展平）
x = F.relu(self.bn4(self.fc1(x)))
x = F.relu(self.bn5(self.fc2(x)))
x = self.fc3(x)
# 展平的单位矩阵：np.array([1, 0, 0, 0, 1, 0, 0, 0, 1])
iden = Variable(torch.from_numpy(np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]).astype(np.float32))).view(1, 9).repeat(batchsize, 1)
if x.is_cuda:
iden = iden.cuda()
x = x + iden # affine transformation（仿射变换）
x = x.view(-1, 3, 3) # batch_size * 3 * 3
return x
# STNkd: T-Net k * k transform (k默认为64)
class STNkd(nn.Module):
def __init__(self, k=64):
super(STNkd, self).__init__()
self.conv1 = torch.nn.Conv1d(k, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, k * k)
self.relu = nn.ReLU()
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(256)
self.k = k
def forward(self, x):
batchsize = x.size()[0]
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.relu(self.bn3(self.conv3(x)))
x = torch.max(x, 2, keepdim=True)[0]
x = x.view(-1, 1024)
x = F.relu(self.bn4(self.fc1(x)))
x = F.relu(self.bn5(self.fc2(x)))
x = self.fc3(x)
# np.eye()：生成单位矩阵
iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1, self.k * self.k).repeat(batchsize, 1)
if x.is_cuda:
iden = iden.cuda()
x = x + iden
x = x.view(-1, self.k, self.k)
return x
# PointNet编码器：input points (-> local embedding, if seg) -> global feature
class PointNetEncoder(nn.Module):
def __init__(self, global_feat=True, feature_transform=False, channel=3):
super(PointNetEncoder, self).__init__()
self.stn = STN3d(channel) # T-Net 3 * 3 transform
self.conv1 = torch.nn.Conv1d(channel, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.global_feat = global_feat
self.feature_transform = feature_transform
if self.feature_transform: # 如果需要特征转换，则使用T-Net 64 * 64 transform
self.fstn = STNkd(k=64)
def forward(self, x):
B, D, N = x.size() # batch_size, 3（xyz坐标）或6（xyz坐标+法向量）, 1024（一个物体所取点的数目）
trans = self.stn(x)
x = x.transpose(2, 1) # 转置，交换维度
if D > 3: # 如果输入点云包含法向量，则将法向量与坐标分开，法向量是第4维及之后的数据，坐标是前3维
feature = x[:, :, 3:]
x = x[:, :, :3]
# 对输入的点云进行输入转换（input transform）：对两个tensor进行矩阵乘法
# torch.bmm()要求两个tensor的维度为[batch_size, n, m]和[batch_size, m, p]，返回的tensor维度为[batch_size, n, p]
x = torch.bmm(x, trans)
if D > 3:
x = torch.cat([x, feature], dim=2) # 将转换后的坐标与法向量拼接
x = x.transpose(2, 1)
x = F.relu(self.bn1(self.conv1(x)))
if self.feature_transform:
trans_feat = self.fstn(x)
x = x.transpose(2, 1)
x = torch.bmm(x, trans_feat) # 对输入的点云进行特征转换（feature transform）
x = x.transpose(2, 1)
else:
trans_feat = None
pointfeat = x # 局部特征
x = F.relu(self.bn2(self.conv2(x)))
x = self.bn3(self.conv3(x))
x = torch.max(x, 2, keepdim=True)[0] # 最大池化获取全局特征
x = x.view(-1, 1024)
if self.global_feat: # 如果需要全局特征（多为分割任务），则返回全局特征x
return x, trans, trans_feat
else: # 否则返回全局特征和局部特征的拼接（torch.cat()）
x = x.view(-1, 1024, 1).repeat(1, 1, N)
return torch.cat([x, pointfeat], 1), trans, trans_feat
# 对特征转换矩阵做正则化（将特征变换矩阵约束为接近正交矩阵，$L_{reg}=||I-AA^T||^2_F$）
def feature_transform_reguliarzer(trans):
d = trans.size()[1]
I = torch.eye(d)[None, :, :] # torch.eye(n, m=None, out=None)：返回一个2维单位矩阵
if trans.is_cuda:
I = I.cuda()
loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2, 1)) - I, dim=(1, 2))) # 正则化损失函数
return loss2.2 物体分类网络模型
pointnet_cls.py
import torch.nn as nn
import torch.utils.data
import torch.nn.functional as F
from pointnet_utils import PointNetEncoder, feature_transform_reguliarzer
class get_model(nn.Module):
def __init__(self, k=40, normal_channel=True):
super(get_model, self).__init__()
if normal_channel: # 使用法向量
channel = 6
else:
channel = 3
self.feat = PointNetEncoder(global_feat=True, feature_transform=True, channel=channel)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, k) # k为类别数（默认为40）
self.dropout = nn.Dropout(p=0.4)
self.bn1 = nn.BatchNorm1d(512)
self.bn2 = nn.BatchNorm1d(256)
self.relu = nn.ReLU()
def forward(self, x):
x, trans, trans_feat = self.feat(x)
x = F.relu(self.bn1(self.fc1(x)))
x = F.relu(self.bn2(self.dropout(self.fc2(x))))
x = self.fc3(x)
x = F.log_softmax(x, dim=1) # 计算对数概率
return x, trans_feat
class get_loss(torch.nn.Module):
def __init__(self, mat_diff_loss_scale=0.001):
super(get_loss, self).__init__()
self.mat_diff_loss_scale = mat_diff_loss_scale
def forward(self, pred, target, trans_feat):
# F.nll_loss(pred, target)输入一个对数概率张量和一个目标标签，不会做softmax，返回一个标量损失
# 适用于最后一层为log_softmax的模型
# 常用的交叉熵损失函数与此类似，唯一不同是其会做softmax
loss = F.nll_loss(pred, target) # 分类损失
mat_diff_loss = feature_transform_reguliarzer(trans_feat) # 特征变换正则化损失
total_loss = loss + mat_diff_loss * self.mat_diff_loss_scale
return total_loss2.3 部件分割网络模型
pointnet_part_seg.py（重写了编码器，总体流程类似pointnet_cls.py）
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
import torch.nn.functional as F
from pointnet_utils import STN3d, STNkd, feature_transform_reguliarzer
class get_model(nn.Module):
def __init__(self, part_num=50, normal_channel=True):
super(get_model, self).__init__()
if normal_channel:
channel = 6
else:
channel = 3
self.part_num = part_num # 部件类别数量，ShapeNet中为50
self.stn = STN3d(channel)
self.conv1 = torch.nn.Conv1d(channel, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 128, 1)
self.conv4 = torch.nn.Conv1d(128, 512, 1)
self.conv5 = torch.nn.Conv1d(512, 2048, 1)
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(128)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(2048)
self.fstn = STNkd(k=128) # 维度为128的feature transform
self.convs1 = torch.nn.Conv1d(4944, 256, 1)
self.convs2 = torch.nn.Conv1d(256, 256, 1)
self.convs3 = torch.nn.Conv1d(256, 128, 1)
self.convs4 = torch.nn.Conv1d(128, part_num, 1)
self.bns1 = nn.BatchNorm1d(256)
self.bns2 = nn.BatchNorm1d(256)
self.bns3 = nn.BatchNorm1d(128)
def forward(self, point_cloud, label):
B, D, N = point_cloud.size()
trans = self.stn(point_cloud)
point_cloud = point_cloud.transpose(2, 1)
if D > 3:
point_cloud, feature = point_cloud.split(3, dim=2)
point_cloud = torch.bmm(point_cloud, trans)
if D > 3:
point_cloud = torch.cat([point_cloud, feature], dim=2)
point_cloud = point_cloud.transpose(2, 1)
out1 = F.relu(self.bn1(self.conv1(point_cloud)))
out2 = F.relu(self.bn2(self.conv2(out1)))
out3 = F.relu(self.bn3(self.conv3(out2)))
trans_feat = self.fstn(out3)
x = out3.transpose(2, 1)
net_transformed = torch.bmm(x, trans_feat)
net_transformed = net_transformed.transpose(2, 1)
out4 = F.relu(self.bn4(self.conv4(net_transformed)))
out5 = self.bn5(self.conv5(out4))
out_max = torch.max(out5, 2, keepdim=True)[0] # 池化后的维度为[B, 2048, 1]（2048个元素）
out_max = out_max.view(-1, 2048)
out_max = torch.cat([out_max,label.squeeze(1)],1)
expand = out_max.view(-1, 2048+16, 1).repeat(1, 1, N) # 16个物体类别
concat = torch.cat([expand, out1, out2, out3, out4, out5], 1) # 局部特征与全局特征拼接
net = F.relu(self.bns1(self.convs1(concat)))
net = F.relu(self.bns2(self.convs2(net)))
net = F.relu(self.bns3(self.convs3(net)))
net = self.convs4(net)
net = net.transpose(2, 1).contiguous()
net = F.log_softmax(net.view(-1, self.part_num), dim=-1)
net = net.view(B, N, self.part_num) # [B, N, 50]
return net, trans_feat
class get_loss(torch.nn.Module):
def __init__(self, mat_diff_loss_scale=0.001):
super(get_loss, self).__init__()
self.mat_diff_loss_scale = mat_diff_loss_scale
def forward(self, pred, target, trans_feat):
loss = F.nll_loss(pred, target)
mat_diff_loss = feature_transform_reguliarzer(trans_feat)
total_loss = loss + mat_diff_loss * self.mat_diff_loss_scale
return total_loss2.4 场景分割网络模型
pointnet_sem_seg.py（基本算法同pointnet_cls.py）
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
import torch.nn.functional as F
from pointnet_utils import PointNetEncoder, feature_transform_reguliarzer
class get_model(nn.Module):
def __init__(self, num_class):
super(get_model, self).__init__()
self.k = num_class # 类别数
self.feat = PointNetEncoder(global_feat=False, feature_transform=True, channel=9) # 注意global_feat=False
self.conv1 = torch.nn.Conv1d(1088, 512, 1)
self.conv2 = torch.nn.Conv1d(512, 256, 1)
self.conv3 = torch.nn.Conv1d(256, 128, 1)
self.conv4 = torch.nn.Conv1d(128, self.k, 1)
self.bn1 = nn.BatchNorm1d(512)
self.bn2 = nn.BatchNorm1d(256)
self.bn3 = nn.BatchNorm1d(128)
def forward(self, x):
batchsize = x.size()[0]
n_pts = x.size()[2]
x, trans, trans_feat = self.feat(x)
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.relu(self.bn3(self.conv3(x)))
x = self.conv4(x)
x = x.transpose(2,1).contiguous()
x = F.log_softmax(x.view(-1,self.k), dim=-1)
x = x.view(batchsize, n_pts, self.k)
return x, trans_feat
class get_loss(torch.nn.Module):
def __init__(self, mat_diff_loss_scale=0.001):
super(get_loss, self).__init__()
self.mat_diff_loss_scale = mat_diff_loss_scale
def forward(self, pred, target, trans_feat, weight):
loss = F.nll_loss(pred, target, weight = weight)
mat_diff_loss = feature_transform_reguliarzer(trans_feat)
total_loss = loss + mat_diff_loss * self.mat_diff_loss_scale
return total_loss
if __name__ == '__main__':
model = get_model(13)
xyz = torch.rand(12, 3, 2048)
(model(xyz))3 PointNet++模型
models
3.1 基础算法实现
pointnet2_utils.py：PointNet++的SA层（最远点采样、球查询）和上采样（线性插值、MLP），以及相关工具函数
import torch
import torch.nn as nn
import torch.nn.functional as F
from time import time
import numpy as np
def timeit(tag, t):
print("{}: {}s".format(tag, time() - t))
return time()
# 归一化点云：使用以centroid为中心，以1为半径的球体
def pc_normalize(pc):
l = pc.shape[0]
centroid = np.mean(pc, axis=0)
pc = pc - centroid
m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
pc = pc / m
return pc
# 计算点云中每个点与查询点之间的距离
# 输入两组点，输出两组点两两之间的距离，即n*m的矩阵
# 在训练中数据以mini-batch的形式输入，因此第1维为B
def square_distance(src, dst):
"""
Calculate Euclid distance between each two points.
src^T * dst = xn * xm + yn * ym + zn * zm;
sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
= sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
Input:
src: source points, [B, N, C]
dst: target points, [B, M, C]
Output:
dist: per-point square distance, [B, N, M]
"""
B, N, _ = src.shape
_, M, _ = dst.shape
dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
dist += torch.sum(src ** 2, -1).view(B, N, 1)
dist += torch.sum(dst ** 2, -1).view(B, 1, M)
return dist
# 按照输入的点云数据和索引，返回索引对应的数据
# 设points形状为[B, N, C]，若idx形状为[B, D1, ..., Dn]，则按idx中的维度结构将其提取为[B, D1, ..., Dn, C]
# 例如points形状为[B, 2408, 3]，idx=[5, 666, 1000, 2000]，则返回batch中第5、666、1000、2000个点组成的B * 4 * 3的点云
def index_points(points, idx):
"""
Input:
points: input points data, [B, N, C]
idx: sample index data, [B, S]
Return:
new_points:, indexed points data, [B, S, C]
"""
device = points.device
B = points.shape[0]
view_shape = list(idx.shape)
view_shape[1:] = [1] * (len(view_shape) - 1)
repeat_shape = list(idx.shape)
repeat_shape[0] = 1
batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
new_points = points[batch_indices, idx, :]
return new_points3.1.1 最远点采样（FPS）
# 最远点采样：从输入点云中按照所需的点的个数npoint采样出足够多的采样点，并且点与点之间的距离尽可能大
# 返回npoint个采样点在原始点云中的索引（找centroids）
def farthest_point_sample(xyz, npoint):
"""
Input:
xyz: pointcloud data, [B, N, 3]
npoint: number of samples
Return:
centroids: sampled pointcloud index, [B, npoint]
"""
device = xyz.device
B, N, C = xyz.shape
# centroids矩阵用于存储npoints个采样点的索引位置，大小为B * npoint
centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
# distance矩阵（B * N）用于存储batch中每个点与采样点之间的距离，初始化为一个很大的数
distance = torch.ones(B, N).to(device) * 1e10
# farthest存储当前距离采样点最远的点的索引，随机初始化范围为0~N；初始化B个（每个batch都随机一个初始最远点）
farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
# batch_indices用于存储当前batch中每个点所属的batch索引，初始化为0~(B-1)的数组
batch_indices = torch.arange(B, dtype=torch.long).to(device)
# 迭代npoint次，每次迭代中找到距离当前采样点最远的点，将其作为新的采样点
for i in range(npoint):
centroids[:, i] = farthest # 将当前采样点（下标i）设为当前的最远点farthest
centroid = xyz[batch_indices, farthest, :].view(B, 1, 3) # 取出该中心点（centroid）的坐标
dist = torch.sum((xyz - centroid) ** 2, -1) # 计算所有点到centroid的欧氏距离，存储于dist
# 建立一个mask：若dist中的元素小于distance矩阵中保存的距离值，则更新distance中的对应值
# 随着迭代的继续，distance矩阵中保存的值会越来越小，其相当于记录着某batch中每个点与所有已出现采样点之间的最小距离
mask = dist < distance
distance[mask] = dist[mask]
# 从distance矩阵取出的最远点为farthest，即当前的最远点，继续下一轮迭代
farthest = torch.max(distance, -1)[1]
return centroids3.1.2 球查询（Ball Query）
# 球查询：寻找球形邻域中的点
# 输入中radius为球半径，nsample为每个邻域中点要采样的点，xyz为所有点（[B, N, 3]），new_xyz为查询点（[B, S, 3]）
# 输出为每个样本的每个球形领域的nsample个采样点集的索引[B, S, nsample]
def query_ball_point(radius, nsample, xyz, new_xyz):
"""
Input:
radius: local region radius
nsample: max sample number in local region
xyz: all points, [B, N, 3]
new_xyz: query points, [B, S, 3]
Return:
group_idx: grouped points index, [B, S, nsample]
"""
device = xyz.device
B, N, C = xyz.shape
_, S, _ = new_xyz.shape
# group_idx矩阵（[B, S, N]）用于存储每个查询点周围球查询中包含的点的索引
group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
# sqrdists矩阵（[B, S, N]）记录S个中心点（new_xyz）与所有点（xyz）之间的欧几里得距离
sqrdists = square_distance(new_xyz, xyz)
# 将所有距离大于radius^2的点的索引直接置为N（球查询中不包含这些点），升序排序筛掉这些点，取出前nsample个点
group_idx[sqrdists > radius ** 2] = N
group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
# 当球形区域内不足nsample个点时，前nsample个点中亦会存在置为N的点，需被舍弃，可直接用第1个点来代替
# group_first实为将第1个点复制至[B, S, K]的维度，便于后续替换
group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
mask = group_idx == N
group_idx[mask] = group_first[mask] # 替换
return group_idx3.1.3 采样+分组
# 采样+分组：将整个点云分散成局部的group，对每一个group都可以用PointNet单独提取局部特征
def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False):
"""
Input:
npoint: points sampled in farthest point sampling
radius: search radius in local region
nsample: how many points in each local region
xyz: input points position data, [B, N, 3]
points: input points data, [B, N, D]
Return:
new_xyz: sampled points position data, [B, npoint, nsample, 3]
new_points: sampled points data, [B, npoint, nsample, 3+D]
"""
B, N, C = xyz.shape
S = npoint
# 进行最远点采样（FPS），挑出的采样点作为new_xyz：
# 先用farthest_point_sample()最远点采样得到采样点的索引fps_idx
# 再用index_points()根据索引从xyz中挑出采样点，作为new_xyz
fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
new_xyz = index_points(xyz, fps_idx) # 中心点（centroids）
# 进行球查询
# 先用query_ball_point()球查询得到npoint个球形区域中每个区域的nsample个采样点的索引，存储于idx（[B, npoint, nsample]）
# 再用index_points()根据索引从xyz中挑出采样点，作为grouped_xyz
idx = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
# grouped_xyz减去采样点（即centroids）
grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
# 若每个点上有新特征的维度，则拼接新特征与旧特征，否则直接为旧特征
if points is not None:
grouped_points = index_points(points, idx)
new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
else:
new_points = grouped_xyz_norm
if returnfps:
return new_xyz, new_points, grouped_xyz, fps_idx
else:
return new_xyz, new_points
# 直接将所有点作为一个group，即npoint=1
def sample_and_group_all(xyz, points):
"""
Input:
xyz: input points position data, [B, N, 3]
points: input points data, [B, N, D]
Return:
new_xyz: sampled points position data, [B, 1, 3]
new_points: sampled points data, [B, 1, N, 3+D]
"""
device = xyz.device
B, N, C = xyz.shape
new_xyz = torch.zeros(B, 1, C).to(device)
grouped_xyz = xyz.view(B, 1, N, C)
if points is not None:
new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
else:
new_points = grouped_xyz
return new_xyz, new_points3.1.4 集合抽象层（SA）
整个集合抽象层（Set Abstraction）的两种封装
# 普通的集合抽象层
# 首先通过sample_and_group()形成局部group，再对局部group中每一个点做MLP操作，最后进行局部的最大池化，得到局部的全局特征
# 其中mlp传入的是MLP的输出通道数，group_all指是否将所有点作为一个group
class PointNetSetAbstraction(nn.Module):
def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
super(PointNetSetAbstraction, self).__init__()
self.npoint = npoint
self.radius = radius
self.nsample = nsample
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
last_channel = in_channel
for out_channel in mlp:
self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
self.mlp_bns.append(nn.BatchNorm2d(out_channel))
last_channel = out_channel
self.group_all = group_all
def forward(self, xyz, points):
"""
Input:
xyz: input points position data, [B, C, N]
points: input points data, [B, D, N]
Return:
new_xyz: sampled points position data, [B, C, S]
new_points_concat: sample points feature data, [B, D', S]
"""
xyz = xyz.permute(0, 2, 1)
if points is not None:
points = points.permute(0, 2, 1)
# 形成局部group
if self.group_all:
new_xyz, new_points = sample_and_group_all(xyz, points)
else:
new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
# new_xyz: sampled points position data, [B, npoint, C]
# new_points: sampled points data, [B, npoint, nsample, C+D]
new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
# PointNet操作：对局部group中的每个点进行MLP操作（1x1卷积相当于全连接层）
for i, conv in enumerate(self.mlp_convs):
bn = self.mlp_bns[i]
new_points = F.relu(bn(conv(new_points)))
# 局部的最大池化，得到局部的全局特征
new_points = torch.max(new_points, 2)[0]
new_xyz = new_xyz.permute(0, 2, 1)
return new_xyz, new_points # 返回采样点（centroids）和局部的全局特征
# 实现多尺度分组（MSG）方法的集合抽象层
# 其中radius_list、nsample_list、mlp_list传入的分别为不同尺度下的球查询半径列表、采样点数列表、MLP输出通道数二维列表
class PointNetSetAbstractionMsg(nn.Module):
def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
super(PointNetSetAbstractionMsg, self).__init__()
self.npoint = npoint
self.radius_list = radius_list
self.nsample_list = nsample_list
self.conv_blocks = nn.ModuleList()
self.bn_blocks = nn.ModuleList()
for i in range(len(mlp_list)):
convs = nn.ModuleList()
bns = nn.ModuleList()
last_channel = in_channel + 3
for out_channel in mlp_list[i]:
convs.append(nn.Conv2d(last_channel, out_channel, 1))
bns.append(nn.BatchNorm2d(out_channel))
last_channel = out_channel
self.conv_blocks.append(convs)
self.bn_blocks.append(bns)
def forward(self, xyz, points):
"""
Input:
xyz: input points position data, [B, C, N]
points: input points data, [B, D, N]
Return:
new_xyz: sampled points position data, [B, C, S]
new_points_concat: sample points feature data, [B, D', S]
"""
xyz = xyz.permute(0, 2, 1)
if points is not None:
points = points.permute(0, 2, 1)
B, N, C = xyz.shape
S = self.npoint
# 最远点采样
new_xyz = index_points(xyz, farthest_point_sample(xyz, S))
# new_points_list用于保存不同半径下的点云特征，最后拼接到一起
new_points_list = []
for i, radius in enumerate(self.radius_list):
K = self.nsample_list[i]
# 球查询
group_idx = query_ball_point(radius, K, xyz, new_xyz)
# 按照输入的点云数据和索引返回索引的点云数据
grouped_xyz = index_points(xyz, group_idx)
grouped_xyz -= new_xyz.view(B, S, 1, C)
if points is not None:
grouped_points = index_points(points, group_idx)
# 拼接点特征数据和点坐标数据
grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
else:
grouped_points = grouped_xyz
grouped_points = grouped_points.permute(0, 3, 2, 1) # [B, D, K, S]
for j in range(len(self.conv_blocks[i])):
conv = self.conv_blocks[i][j]
bn = self.bn_blocks[i][j]
grouped_points = F.relu(bn(conv(grouped_points)))
# 最大池化，获得局部的全局特征
new_points = torch.max(grouped_points, 2)[0] # [B, D', S]
new_points_list.append(new_points)
new_xyz = new_xyz.permute(0, 2, 1)
# 最后拼接不同半径下的点云特征
new_points_concat = torch.cat(new_points_list, dim=1)
return new_xyz, new_points_concat3.1.5 特征传播（Feature Propagation）
# Feature Propagation（分割时的上采样）主要通过线性插值（基于距离插值）和MLP完成
# 当点的个数为1时，使用repeat()直接复制成N个点
# 当点的个数大于1时，使用线性插值进行上采样
# 拼接上下采样对应点的SA层的特征，再对拼接后的每个点都做一个MLP
class PointNetFeaturePropagation(nn.Module):
def __init__(self, in_channel, mlp):
super(PointNetFeaturePropagation, self).__init__()
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
last_channel = in_channel
for out_channel in mlp:
self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
self.mlp_bns.append(nn.BatchNorm1d(out_channel))
last_channel = out_channel
def forward(self, xyz1, xyz2, points1, points2):
"""
Input:
xyz1: input points position data, [B, C, N]
xyz2: sampled input points position data, [B, C, S]
points1: input points data, [B, D, N]
points2: input points data, [B, D, S]
Return:
new_points: upsampled points data, [B, D', N]
"""
xyz1 = xyz1.permute(0, 2, 1)
xyz2 = xyz2.permute(0, 2, 1)
points2 = points2.permute(0, 2, 1)
B, N, C = xyz1.shape
_, S, _ = xyz2.shape
if S == 1: # 直接复制
interpolated_points = points2.repeat(1, N, 1)
else: # 线性插值
dists = square_distance(xyz1, xyz2)
dists, idx = dists.sort(dim=-1)
dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3]
dist_recip = 1.0 / (dists + 1e-8) # 距离越远的点权重越小
norm = torch.sum(dist_recip, dim=2, keepdim=True)
weight = dist_recip / norm # 对于每个点的权重做全局归一化
# 获得插值点
interpolated_points = torch.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2)
if points1 is not None:
points1 = points1.permute(0, 2, 1)
# 拼接上下采样前对应点SA层的特征
new_points = torch.cat([points1, interpolated_points], dim=-1)
else:
new_points = interpolated_points
new_points = new_points.permute(0, 2, 1)
# 对拼接后每个点都做一个MLP
for i, conv in enumerate(self.mlp_convs):
bn = self.mlp_bns[i]
new_points = F.relu(bn(conv(new_points)))
return new_points3.2 物体分类网络模型
pointnet2_cls_ssg.py（单尺度）
import torch.nn as nn
import torch.nn.functional as F
from pointnet2_utils import PointNetSetAbstraction
class get_model(nn.Module):
def __init__(self, num_class, normal_channel=True):
super(get_model, self).__init__()
in_channel = 6 if normal_channel else 3
self.normal_channel = normal_channel
self.sa1 = PointNetSetAbstraction(npoint=512, radius=0.2, nsample=32, in_channel=in_channel, mlp=[64, 64, 128], group_all=False)
self.sa2 = PointNetSetAbstraction(npoint=128, radius=0.4, nsample=64, in_channel=128 + 3, mlp=[128, 128, 256], group_all=False)
self.sa3 = PointNetSetAbstraction(npoint=None, radius=None, nsample=None, in_channel=256 + 3, mlp=[256, 512, 1024], group_all=True)
self.fc1 = nn.Linear(1024, 512)
self.bn1 = nn.BatchNorm1d(512)
self.drop1 = nn.Dropout(0.4)
self.fc2 = nn.Linear(512, 256)
self.bn2 = nn.BatchNorm1d(256)
self.drop2 = nn.Dropout(0.4)
self.fc3 = nn.Linear(256, num_class)
def forward(self, xyz):
B, _, _ = xyz.shape
if self.normal_channel: # 包含法向量
norm = xyz[:, 3:, :]
xyz = xyz[:, :3, :]
else:
norm = None
l1_xyz, l1_points = self.sa1(xyz, norm)
l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
x = l3_points.view(B, 1024)
x = self.drop1(F.relu(self.bn1(self.fc1(x))))
x = self.drop2(F.relu(self.bn2(self.fc2(x))))
x = self.fc3(x)
x = F.log_softmax(x, -1) # 计算对数概率
return x, l3_points
class get_loss(nn.Module):
def __init__(self):
super(get_loss, self).__init__()
def forward(self, pred, target, trans_feat):
total_loss = F.nll_loss(pred, target)
return total_losspointnet2_cls_msg.py（多尺度）
import torch.nn as nn
import torch.nn.functional as F
from pointnet2_utils import PointNetSetAbstractionMsg, PointNetSetAbstraction
class get_model(nn.Module):
def __init__(self, num_class, normal_channel=True):
super(get_model, self).__init__()
in_channel = 3 if normal_channel else 0 # 注意与SSG不同，此处in_channel为0或3
self.normal_channel = normal_channel
self.sa1 = PointNetSetAbstractionMsg(512, [0.1, 0.2, 0.4], [16, 32, 128], in_channel, [[32, 32, 64], [64, 64, 128], [64, 96, 128]])
self.sa2 = PointNetSetAbstractionMsg(128, [0.2, 0.4, 0.8], [32, 64, 128], 320, [[64, 64, 128], [128, 128, 256], [128, 128, 256]])
self.sa3 = PointNetSetAbstraction(None, None, None, 640 + 3, [256, 512, 1024], True)
self.fc1 = nn.Linear(1024, 512)
self.bn1 = nn.BatchNorm1d(512)
self.drop1 = nn.Dropout(0.4)
self.fc2 = nn.Linear(512, 256)
self.bn2 = nn.BatchNorm1d(256)
self.drop2 = nn.Dropout(0.5)
self.fc3 = nn.Linear(256, num_class)
def forward(self, xyz):
B, _, _ = xyz.shape # 24, _, 1024（其中24为测试时默认值）
if self.normal_channel:
norm = xyz[:, 3:, :]
xyz = xyz[:, :3, :]
else:
norm = None
l1_xyz, l1_points = self.sa1(xyz, norm)
l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
x = l3_points.view(B, 1024)
x = self.drop1(F.relu(self.bn1(self.fc1(x))))
x = self.drop2(F.relu(self.bn2(self.fc2(x))))
x = self.fc3(x)
x = F.log_softmax(x, -1)
return x,l3_points
class get_loss(nn.Module):
def __init__(self):
super(get_loss, self).__init__()
def forward(self, pred, target, trans_feat):
total_loss = F.nll_loss(pred, target)
return total_loss3.3 部件分割网络模型
pointnet2_part_seg_ssg.py（单尺度）
import torch.nn as nn
import torch
import torch.nn.functional as F
from models.pointnet2_utils import PointNetSetAbstraction,PointNetFeaturePropagation
class get_model(nn.Module):
def __init__(self, num_classes, normal_channel=False): # num_part = 50
super(get_model, self).__init__()
if normal_channel:
additional_channel = 3
else:
additional_channel = 0
self.normal_channel = normal_channel
self.sa1 = PointNetSetAbstraction(npoint=512, radius=0.2, nsample=32, in_channel=6 + additional_channel, mlp=[64, 64, 128], group_all=False)
self.sa2 = PointNetSetAbstraction(npoint=128, radius=0.4, nsample=64, in_channel=128 + 3, mlp=[128, 128, 256], group_all=False)
self.sa3 = PointNetSetAbstraction(npoint=None, radius=None, nsample=None, in_channel=256 + 3, mlp=[256, 512, 1024], group_all=True)
self.fp3 = PointNetFeaturePropagation(in_channel=1280, mlp=[256, 256])
self.fp2 = PointNetFeaturePropagation(in_channel=384, mlp=[256, 128])
self.fp1 = PointNetFeaturePropagation(in_channel=128 + 16 + 6 + additional_channel, mlp=[128, 128, 128])
self.conv1 = nn.Conv1d(128, 128, 1)
self.bn1 = nn.BatchNorm1d(128)
self.drop1 = nn.Dropout(0.5)
self.conv2 = nn.Conv1d(128, num_classes, 1)
def forward(self, xyz, cls_label):
# Set Abstraction layers
B, C, N = xyz.shape
if self.normal_channel:
l0_points = xyz
l0_xyz = xyz[:,:3,:]
else:
l0_points = xyz
l0_xyz = xyz
l1_xyz, l1_points = self.sa1(l0_xyz, l0_points)
l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
# Feature Propagation layers
l2_points = self.fp3(l2_xyz, l3_xyz, l2_points, l3_points)
l1_points = self.fp2(l1_xyz, l2_xyz, l1_points, l2_points)
cls_label_one_hot = cls_label.view(B,16,1).repeat(1,1,N)
l0_points = self.fp1(l0_xyz, l1_xyz, torch.cat([cls_label_one_hot,l0_xyz,l0_points],1), l1_points)
# FC layers
feat = F.relu(self.bn1(self.conv1(l0_points)))
x = self.drop1(feat)
x = self.conv2(x)
x = F.log_softmax(x, dim=1)
x = x.permute(0, 2, 1)
return x, l3_points
class get_loss(nn.Module):
def __init__(self):
super(get_loss, self).__init__()
def forward(self, pred, target, trans_feat):
total_loss = F.nll_loss(pred, target)
return total_losspointnet2_part_seg_msg.py（多尺度）
import torch.nn as nn
import torch
import torch.nn.functional as F
from models.pointnet2_utils import PointNetSetAbstractionMsg,PointNetSetAbstraction,PointNetFeaturePropagation
class get_model(nn.Module):
def __init__(self, num_classes, normal_channel=False): # num_part = 50
super(get_model, self).__init__()
if normal_channel:
additional_channel = 3
else:
additional_channel = 0
self.normal_channel = normal_channel
self.sa1 = PointNetSetAbstractionMsg(512, [0.1, 0.2, 0.4], [32, 64, 128], 3 + additional_channel, [[32, 32, 64], [64, 64, 128], [64, 96, 128]])
self.sa2 = PointNetSetAbstractionMsg(128, [0.4,0.8], [64, 128], 128 + 128 + 64, [[128, 128, 256], [128, 196, 256]])
self.sa3 = PointNetSetAbstraction(npoint=None, radius=None, nsample=None, in_channel=512 + 3, mlp=[256, 512, 1024], group_all=True)
self.fp3 = PointNetFeaturePropagation(in_channel=1536, mlp=[256, 256])
self.fp2 = PointNetFeaturePropagation(in_channel=576, mlp=[256, 128])
self.fp1 = PointNetFeaturePropagation(in_channel=150 + additional_channel, mlp=[128, 128])
self.conv1 = nn.Conv1d(128, 128, 1)
self.bn1 = nn.BatchNorm1d(128)
self.drop1 = nn.Dropout(0.5)
self.conv2 = nn.Conv1d(128, num_classes, 1)
def forward(self, xyz, cls_label):
# Set Abstraction layers
B, C, N = xyz.shape
if self.normal_channel:
l0_points = xyz
l0_xyz = xyz[:,:3,:]
else:
l0_points = xyz
l0_xyz = xyz
l1_xyz, l1_points = self.sa1(l0_xyz, l0_points)
l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
# Feature Propagation layers
l2_points = self.fp3(l2_xyz, l3_xyz, l2_points, l3_points)
l1_points = self.fp2(l1_xyz, l2_xyz, l1_points, l2_points)
cls_label_one_hot = cls_label.view(B,16,1).repeat(1,1,N)
l0_points = self.fp1(l0_xyz, l1_xyz, torch.cat([cls_label_one_hot,l0_xyz,l0_points],1), l1_points)
# FC layers
feat = F.relu(self.bn1(self.conv1(l0_points)))
x = self.drop1(feat)
x = self.conv2(x)
x = F.log_softmax(x, dim=1)
x = x.permute(0, 2, 1)
return x, l3_points
class get_loss(nn.Module):
def __init__(self):
super(get_loss, self).__init__()
def forward(self, pred, target, trans_feat):
total_loss = F.nll_loss(pred, target)
return total_loss3.4 场景分割网络模型
pointnet2_sem_seg.py
import torch.nn as nn
import torch.nn.functional as F
from models.pointnet2_utils import PointNetSetAbstraction,PointNetFeaturePropagation
class get_model(nn.Module):
def __init__(self, num_classes):
super(get_model, self).__init__()
self.sa1 = PointNetSetAbstraction(1024, 0.1, 32, 9 + 3, [32, 32, 64], False)
self.sa2 = PointNetSetAbstraction(256, 0.2, 32, 64 + 3, [64, 64, 128], False)
self.sa3 = PointNetSetAbstraction(64, 0.4, 32, 128 + 3, [128, 128, 256], False)
self.sa4 = PointNetSetAbstraction(16, 0.8, 32, 256 + 3, [256, 256, 512], False)
self.fp4 = PointNetFeaturePropagation(768, [256, 256])
self.fp3 = PointNetFeaturePropagation(384, [256, 256])
self.fp2 = PointNetFeaturePropagation(320, [256, 128])
self.fp1 = PointNetFeaturePropagation(128, [128, 128, 128])
self.conv1 = nn.Conv1d(128, 128, 1)
self.bn1 = nn.BatchNorm1d(128)
self.drop1 = nn.Dropout(0.5)
self.conv2 = nn.Conv1d(128, num_classes, 1)
def forward(self, xyz):
l0_points = xyz
l0_xyz = xyz[:,:3,:]
l1_xyz, l1_points = self.sa1(l0_xyz, l0_points)
l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
l4_xyz, l4_points = self.sa4(l3_xyz, l3_points)
l3_points = self.fp4(l3_xyz, l4_xyz, l3_points, l4_points)
l2_points = self.fp3(l2_xyz, l3_xyz, l2_points, l3_points)
l1_points = self.fp2(l1_xyz, l2_xyz, l1_points, l2_points)
l0_points = self.fp1(l0_xyz, l1_xyz, None, l1_points)
x = self.drop1(F.relu(self.bn1(self.conv1(l0_points))))
x = self.conv2(x)
x = F.log_softmax(x, dim=1)
x = x.permute(0, 2, 1)
return x, l4_points
class get_loss(nn.Module):
def __init__(self):
super(get_loss, self).__init__()
def forward(self, pred, target, trans_feat, weight):
total_loss = F.nll_loss(pred, target, weight=weight)
return total_loss
if __name__ == '__main__':
import torch
model = get_model(13)
xyz = torch.rand(6, 9, 2048)
(model(xyz))pointnet2_sem_seg_msg.py（多尺度）
import torch.nn as nn
import torch.nn.functional as F
from models.pointnet2_utils import PointNetSetAbstractionMsg,PointNetFeaturePropagation
class get_model(nn.Module):
def __init__(self, num_classes):
super(get_model, self).__init__()
self.sa1 = PointNetSetAbstractionMsg(1024, [0.05, 0.1], [16, 32], 9, [[16, 16, 32], [32, 32, 64]])
self.sa2 = PointNetSetAbstractionMsg(256, [0.1, 0.2], [16, 32], 32+64, [[64, 64, 128], [64, 96, 128]])
self.sa3 = PointNetSetAbstractionMsg(64, [0.2, 0.4], [16, 32], 128+128, [[128, 196, 256], [128, 196, 256]])
self.sa4 = PointNetSetAbstractionMsg(16, [0.4, 0.8], [16, 32], 256+256, [[256, 256, 512], [256, 384, 512]])
self.fp4 = PointNetFeaturePropagation(512 + 512 + 256 + 256, [256, 256])
self.fp3 = PointNetFeaturePropagation(128 + 128 + 256, [256, 256])
self.fp2 = PointNetFeaturePropagation(32 + 64 + 256, [256, 128])
self.fp1 = PointNetFeaturePropagation(128, [128, 128, 128])
self.conv1 = nn.Conv1d(128, 128, 1)
self.bn1 = nn.BatchNorm1d(128)
self.drop1 = nn.Dropout(0.5)
self.conv2 = nn.Conv1d(128, num_classes, 1)
def forward(self, xyz):
l0_points = xyz
l0_xyz = xyz[:,:3,:]
l1_xyz, l1_points = self.sa1(l0_xyz, l0_points)
l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
l4_xyz, l4_points = self.sa4(l3_xyz, l3_points)
l3_points = self.fp4(l3_xyz, l4_xyz, l3_points, l4_points)
l2_points = self.fp3(l2_xyz, l3_xyz, l2_points, l3_points)
l1_points = self.fp2(l1_xyz, l2_xyz, l1_points, l2_points)
l0_points = self.fp1(l0_xyz, l1_xyz, None, l1_points)
x = self.drop1(F.relu(self.bn1(self.conv1(l0_points))))
x = self.conv2(x)
x = F.log_softmax(x, dim=1)
x = x.permute(0, 2, 1)
return x, l4_points
class get_loss(nn.Module):
def __init__(self):
super(get_loss, self).__init__()
def forward(self, pred, target, trans_feat, weight):
total_loss = F.nll_loss(pred, target, weight=weight)
return total_loss
if __name__ == '__main__':
import torch
model = get_model(13)
xyz = torch.rand(6, 9, 2048)
(model(xyz))4 数据加载
data_utils
4.1 ModelNetDataLoader
ModelNetDataLoader.py：ModelNet40用于训练物体形状分类（40个），训练集有9843个点云，测试集有2468个点云。
import os
import numpy as np
import warnings
import pickle
from tqdm import tqdm
from torch.utils.data import Dataset
warnings.filterwarnings('ignore')
# ModelNet40用于训练物体形状分类（40个），训练集有9843个点云，测试集有2468个点云
# 点云归一化，以centroid为中心，半径为1
def pc_normalize(pc):
centroid = np.mean(pc, axis=0)
pc = pc - centroid
m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
pc = pc / m
return pc
# 最远点采样
def farthest_point_sample(point, npoint):
"""
Input:
xyz: pointcloud data, [N, D]
npoint: number of samples
Return:
centroids: sampled pointcloud index, [npoint, D]
"""
N, D = point.shape
xyz = point[:,:3]
centroids = np.zeros((npoint,))
distance = np.ones((N,)) * 1e10
farthest = np.random.randint(0, N)
for i in range(npoint):
centroids[i] = farthest
centroid = xyz[farthest, :]
dist = np.sum((xyz - centroid) ** 2, -1)
mask = dist < distance
distance[mask] = dist[mask]
farthest = np.argmax(distance, -1)
point = point[centroids.astype(np.int32)]
return point
class ModelNetDataLoader(Dataset):
def __init__(self, root, args, split='train', process_data=False):
self.root = root
self.npoints = args.num_point
self.process_data = process_data
self.uniform = args.use_uniform_sample
self.use_normals = args.use_normals
self.num_category = args.num_category
if self.num_category == 10:
self.catfile = os.path.join(self.root, 'modelnet10_shape_names.txt')
else:
self.catfile = os.path.join(self.root, 'modelnet40_shape_names.txt')
self.cat = [line.rstrip() for line in open(self.catfile)]
self.classes = dict(zip(self.cat, range(len(self.cat))))
shape_ids = {}
if self.num_category == 10:
shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_train.txt'))]
shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_test.txt'))]
else:
shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_train.txt'))]
shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_test.txt'))]
assert (split == 'train' or split == 'test')
shape_names = ['_'.join(x.split('_')[0:-1]) for x in shape_ids[split]]
self.datapath = [(shape_names[i], os.path.join(self.root, shape_names[i], shape_ids[split][i]) + '.txt') for i
in range(len(shape_ids[split]))]
print('The size of %s data is %d' % (split, len(self.datapath)))
if self.uniform:
self.save_path = os.path.join(root, 'modelnet%d_%s_%dpts_fps.dat' % (self.num_category, split, self.npoints))
else:
self.save_path = os.path.join(root, 'modelnet%d_%s_%dpts.dat' % (self.num_category, split, self.npoints))
if self.process_data:
if not os.path.exists(self.save_path):
print('Processing data %s (only running in the first time)...' % self.save_path)
self.list_of_points = [None] * len(self.datapath)
self.list_of_labels = [None] * len(self.datapath)
for index in tqdm(range(len(self.datapath)), total=len(self.datapath)):
fn = self.datapath[index]
cls = self.classes[self.datapath[index][0]]
cls = np.array([cls]).astype(np.int32)
point_set = np.loadtxt(fn[1], delimiter=',').astype(np.float32)
if self.uniform:
point_set = farthest_point_sample(point_set, self.npoints)
else:
point_set = point_set[0:self.npoints, :]
self.list_of_points[index] = point_set
self.list_of_labels[index] = cls
with open(self.save_path, 'wb') as f:
pickle.dump([self.list_of_points, self.list_of_labels], f)
else:
print('Load processed data from %s...' % self.save_path)
with open(self.save_path, 'rb') as f:
self.list_of_points, self.list_of_labels = pickle.load(f)
def __len__(self):
return len(self.datapath)
def _get_item(self, index):
if self.process_data:
point_set, label = self.list_of_points[index], self.list_of_labels[index]
else:
fn = self.datapath[index]
cls = self.classes[self.datapath[index][0]]
label = np.array([cls]).astype(np.int32)
point_set = np.loadtxt(fn[1], delimiter=',').astype(np.float32)
# 数据集采样npoints个点送入网络
if self.uniform:
point_set = farthest_point_sample(point_set, self.npoints)
else:
point_set = point_set[0:self.npoints, :]
point_set[:, 0:3] = pc_normalize(point_set[:, 0:3])
if not self.use_normals:
point_set = point_set[:, 0:3]
return point_set, label[0]
def __getitem__(self, index):
return self._get_item(index)
if __name__ == '__main__':
import torch
data = ModelNetDataLoader('/data/modelnet40_normal_resampled/', split='train')
DataLoader = torch.utils.data.DataLoader(data, batch_size=12, shuffle=True)
for point, label in DataLoader:
print(point.shape)
print(label.shape)4.2 ShapeNetDataLoader
ShapeNetDataLoader.py：ShapeNet用于训练部件分割，训练集有14007个点云，测试集有2874个点云。
import os
import json
import warnings
import numpy as np
from torch.utils.data import Dataset
warnings.filterwarnings('ignore')
def pc_normalize(pc):
centroid = np.mean(pc, axis=0)
pc = pc - centroid
m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
pc = pc / m
return pc
class PartNormalDataset(Dataset):
def __init__(self,root = './data/shapenetcore_partanno_segmentation_benchmark_v0_normal', npoints=2500, split='train', class_choice=None, normal_channel=False):
self.npoints = npoints
self.root = root
self.catfile = os.path.join(self.root, 'synsetoffset2category.txt')
self.cat = {}
self.normal_channel = normal_channel
with open(self.catfile, 'r') as f:
for line in f:
ls = line.strip().split()
self.cat[ls[0]] = ls[1]
self.cat = {k: v for k, v in self.cat.items()}
self.classes_original = dict(zip(self.cat, range(len(self.cat))))
if not class_choice is None:
self.cat = {k:v for k,v in self.cat.items() if k in class_choice}
# print(self.cat)
self.meta = {}
with open(os.path.join(self.root, 'train_test_split', 'shuffled_train_file_list.json'), 'r') as f:
train_ids = set([str(d.split('/')[2]) for d in json.load(f)])
with open(os.path.join(self.root, 'train_test_split', 'shuffled_val_file_list.json'), 'r') as f:
val_ids = set([str(d.split('/')[2]) for d in json.load(f)])
with open(os.path.join(self.root, 'train_test_split', 'shuffled_test_file_list.json'), 'r') as f:
test_ids = set([str(d.split('/')[2]) for d in json.load(f)])
for item in self.cat:
# print('category', item)
self.meta[item] = []
dir_point = os.path.join(self.root, self.cat[item])
fns = sorted(os.listdir(dir_point))
# print(fns[0][0:-4])
if split == 'trainval':
fns = [fn for fn in fns if ((fn[0:-4] in train_ids) or (fn[0:-4] in val_ids))]
elif split == 'train':
fns = [fn for fn in fns if fn[0:-4] in train_ids]
elif split == 'val':
fns = [fn for fn in fns if fn[0:-4] in val_ids]
elif split == 'test':
fns = [fn for fn in fns if fn[0:-4] in test_ids]
else:
print('Unknown split: %s. Exiting..' % (split))
exit(-1)
# print(os.path.basename(fns))
for fn in fns:
token = (os.path.splitext(os.path.basename(fn))[0])
self.meta[item].append(os.path.join(dir_point, token + '.txt'))
self.datapath = []
for item in self.cat:
for fn in self.meta[item]:
self.datapath.append((item, fn))
self.classes = {}
for i in self.cat.keys():
self.classes[i] = self.classes_original[i]
# Mapping from category ('Chair') to a list of int [10,11,12,13] as segmentation labels
self.seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}
# for cat in sorted(self.seg_classes.keys()):
# print(cat, self.seg_classes[cat])
self.cache = {} # from index to (point_set, cls, seg) tuple
self.cache_size = 20000
def __getitem__(self, index):
if index in self.cache:
point_set, cls, seg = self.cache[index]
else:
fn = self.datapath[index]
cat = self.datapath[index][0]
cls = self.classes[cat]
cls = np.array([cls]).astype(np.int32)
data = np.loadtxt(fn[1]).astype(np.float32)
if not self.normal_channel:
point_set = data[:, 0:3]
else:
point_set = data[:, 0:6]
seg = data[:, -1].astype(np.int32)
if len(self.cache) < self.cache_size:
self.cache[index] = (point_set, cls, seg)
point_set[:, 0:3] = pc_normalize(point_set[:, 0:3])
choice = np.random.choice(len(seg), self.npoints, replace=True)
# resample
point_set = point_set[choice, :]
seg = seg[choice]
return point_set, cls, seg
def __len__(self):
return len(self.datapath)4.3 S3DISDataLoader
S3DISDataLoader.py：S3DIS用于训练语义分割，其中分为6个area，每个area内有若干个room（场景）；room被切割成1 * 1平方米的block，每个block采样4096个点。
import os
import numpy as np
from tqdm import tqdm
from torch.utils.data import Dataset
class S3DISDataset(Dataset):
def __init__(self, split='train', data_root='trainval_fullarea', num_point=4096, test_area=5, block_size=1.0, sample_rate=1.0, transform=None):
super().__init__()
self.num_point = num_point
self.block_size = block_size
self.transform = transform
rooms = sorted(os.listdir(data_root))
rooms = [room for room in rooms if 'Area_' in room]
if split == 'train':
rooms_split = [room for room in rooms if not 'Area_{}'.format(test_area) in room]
else:
rooms_split = [room for room in rooms if 'Area_{}'.format(test_area) in room]
self.room_points, self.room_labels = [], []
self.room_coord_min, self.room_coord_max = [], []
num_point_all = []
labelweights = np.zeros(13)
for room_name in tqdm(rooms_split, total=len(rooms_split)):
room_path = os.path.join(data_root, room_name)
room_data = np.load(room_path) # xyzrgbl, N*7
points, labels = room_data[:, 0:6], room_data[:, 6] # xyzrgb, N*6; l, N
tmp, _ = np.histogram(labels, range(14))
labelweights += tmp
coord_min, coord_max = np.amin(points, axis=0)[:3], np.amax(points, axis=0)[:3]
self.room_points.append(points), self.room_labels.append(labels)
self.room_coord_min.append(coord_min), self.room_coord_max.append(coord_max)
num_point_all.append(labels.size)
labelweights = labelweights.astype(np.float32)
labelweights = labelweights / np.sum(labelweights)
self.labelweights = np.power(np.amax(labelweights) / labelweights, 1 / 3.0)
print(self.labelweights)
sample_prob = num_point_all / np.sum(num_point_all)
num_iter = int(np.sum(num_point_all) * sample_rate / num_point)
room_idxs = []
for index in range(len(rooms_split)):
room_idxs.extend([index] * int(round(sample_prob[index] * num_iter)))
self.room_idxs = np.array(room_idxs)
print("Totally {} samples in {} set.".format(len(self.room_idxs), split))
def __getitem__(self, idx):
room_idx = self.room_idxs[idx]
points = self.room_points[room_idx] # N * 6
labels = self.room_labels[room_idx] # N
N_points = points.shape[0]
while (True):
center = points[np.random.choice(N_points)][:3]
block_min = center - [self.block_size / 2.0, self.block_size / 2.0, 0]
block_max = center + [self.block_size / 2.0, self.block_size / 2.0, 0]
point_idxs = np.where((points[:, 0] >= block_min[0]) & (points[:, 0] <= block_max[0]) & (points[:, 1] >= block_min[1]) & (points[:, 1] <= block_max[1]))[0]
if point_idxs.size > 1024:
break
if point_idxs.size >= self.num_point:
selected_point_idxs = np.random.choice(point_idxs, self.num_point, replace=False)
else:
selected_point_idxs = np.random.choice(point_idxs, self.num_point, replace=True)
# normalize
selected_points = points[selected_point_idxs, :] # num_point * 6
current_points = np.zeros((self.num_point, 9)) # num_point * 9
current_points[:, 6] = selected_points[:, 0] / self.room_coord_max[room_idx][0]
current_points[:, 7] = selected_points[:, 1] / self.room_coord_max[room_idx][1]
current_points[:, 8] = selected_points[:, 2] / self.room_coord_max[room_idx][2]
selected_points[:, 0] = selected_points[:, 0] - center[0]
selected_points[:, 1] = selected_points[:, 1] - center[1]
selected_points[:, 3:6] /= 255.0
current_points[:, 0:6] = selected_points
current_labels = labels[selected_point_idxs]
if self.transform is not None:
current_points, current_labels = self.transform(current_points, current_labels)
return current_points, current_labels
def __len__(self):
return len(self.room_idxs)
class ScannetDatasetWholeScene():
# prepare to give prediction on each points
def __init__(self, root, block_points=4096, split='test', test_area=5, stride=0.5, block_size=1.0, padding=0.001):
self.block_points = block_points
self.block_size = block_size
self.padding = padding
self.root = root
self.split = split
self.stride = stride
self.scene_points_num = []
assert split in ['train', 'test']
if self.split == 'train':
self.file_list = [d for d in os.listdir(root) if d.find('Area_%d' % test_area) is -1]
else:
self.file_list = [d for d in os.listdir(root) if d.find('Area_%d' % test_area) is not -1]
self.scene_points_list = []
self.semantic_labels_list = []
self.room_coord_min, self.room_coord_max = [], []
for file in self.file_list:
data = np.load(root + file)
points = data[:, :3]
self.scene_points_list.append(data[:, :6])
self.semantic_labels_list.append(data[:, 6])
coord_min, coord_max = np.amin(points, axis=0)[:3], np.amax(points, axis=0)[:3]
self.room_coord_min.append(coord_min), self.room_coord_max.append(coord_max)
assert len(self.scene_points_list) == len(self.semantic_labels_list)
labelweights = np.zeros(13)
for seg in self.semantic_labels_list:
tmp, _ = np.histogram(seg, range(14))
self.scene_points_num.append(seg.shape[0])
labelweights += tmp
labelweights = labelweights.astype(np.float32)
labelweights = labelweights / np.sum(labelweights)
self.labelweights = np.power(np.amax(labelweights) / labelweights, 1 / 3.0)
def __getitem__(self, index):
point_set_ini = self.scene_points_list[index]
points = point_set_ini[:,:6]
labels = self.semantic_labels_list[index]
coord_min, coord_max = np.amin(points, axis=0)[:3], np.amax(points, axis=0)[:3]
grid_x = int(np.ceil(float(coord_max[0] - coord_min[0] - self.block_size) / self.stride) + 1)
grid_y = int(np.ceil(float(coord_max[1] - coord_min[1] - self.block_size) / self.stride) + 1)
data_room, label_room, sample_weight, index_room = np.array([]), np.array([]), np.array([]), np.array([])
for index_y in range(0, grid_y):
for index_x in range(0, grid_x):
s_x = coord_min[0] + index_x * self.stride
e_x = min(s_x + self.block_size, coord_max[0])
s_x = e_x - self.block_size
s_y = coord_min[1] + index_y * self.stride
e_y = min(s_y + self.block_size, coord_max[1])
s_y = e_y - self.block_size
point_idxs = np.where(
(points[:, 0] >= s_x - self.padding) & (points[:, 0] <= e_x + self.padding) & (points[:, 1] >= s_y - self.padding) & (
points[:, 1] <= e_y + self.padding))[0]
if point_idxs.size == 0:
continue
num_batch = int(np.ceil(point_idxs.size / self.block_points))
point_size = int(num_batch * self.block_points)
replace = False if (point_size - point_idxs.size <= point_idxs.size) else True
point_idxs_repeat = np.random.choice(point_idxs, point_size - point_idxs.size, replace=replace)
point_idxs = np.concatenate((point_idxs, point_idxs_repeat))
np.random.shuffle(point_idxs)
data_batch = points[point_idxs, :]
normlized_xyz = np.zeros((point_size, 3))
normlized_xyz[:, 0] = data_batch[:, 0] / coord_max[0]
normlized_xyz[:, 1] = data_batch[:, 1] / coord_max[1]
normlized_xyz[:, 2] = data_batch[:, 2] / coord_max[2]
data_batch[:, 0] = data_batch[:, 0] - (s_x + self.block_size / 2.0)
data_batch[:, 1] = data_batch[:, 1] - (s_y + self.block_size / 2.0)
data_batch[:, 3:6] /= 255.0
data_batch = np.concatenate((data_batch, normlized_xyz), axis=1)
label_batch = labels[point_idxs].astype(int)
batch_weight = self.labelweights[label_batch]
data_room = np.vstack([data_room, data_batch]) if data_room.size else data_batch
label_room = np.hstack([label_room, label_batch]) if label_room.size else label_batch
sample_weight = np.hstack([sample_weight, batch_weight]) if label_room.size else batch_weight
index_room = np.hstack([index_room, point_idxs]) if index_room.size else point_idxs
data_room = data_room.reshape((-1, self.block_points, data_room.shape[1]))
label_room = label_room.reshape((-1, self.block_points))
sample_weight = sample_weight.reshape((-1, self.block_points))
index_room = index_room.reshape((-1, self.block_points))
return data_room, label_room, sample_weight, index_room
def __len__(self):
return len(self.scene_points_list)
if __name__ == '__main__':
data_root = '/data/yxu/PointNonLocal/data/stanford_indoor3d/'
num_point, test_area, block_size, sample_rate = 4096, 5, 1.0, 0.01
point_data = S3DISDataset(split='train', data_root=data_root, num_point=num_point, test_area=test_area, block_size=block_size, sample_rate=sample_rate, transform=None)
print('point data size:', point_data.__len__())
print('point data 0 shape:', point_data.__getitem__(0)[0].shape)
print('point label 0 shape:', point_data.__getitem__(0)[1].shape)
import torch, time, random
manual_seed = 123
random.seed(manual_seed)
np.random.seed(manual_seed)
torch.manual_seed(manual_seed)
torch.cuda.manual_seed_all(manual_seed)
def worker_init_fn(worker_id):
random.seed(manual_seed + worker_id)
train_loader = torch.utils.data.DataLoader(point_data, batch_size=16, shuffle=True, num_workers=16, pin_memory=True, worker_init_fn=worker_init_fn)
for idx in range(4):
end = time.time()
for i, (input, target) in enumerate(train_loader):
print('time: {}/{}--{}'.format(i+1, len(train_loader), time.time() - end))
end = time.time()5 数据预处理
provider.py
import numpy as np
# 归一化batch_data，使用以centroid为中心的块的坐标
def normalize_data(batch_data):
""" Normalize the batch data, use coordinates of the block centered at origin,
Input:
BxNxC array
Output:
BxNxC array
"""
B, N, C = batch_data.shape
normal_data = np.zeros((B, N, C))
for b in range(B):
pc = batch_data[b]
centroid = np.mean(pc, axis=0)
pc = pc - centroid
m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
pc = pc / m
normal_data[b] = pc
return normal_data
def shuffle_data(data, labels):
""" Shuffle data and labels.
Input:
data: B,N,... numpy array
label: B,... numpy array
Return:
shuffled data, label and shuffle indices
"""
idx = np.arange(len(labels)) # 创建0 ~ len(labels)-1的等差数列
np.random.shuffle(idx) # 打乱idx
return data[idx, ...], labels[idx], idx
# 打乱每个点云中的点顺序，用于更改FPS行为，对整个batch使用相同的打乱索引idx
def shuffle_points(batch_data):
""" Shuffle orders of points in each point cloud -- changes FPS behavior.
Use the same shuffling idx for the entire batch.
Input:
BxNxC array
Output:
BxNxC array
"""
idx = np.arange(batch_data.shape[1])
np.random.shuffle(idx)
return batch_data[:,idx,:]
# 随机旋转点云进行数据集增强：每个形状沿向上方向旋转
def rotate_point_cloud(batch_data):
""" Randomly rotate the point clouds to augument the dataset
rotation is per shape based along up direction
Input:
BxNx3 array, original batch of point clouds
Return:
BxNx3 array, rotated batch of point clouds
"""
rotated_data = np.zeros(batch_data.shape, dtype=np.float32) # 根据batch_data的矩阵结构，构建零矩阵
for k in range(batch_data.shape[0]):
# 生成一个0~1之间的随机数，乘以2 * pi，得到一个角度
rotation_angle = np.random.uniform() * 2 * np.pi
# 求此角度的cos和sin
cosval = np.cos(rotation_angle)
sinval = np.sin(rotation_angle)
# 组成一个3 x 3的旋转矩阵
rotation_matrix = np.array([[cosval, 0, sinval],
[0, 1, 0],
[-sinval, 0, cosval]])
# 旋转点云数据：乘上旋转矩阵
shape_pc = batch_data[k, ...]
rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
return rotated_data
# 沿z轴做数据增强
def rotate_point_cloud_z(batch_data):
""" Randomly rotate the point clouds to augument the dataset
rotation is per shape based along up direction
Input:
BxNx3 array, original batch of point clouds
Return:
BxNx3 array, rotated batch of point clouds
"""
rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
for k in range(batch_data.shape[0]):
rotation_angle = np.random.uniform() * 2 * np.pi
cosval = np.cos(rotation_angle)
sinval = np.sin(rotation_angle)
rotation_matrix = np.array([[cosval, sinval, 0],
[-sinval, cosval, 0],
[0, 0, 1]])
shape_pc = batch_data[k, ...]
rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
return rotated_data
# 旋转具有法向量信息的点云做数据增强
def rotate_point_cloud_with_normal(batch_xyz_normal):
''' Randomly rotate XYZ, normal point cloud.
Input:
batch_xyz_normal: B,N,6, first three channels are XYZ, last 3 all normal
Output:
B,N,6, rotated XYZ, normal point cloud
'''
for k in range(batch_xyz_normal.shape[0]):
rotation_angle = np.random.uniform() * 2 * np.pi
cosval = np.cos(rotation_angle)
sinval = np.sin(rotation_angle)
rotation_matrix = np.array([[cosval, 0, sinval],
[0, 1, 0],
[-sinval, 0, cosval]])
shape_pc = batch_xyz_normal[k,:,0:3]
shape_normal = batch_xyz_normal[k,:,3:6]
batch_xyz_normal[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
batch_xyz_normal[k,:,3:6] = np.dot(shape_normal.reshape((-1, 3)), rotation_matrix)
return batch_xyz_normal
# 通过小的旋转随机扰动点云
def rotate_perturbation_point_cloud_with_normal(batch_data, angle_sigma=0.06, angle_clip=0.18):
""" Randomly perturb the point clouds by small rotations
Input:
BxNx6 array, original batch of point clouds and point normals
Return:
BxNx3 array, rotated batch of point clouds
"""
rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
for k in range(batch_data.shape[0]):
angles = np.clip(angle_sigma*np.random.randn(3), -angle_clip, angle_clip)
Rx = np.array([[1,0,0],
[0,np.cos(angles[0]),-np.sin(angles[0])],
[0,np.sin(angles[0]),np.cos(angles[0])]])
Ry = np.array([[np.cos(angles[1]),0,np.sin(angles[1])],
[0,1,0],
[-np.sin(angles[1]),0,np.cos(angles[1])]])
Rz = np.array([[np.cos(angles[2]),-np.sin(angles[2]),0],
[np.sin(angles[2]),np.cos(angles[2]),0],
[0,0,1]])
R = np.dot(Rz, np.dot(Ry,Rx))
shape_pc = batch_data[k,:,0:3]
shape_normal = batch_data[k,:,3:6]
rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), R)
rotated_data[k,:,3:6] = np.dot(shape_normal.reshape((-1, 3)), R)
return rotated_data
# 将点云向上旋转指定角度
def rotate_point_cloud_by_angle(batch_data, rotation_angle):
""" Rotate the point cloud along up direction with certain angle.
Input:
BxNx3 array, original batch of point clouds
Return:
BxNx3 array, rotated batch of point clouds
"""
rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
for k in range(batch_data.shape[0]):
#rotation_angle = np.random.uniform() * 2 * np.pi
cosval = np.cos(rotation_angle)
sinval = np.sin(rotation_angle)
rotation_matrix = np.array([[cosval, 0, sinval],
[0, 1, 0],
[-sinval, 0, cosval]])
shape_pc = batch_data[k,:,0:3]
rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
return rotated_data
# 将具有法向量信息的点云向上方向旋转一定角度
def rotate_point_cloud_by_angle_with_normal(batch_data, rotation_angle):
""" Rotate the point cloud along up direction with certain angle.
Input:
BxNx6 array, original batch of point clouds with normal
scalar, angle of rotation
Return:
BxNx6 array, rotated batch of point clouds iwth normal
"""
rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
for k in range(batch_data.shape[0]):
#rotation_angle = np.random.uniform() * 2 * np.pi
cosval = np.cos(rotation_angle)
sinval = np.sin(rotation_angle)
rotation_matrix = np.array([[cosval, 0, sinval],
[0, 1, 0],
[-sinval, 0, cosval]])
shape_pc = batch_data[k,:,0:3]
shape_normal = batch_data[k,:,3:6]
rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
rotated_data[k,:,3:6] = np.dot(shape_normal.reshape((-1,3)), rotation_matrix)
return rotated_data
# 过小的旋转随机扰动点云
def rotate_perturbation_point_cloud(batch_data, angle_sigma=0.06, angle_clip=0.18):
""" Randomly perturb the point clouds by small rotations
Input:
BxNx3 array, original batch of point clouds
Return:
BxNx3 array, rotated batch of point clouds
"""
rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
for k in range(batch_data.shape[0]):
angles = np.clip(angle_sigma*np.random.randn(3), -angle_clip, angle_clip)
Rx = np.array([[1,0,0],
[0,np.cos(angles[0]),-np.sin(angles[0])],
[0,np.sin(angles[0]),np.cos(angles[0])]])
Ry = np.array([[np.cos(angles[1]),0,np.sin(angles[1])],
[0,1,0],
[-np.sin(angles[1]),0,np.cos(angles[1])]])
Rz = np.array([[np.cos(angles[2]),-np.sin(angles[2]),0],
[np.sin(angles[2]),np.cos(angles[2]),0],
[0,0,1]])
R = np.dot(Rz, np.dot(Ry,Rx))
shape_pc = batch_data[k, ...]
rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), R)
return rotated_data
# 随机抖动点
def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05):
""" Randomly jitter points. jittering is per point.
Input:
BxNx3 array, original batch of point clouds
Return:
BxNx3 array, jittered batch of point clouds
"""
B, N, C = batch_data.shape
assert(clip > 0)
jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1*clip, clip)
jittered_data += batch_data
return jittered_data
# 随机移位点云（针对每个点云）
def shift_point_cloud(batch_data, shift_range=0.1):
""" Randomly shift point cloud. Shift is per point cloud.
Input:
BxNx3 array, original batch of point clouds
Return:
BxNx3 array, shifted batch of point clouds
"""
B, N, C = batch_data.shape
shifts = np.random.uniform(-shift_range, shift_range, (B,3))
for batch_index in range(B):
batch_data[batch_index,:,:] += shifts[batch_index,:]
return batch_data
# 随机缩放点云（针对每个点云）
def random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25):
""" Randomly scale the point cloud. Scale is per point cloud.
Input:
BxNx3 array, original batch of point clouds
Return:
BxNx3 array, scaled batch of point clouds
"""
B, N, C = batch_data.shape
scales = np.random.uniform(scale_low, scale_high, B)
for batch_index in range(B):
batch_data[batch_index,:,:] *= scales[batch_index]
return batch_data
# 随机丢弃点云中的点
def random_point_dropout(batch_pc, max_dropout_ratio=0.875):
''' batch_pc: BxNx3 '''
for b in range(batch_pc.shape[0]):
dropout_ratio = np.random.random()*max_dropout_ratio # 0~0.875
drop_idx = np.where(np.random.random((batch_pc.shape[1]))<=dropout_ratio)[0]
if len(drop_idx)>0:
batch_pc[b,drop_idx,:] = batch_pc[b,0,:] # set to the first point
return batch_pc6 训练与测试
6.1 物体分类
train_classification.py
import os
import sys
import torch
import numpy as np
import datetime
import logging # 日志处理
import provider
import importlib
import shutil
import argparse # 命令行解析
from pathlib import Path
from tqdm import tqdm # 进度条
from data_utils.ModelNetDataLoader import ModelNetDataLoader
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))
def parse_args():
'''PARAMETERS'''
# 建立参数解析对象
parser = argparse.ArgumentParser('training')
# 给实例添加属性
parser.add_argument('--use_cpu', action='store_true', default=False, help='use cpu mode')
parser.add_argument('--gpu', type=str, default='0', help='specify gpu device')
parser.add_argument('--batch_size', type=int, default=24, help='batch size in training')
parser.add_argument('--model', default='pointnet_cls', help='model name [default: pointnet_cls]')
parser.add_argument('--num_category', default=40, type=int, choices=[10, 40], help='training on ModelNet10/40')
parser.add_argument('--epoch', default=200, type=int, help='number of epoch in training')
parser.add_argument('--learning_rate', default=0.001, type=float, help='learning rate in training')
parser.add_argument('--num_point', type=int, default=1024, help='Point Number')
parser.add_argument('--optimizer', type=str, default='Adam', help='optimizer for training')
parser.add_argument('--log_dir', type=str, default=None, help='experiment root')
parser.add_argument('--decay_rate', type=float, default=1e-4, help='decay rate')
parser.add_argument('--use_normals', action='store_true', default=False, help='use normals')
parser.add_argument('--process_data', action='store_true', default=False, help='save data offline')
parser.add_argument('--use_uniform_sample', action='store_true', default=False, help='use uniform sampiling')
# 获取对象解析的参数
return parser.parse_args()
def inplace_relu(m):
classname = m.__class__.__name__
if classname.find('ReLU') != -1:
m.inplace=True
def test(model, loader, num_class=40):
mean_correct = []
class_acc = np.zeros((num_class, 3))
classifier = model.eval()
for j, (points, target) in tqdm(enumerate(loader), total=len(loader)):
if not args.use_cpu:
points, target = points.cuda(), target.cuda()
points = points.transpose(2, 1)
pred, _ = classifier(points)
pred_choice = pred.data.max(1)[1]
for cat in np.unique(target.cpu()):
classacc = pred_choice[target == cat].eq(target[target == cat].long().data).cpu().sum()
class_acc[cat, 0] += classacc.item() / float(points[target == cat].size()[0])
class_acc[cat, 1] += 1
correct = pred_choice.eq(target.long().data).cpu().sum()
mean_correct.append(correct.item() / float(points.size()[0]))
class_acc[:, 2] = class_acc[:, 0] / class_acc[:, 1]
class_acc = np.mean(class_acc[:, 2])
instance_acc = np.mean(mean_correct)
return instance_acc, class_acc
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
exp_dir = Path('./log/')
exp_dir.mkdir(exist_ok=True)
exp_dir = exp_dir.joinpath('classification')
exp_dir.mkdir(exist_ok=True)
if args.log_dir is None:
exp_dir = exp_dir.joinpath(timestr)
else:
exp_dir = exp_dir.joinpath(args.log_dir)
exp_dir.mkdir(exist_ok=True)
checkpoints_dir = exp_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = exp_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
'''DATA LOADING'''
log_string('Load dataset ...')
data_path = 'data/modelnet40_normal_resampled/'
train_dataset = ModelNetDataLoader(root=data_path, args=args, split='train', process_data=args.process_data)
test_dataset = ModelNetDataLoader(root=data_path, args=args, split='test', process_data=args.process_data)
trainDataLoader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=10, drop_last=True)
testDataLoader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=10)
'''MODEL LOADING'''
num_class = args.num_category
model = importlib.import_module(args.model)
shutil.copy('./models/%s.py' % args.model, str(exp_dir))
shutil.copy('models/pointnet2_utils.py', str(exp_dir))
shutil.copy('./train_classification.py', str(exp_dir))
classifier = model.get_model(num_class, normal_channel=args.use_normals)
criterion = model.get_loss()
classifier.apply(inplace_relu)
if not args.use_cpu:
classifier = classifier.cuda()
criterion = criterion.cuda()
try:
checkpoint = torch.load(str(exp_dir) + '/checkpoints/best_model.pth')
start_epoch = checkpoint['epoch']
classifier.load_state_dict(checkpoint['model_state_dict'])
log_string('Use pretrain model')
except:
log_string('No existing model, starting training from scratch...')
start_epoch = 0
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(
classifier.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate
)
else:
optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7)
global_epoch = 0
global_step = 0
best_instance_acc = 0.0
best_class_acc = 0.0
'''TRANING'''
logger.info('Start training...')
for epoch in range(start_epoch, args.epoch):
log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch))
mean_correct = []
classifier = classifier.train()
scheduler.step()
for batch_id, (points, target) in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9):
optimizer.zero_grad()
points = points.data.numpy()
points = provider.random_point_dropout(points)
points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3])
points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3])
points = torch.Tensor(points)
points = points.transpose(2, 1)
if not args.use_cpu:
points, target = points.cuda(), target.cuda()
pred, trans_feat = classifier(points)
loss = criterion(pred, target.long(), trans_feat)
pred_choice = pred.data.max(1)[1]
correct = pred_choice.eq(target.long().data).cpu().sum()
mean_correct.append(correct.item() / float(points.size()[0]))
loss.backward()
optimizer.step()
global_step += 1
train_instance_acc = np.mean(mean_correct)
log_string('Train Instance Accuracy: %f' % train_instance_acc)
with torch.no_grad():
instance_acc, class_acc = test(classifier.eval(), testDataLoader, num_class=num_class)
if (instance_acc >= best_instance_acc):
best_instance_acc = instance_acc
best_epoch = epoch + 1
if (class_acc >= best_class_acc):
best_class_acc = class_acc
log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc))
log_string('Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc))
if (instance_acc >= best_instance_acc):
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/best_model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': best_epoch,
'instance_acc': instance_acc,
'class_acc': class_acc,
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
global_epoch += 1
logger.info('End of training...')
if __name__ == '__main__':
args = parse_args()
main(args)test_classification.py
from data_utils.ModelNetDataLoader import ModelNetDataLoader
import argparse # 命令行解析
import numpy as np
import os
import torch
import logging # 日志处理
from tqdm import tqdm # 进度条
import sys
import importlib
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))
def parse_args():
'''PARAMETERS'''
# 建立参数解析对象
parser = argparse.ArgumentParser('Testing')
# 给实例添加属性
parser.add_argument('--use_cpu', action='store_true', default=False, help='use cpu mode')
parser.add_argument('--gpu', type=str, default='0', help='specify gpu device')
parser.add_argument('--batch_size', type=int, default=24, help='batch size in training')
parser.add_argument('--num_category', default=40, type=int, choices=[10, 40], help='training on ModelNet10/40')
parser.add_argument('--num_point', type=int, default=1024, help='Point Number')
parser.add_argument('--log_dir', type=str, required=True, help='Experiment root')
parser.add_argument('--use_normals', action='store_true', default=False, help='use normals')
parser.add_argument('--use_uniform_sample', action='store_true', default=False, help='use uniform sampiling')
parser.add_argument('--num_votes', type=int, default=3, help='Aggregate classification scores with voting')
# 获取对象解析的参数
return parser.parse_args()
def test(model, loader, num_class=40, vote_num=1):
mean_correct = []
classifier = model.eval()
class_acc = np.zeros((num_class, 3)) # (40, 3)
for j, (points, target) in tqdm(enumerate(loader), total=len(loader)):
if not args.use_cpu:
points, target = points.cuda(), target.cuda()
points = points.transpose(2, 1)
if args.use_cpu:
vote_pool = torch.zeros(target.size()[0], num_class)
else:
vote_pool = torch.zeros(target.size()[0], num_class).cuda()
for _ in range(vote_num): # default: 3
pred, _ = classifier(points)
vote_pool += pred
pred = vote_pool / vote_num # 求平均
pred_choice = pred.data.max(1)[1] # pred_choice.shape: (batch_size,)（batch_size默认24）
for cat in np.unique(target.cpu()):
classacc = pred_choice[target == cat].eq(target[target == cat].long().data).cpu().sum()
class_acc[cat, 0] += classacc.item() / float(points[target == cat].size()[0])
class_acc[cat, 1] += 1
correct = pred_choice.eq(target.long().data).cpu().sum()
mean_correct.append(correct.item() / float(points.size()[0]))
# 求类别的accuracy
class_acc[:, 2] = class_acc[:, 0] / class_acc[:, 1]
class_acc = np.mean(class_acc[:, 2])
# 求instance的accuracy
instance_acc = np.mean(mean_correct)
return instance_acc, class_acc
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
experiment_dir = 'log/classification/' + args.log_dir
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir)
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
'''DATA LOADING'''
log_string('Load dataset ...')
data_path = 'data/modelnet40_normal_resampled/'
test_dataset = ModelNetDataLoader(root=data_path, args=args, split='test', process_data=False)
testDataLoader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=10)
'''MODEL LOADING'''
num_class = args.num_category
model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0]
model = importlib.import_module(model_name)
classifier = model.get_model(num_class, normal_channel=args.use_normals)
if not args.use_cpu:
classifier = classifier.cuda()
checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth', weights_only=False) # TEMP forced
classifier.load_state_dict(checkpoint['model_state_dict'])
with torch.no_grad():
instance_acc, class_acc = test(classifier.eval(), testDataLoader, vote_num=args.num_votes, num_class=num_class)
log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc))
if __name__ == '__main__':
args = parse_args()
main(args)6.2 部件分割
train_partseg.py
import argparse
import os
import torch
import datetime
import logging
import sys
import importlib
import shutil
import provider
import numpy as np
from pathlib import Path
from tqdm import tqdm
from data_utils.ShapeNetDataLoader import PartNormalDataset
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))
seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37],
'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49],
'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}
seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table}
for cat in seg_classes.keys():
for label in seg_classes[cat]:
seg_label_to_cat[label] = cat
def inplace_relu(m):
classname = m.__class__.__name__
if classname.find('ReLU') != -1:
m.inplace=True
def to_categorical(y, num_classes):
""" 1-hot encodes a tensor """
new_y = torch.eye(num_classes)[y.cpu().data.numpy(),]
if (y.is_cuda):
return new_y.cuda()
return new_y
def parse_args():
parser = argparse.ArgumentParser('Model')
parser.add_argument('--model', type=str, default='pointnet_part_seg', help='model name')
parser.add_argument('--batch_size', type=int, default=16, help='batch Size during training')
parser.add_argument('--epoch', default=251, type=int, help='epoch to run')
parser.add_argument('--learning_rate', default=0.001, type=float, help='initial learning rate')
parser.add_argument('--gpu', type=str, default='0', help='specify GPU devices')
parser.add_argument('--optimizer', type=str, default='Adam', help='Adam or SGD')
parser.add_argument('--log_dir', type=str, default=None, help='log path')
parser.add_argument('--decay_rate', type=float, default=1e-4, help='weight decay')
parser.add_argument('--npoint', type=int, default=2048, help='point Number')
parser.add_argument('--normal', action='store_true', default=False, help='use normals')
parser.add_argument('--step_size', type=int, default=20, help='decay step for lr decay')
parser.add_argument('--lr_decay', type=float, default=0.5, help='decay rate for lr decay')
return parser.parse_args()
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
exp_dir = Path('./log/')
exp_dir.mkdir(exist_ok=True)
exp_dir = exp_dir.joinpath('part_seg')
exp_dir.mkdir(exist_ok=True)
if args.log_dir is None:
exp_dir = exp_dir.joinpath(timestr)
else:
exp_dir = exp_dir.joinpath(args.log_dir)
exp_dir.mkdir(exist_ok=True)
checkpoints_dir = exp_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = exp_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
root = 'data/shapenetcore_partanno_segmentation_benchmark_v0_normal/'
TRAIN_DATASET = PartNormalDataset(root=root, npoints=args.npoint, split='trainval', normal_channel=args.normal)
trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=10, drop_last=True)
TEST_DATASET = PartNormalDataset(root=root, npoints=args.npoint, split='test', normal_channel=args.normal)
testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=10)
log_string("The number of training data is: %d" % len(TRAIN_DATASET))
log_string("The number of test data is: %d" % len(TEST_DATASET))
num_classes = 16
num_part = 50
'''MODEL LOADING'''
MODEL = importlib.import_module(args.model)
shutil.copy('models/%s.py' % args.model, str(exp_dir))
shutil.copy('models/pointnet2_utils.py', str(exp_dir))
classifier = MODEL.get_model(num_part, normal_channel=args.normal).cuda()
criterion = MODEL.get_loss().cuda()
classifier.apply(inplace_relu)
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
elif classname.find('Linear') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
try:
checkpoint = torch.load(str(exp_dir) + '/checkpoints/best_model.pth')
start_epoch = checkpoint['epoch']
classifier.load_state_dict(checkpoint['model_state_dict'])
log_string('Use pretrain model')
except:
log_string('No existing model, starting training from scratch...')
start_epoch = 0
classifier = classifier.apply(weights_init)
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(
classifier.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate
)
else:
optimizer = torch.optim.SGD(classifier.parameters(), lr=args.learning_rate, momentum=0.9)
def bn_momentum_adjust(m, momentum):
if isinstance(m, torch.nn.BatchNorm2d) or isinstance(m, torch.nn.BatchNorm1d):
m.momentum = momentum
LEARNING_RATE_CLIP = 1e-5
MOMENTUM_ORIGINAL = 0.1
MOMENTUM_DECCAY = 0.5
MOMENTUM_DECCAY_STEP = args.step_size
best_acc = 0
global_epoch = 0
best_class_avg_iou = 0
best_inctance_avg_iou = 0
for epoch in range(start_epoch, args.epoch):
mean_correct = []
log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch))
'''Adjust learning rate and BN momentum'''
lr = max(args.learning_rate * (args.lr_decay ** (epoch // args.step_size)), LEARNING_RATE_CLIP)
log_string('Learning rate:%f' % lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY ** (epoch // MOMENTUM_DECCAY_STEP))
if momentum < 0.01:
momentum = 0.01
print('BN momentum updated to: %f' % momentum)
classifier = classifier.apply(lambda x: bn_momentum_adjust(x, momentum))
classifier = classifier.train()
'''learning one epoch'''
for i, (points, label, target) in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9):
optimizer.zero_grad()
points = points.data.numpy()
points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3])
points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3])
points = torch.Tensor(points)
points, label, target = points.float().cuda(), label.long().cuda(), target.long().cuda()
points = points.transpose(2, 1)
seg_pred, trans_feat = classifier(points, to_categorical(label, num_classes))
seg_pred = seg_pred.contiguous().view(-1, num_part)
target = target.view(-1, 1)[:, 0]
pred_choice = seg_pred.data.max(1)[1]
correct = pred_choice.eq(target.data).cpu().sum()
mean_correct.append(correct.item() / (args.batch_size * args.npoint))
loss = criterion(seg_pred, target, trans_feat)
loss.backward()
optimizer.step()
train_instance_acc = np.mean(mean_correct)
log_string('Train accuracy is: %.5f' % train_instance_acc)
with torch.no_grad():
test_metrics = {}
total_correct = 0
total_seen = 0
total_seen_class = [0 for _ in range(num_part)]
total_correct_class = [0 for _ in range(num_part)]
shape_ious = {cat: [] for cat in seg_classes.keys()}
seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table}
for cat in seg_classes.keys():
for label in seg_classes[cat]:
seg_label_to_cat[label] = cat
classifier = classifier.eval()
for batch_id, (points, label, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader), smoothing=0.9):
cur_batch_size, NUM_POINT, _ = points.size()
points, label, target = points.float().cuda(), label.long().cuda(), target.long().cuda()
points = points.transpose(2, 1)
seg_pred, _ = classifier(points, to_categorical(label, num_classes))
cur_pred_val = seg_pred.cpu().data.numpy()
cur_pred_val_logits = cur_pred_val
cur_pred_val = np.zeros((cur_batch_size, NUM_POINT)).astype(np.int32)
target = target.cpu().data.numpy()
for i in range(cur_batch_size):
cat = seg_label_to_cat[target[i, 0]]
logits = cur_pred_val_logits[i, :, :]
cur_pred_val[i, :] = np.argmax(logits[:, seg_classes[cat]], 1) + seg_classes[cat][0]
correct = np.sum(cur_pred_val == target)
total_correct += correct
total_seen += (cur_batch_size * NUM_POINT)
for l in range(num_part):
total_seen_class[l] += np.sum(target == l)
total_correct_class[l] += (np.sum((cur_pred_val == l) & (target == l)))
for i in range(cur_batch_size):
segp = cur_pred_val[i, :]
segl = target[i, :]
cat = seg_label_to_cat[segl[0]]
part_ious = [0.0 for _ in range(len(seg_classes[cat]))]
for l in seg_classes[cat]:
if (np.sum(segl == l) == 0) and (
np.sum(segp == l) == 0): # part is not present, no prediction as well
part_ious[l - seg_classes[cat][0]] = 1.0
else:
part_ious[l - seg_classes[cat][0]] = np.sum((segl == l) & (segp == l)) / float(
np.sum((segl == l) | (segp == l)))
shape_ious[cat].append(np.mean(part_ious))
all_shape_ious = []
for cat in shape_ious.keys():
for iou in shape_ious[cat]:
all_shape_ious.append(iou)
shape_ious[cat] = np.mean(shape_ious[cat])
mean_shape_ious = np.mean(list(shape_ious.values()))
test_metrics['accuracy'] = total_correct / float(total_seen)
test_metrics['class_avg_accuracy'] = np.mean(
np.array(total_correct_class) / np.array(total_seen_class, dtype=np.float))
for cat in sorted(shape_ious.keys()):
log_string('eval mIoU of %s %f' % (cat + ' ' * (14 - len(cat)), shape_ious[cat]))
test_metrics['class_avg_iou'] = mean_shape_ious
test_metrics['inctance_avg_iou'] = np.mean(all_shape_ious)
log_string('Epoch %d test Accuracy: %f Class avg mIOU: %f Inctance avg mIOU: %f' % (
epoch + 1, test_metrics['accuracy'], test_metrics['class_avg_iou'], test_metrics['inctance_avg_iou']))
if (test_metrics['inctance_avg_iou'] >= best_inctance_avg_iou):
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/best_model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'train_acc': train_instance_acc,
'test_acc': test_metrics['accuracy'],
'class_avg_iou': test_metrics['class_avg_iou'],
'inctance_avg_iou': test_metrics['inctance_avg_iou'],
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
if test_metrics['accuracy'] > best_acc:
best_acc = test_metrics['accuracy']
if test_metrics['class_avg_iou'] > best_class_avg_iou:
best_class_avg_iou = test_metrics['class_avg_iou']
if test_metrics['inctance_avg_iou'] > best_inctance_avg_iou:
best_inctance_avg_iou = test_metrics['inctance_avg_iou']
log_string('Best accuracy is: %.5f' % best_acc)
log_string('Best class avg mIOU is: %.5f' % best_class_avg_iou)
log_string('Best inctance avg mIOU is: %.5f' % best_inctance_avg_iou)
global_epoch += 1
if __name__ == '__main__':
args = parse_args()
main(args)test_partseg.py
import argparse
import os
from data_utils.ShapeNetDataLoader import PartNormalDataset
import torch
import logging
import sys
import importlib
from tqdm import tqdm
import numpy as np
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))
seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37],
'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49],
'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}
seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table}
for cat in seg_classes.keys():
for label in seg_classes[cat]:
seg_label_to_cat[label] = cat
def to_categorical(y, num_classes):
""" 1-hot encodes a tensor """
new_y = torch.eye(num_classes)[y.cpu().data.numpy(),]
if (y.is_cuda):
return new_y.cuda()
return new_y
def parse_args():
'''PARAMETERS'''
parser = argparse.ArgumentParser('PointNet')
parser.add_argument('--batch_size', type=int, default=24, help='batch size in testing')
parser.add_argument('--gpu', type=str, default='0', help='specify gpu device')
parser.add_argument('--num_point', type=int, default=2048, help='point Number')
parser.add_argument('--log_dir', type=str, required=True, help='experiment root')
parser.add_argument('--normal', action='store_true', default=False, help='use normals')
parser.add_argument('--num_votes', type=int, default=3, help='aggregate segmentation scores with voting')
return parser.parse_args()
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
experiment_dir = 'log/part_seg/' + args.log_dir
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir)
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
root = 'data/shapenetcore_partanno_segmentation_benchmark_v0_normal/'
TEST_DATASET = PartNormalDataset(root=root, npoints=args.num_point, split='test', normal_channel=args.normal)
testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4)
log_string("The number of test data is: %d" % len(TEST_DATASET))
num_classes = 16
num_part = 50
'''MODEL LOADING'''
model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0]
MODEL = importlib.import_module(model_name)
classifier = MODEL.get_model(num_part, normal_channel=args.normal).cuda()
checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth')
classifier.load_state_dict(checkpoint['model_state_dict'])
with torch.no_grad():
test_metrics = {}
total_correct = 0
total_seen = 0
total_seen_class = [0 for _ in range(num_part)]
total_correct_class = [0 for _ in range(num_part)]
shape_ious = {cat: [] for cat in seg_classes.keys()}
seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table}
for cat in seg_classes.keys():
for label in seg_classes[cat]:
seg_label_to_cat[label] = cat
classifier = classifier.eval()
for batch_id, (points, label, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader),
smoothing=0.9):
batchsize, num_point, _ = points.size()
cur_batch_size, NUM_POINT, _ = points.size()
points, label, target = points.float().cuda(), label.long().cuda(), target.long().cuda()
points = points.transpose(2, 1)
vote_pool = torch.zeros(target.size()[0], target.size()[1], num_part).cuda()
for _ in range(args.num_votes):
seg_pred, _ = classifier(points, to_categorical(label, num_classes))
vote_pool += seg_pred
seg_pred = vote_pool / args.num_votes
cur_pred_val = seg_pred.cpu().data.numpy()
cur_pred_val_logits = cur_pred_val
cur_pred_val = np.zeros((cur_batch_size, NUM_POINT)).astype(np.int32)
target = target.cpu().data.numpy()
for i in range(cur_batch_size):
cat = seg_label_to_cat[target[i, 0]]
logits = cur_pred_val_logits[i, :, :]
cur_pred_val[i, :] = np.argmax(logits[:, seg_classes[cat]], 1) + seg_classes[cat][0]
correct = np.sum(cur_pred_val == target)
total_correct += correct
total_seen += (cur_batch_size * NUM_POINT)
for l in range(num_part):
total_seen_class[l] += np.sum(target == l)
total_correct_class[l] += (np.sum((cur_pred_val == l) & (target == l)))
for i in range(cur_batch_size):
segp = cur_pred_val[i, :]
segl = target[i, :]
cat = seg_label_to_cat[segl[0]]
part_ious = [0.0 for _ in range(len(seg_classes[cat]))]
for l in seg_classes[cat]:
if (np.sum(segl == l) == 0) and (
np.sum(segp == l) == 0): # part is not present, no prediction as well
part_ious[l - seg_classes[cat][0]] = 1.0
else:
part_ious[l - seg_classes[cat][0]] = np.sum((segl == l) & (segp == l)) / float(
np.sum((segl == l) | (segp == l)))
shape_ious[cat].append(np.mean(part_ious))
all_shape_ious = []
for cat in shape_ious.keys():
for iou in shape_ious[cat]:
all_shape_ious.append(iou)
shape_ious[cat] = np.mean(shape_ious[cat])
mean_shape_ious = np.mean(list(shape_ious.values()))
test_metrics['accuracy'] = total_correct / float(total_seen)
test_metrics['class_avg_accuracy'] = np.mean(
np.array(total_correct_class) / np.array(total_seen_class, dtype=np.float))
for cat in sorted(shape_ious.keys()):
log_string('eval mIoU of %s %f' % (cat + ' ' * (14 - len(cat)), shape_ious[cat]))
test_metrics['class_avg_iou'] = mean_shape_ious
test_metrics['inctance_avg_iou'] = np.mean(all_shape_ious)
log_string('Accuracy is: %.5f' % test_metrics['accuracy'])
log_string('Class avg accuracy is: %.5f' % test_metrics['class_avg_accuracy'])
log_string('Class avg mIOU is: %.5f' % test_metrics['class_avg_iou'])
log_string('Inctance avg mIOU is: %.5f' % test_metrics['inctance_avg_iou'])
if __name__ == '__main__':
args = parse_args()
main(args)6.3 场景分割
train_semseg.py
import argparse
import os
from data_utils.S3DISDataLoader import S3DISDataset
import torch
import datetime
import logging
from pathlib import Path
import sys
import importlib
import shutil
from tqdm import tqdm
import provider
import numpy as np
import time
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))
classes = ['ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door', 'table', 'chair', 'sofa', 'bookcase',
'board', 'clutter']
class2label = {cls: i for i, cls in enumerate(classes)}
seg_classes = class2label
seg_label_to_cat = {}
for i, cat in enumerate(seg_classes.keys()):
seg_label_to_cat[i] = cat
def inplace_relu(m):
classname = m.__class__.__name__
if classname.find('ReLU') != -1:
m.inplace=True
def parse_args():
parser = argparse.ArgumentParser('Model')
parser.add_argument('--model', type=str, default='pointnet_sem_seg', help='model name [default: pointnet_sem_seg]')
parser.add_argument('--batch_size', type=int, default=16, help='Batch Size during training [default: 16]')
parser.add_argument('--epoch', default=32, type=int, help='Epoch to run [default: 32]')
parser.add_argument('--learning_rate', default=0.001, type=float, help='Initial learning rate [default: 0.001]')
parser.add_argument('--gpu', type=str, default='0', help='GPU to use [default: GPU 0]')
parser.add_argument('--optimizer', type=str, default='Adam', help='Adam or SGD [default: Adam]')
parser.add_argument('--log_dir', type=str, default=None, help='Log path [default: None]')
parser.add_argument('--decay_rate', type=float, default=1e-4, help='weight decay [default: 1e-4]')
parser.add_argument('--npoint', type=int, default=4096, help='Point Number [default: 4096]')
parser.add_argument('--step_size', type=int, default=10, help='Decay step for lr decay [default: every 10 epochs]')
parser.add_argument('--lr_decay', type=float, default=0.7, help='Decay rate for lr decay [default: 0.7]')
parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]')
return parser.parse_args()
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
experiment_dir = Path('./log/')
experiment_dir.mkdir(exist_ok=True)
experiment_dir = experiment_dir.joinpath('sem_seg')
experiment_dir.mkdir(exist_ok=True)
if args.log_dir is None:
experiment_dir = experiment_dir.joinpath(timestr)
else:
experiment_dir = experiment_dir.joinpath(args.log_dir)
experiment_dir.mkdir(exist_ok=True)
checkpoints_dir = experiment_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = experiment_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
root = 'data/stanford_indoor3d/'
NUM_CLASSES = 13
NUM_POINT = args.npoint
BATCH_SIZE = args.batch_size
print("start loading training data ...")
TRAIN_DATASET = S3DISDataset(split='train', data_root=root, num_point=NUM_POINT, test_area=args.test_area, block_size=1.0, sample_rate=1.0, transform=None)
print("start loading test data ...")
TEST_DATASET = S3DISDataset(split='test', data_root=root, num_point=NUM_POINT, test_area=args.test_area, block_size=1.0, sample_rate=1.0, transform=None)
trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=BATCH_SIZE, shuffle=True, num_workers=10,
pin_memory=True, drop_last=True,
worker_init_fn=lambda x: np.random.seed(x + int(time.time())))
testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=BATCH_SIZE, shuffle=False, num_workers=10,
pin_memory=True, drop_last=True)
weights = torch.Tensor(TRAIN_DATASET.labelweights).cuda()
log_string("The number of training data is: %d" % len(TRAIN_DATASET))
log_string("The number of test data is: %d" % len(TEST_DATASET))
'''MODEL LOADING'''
MODEL = importlib.import_module(args.model)
shutil.copy('models/%s.py' % args.model, str(experiment_dir))
shutil.copy('models/pointnet2_utils.py', str(experiment_dir))
classifier = MODEL.get_model(NUM_CLASSES).cuda()
criterion = MODEL.get_loss().cuda()
classifier.apply(inplace_relu)
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
elif classname.find('Linear') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
try:
checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth')
start_epoch = checkpoint['epoch']
classifier.load_state_dict(checkpoint['model_state_dict'])
log_string('Use pretrain model')
except:
log_string('No existing model, starting training from scratch...')
start_epoch = 0
classifier = classifier.apply(weights_init)
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(
classifier.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate
)
else:
optimizer = torch.optim.SGD(classifier.parameters(), lr=args.learning_rate, momentum=0.9)
def bn_momentum_adjust(m, momentum):
if isinstance(m, torch.nn.BatchNorm2d) or isinstance(m, torch.nn.BatchNorm1d):
m.momentum = momentum
LEARNING_RATE_CLIP = 1e-5
MOMENTUM_ORIGINAL = 0.1
MOMENTUM_DECCAY = 0.5
MOMENTUM_DECCAY_STEP = args.step_size
global_epoch = 0
best_iou = 0
for epoch in range(start_epoch, args.epoch):
'''Train on chopped scenes'''
log_string('**** Epoch %d (%d/%s) ****' % (global_epoch + 1, epoch + 1, args.epoch))
lr = max(args.learning_rate * (args.lr_decay ** (epoch // args.step_size)), LEARNING_RATE_CLIP)
log_string('Learning rate:%f' % lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY ** (epoch // MOMENTUM_DECCAY_STEP))
if momentum < 0.01:
momentum = 0.01
print('BN momentum updated to: %f' % momentum)
classifier = classifier.apply(lambda x: bn_momentum_adjust(x, momentum))
num_batches = len(trainDataLoader)
total_correct = 0
total_seen = 0
loss_sum = 0
classifier = classifier.train()
for i, (points, target) in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9):
optimizer.zero_grad()
points = points.data.numpy()
points[:, :, :3] = provider.rotate_point_cloud_z(points[:, :, :3])
points = torch.Tensor(points)
points, target = points.float().cuda(), target.long().cuda()
points = points.transpose(2, 1)
seg_pred, trans_feat = classifier(points)
seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)
batch_label = target.view(-1, 1)[:, 0].cpu().data.numpy()
target = target.view(-1, 1)[:, 0]
loss = criterion(seg_pred, target, trans_feat, weights)
loss.backward()
optimizer.step()
pred_choice = seg_pred.cpu().data.max(1)[1].numpy()
correct = np.sum(pred_choice == batch_label)
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
loss_sum += loss
log_string('Training mean loss: %f' % (loss_sum / num_batches))
log_string('Training accuracy: %f' % (total_correct / float(total_seen)))
if epoch % 5 == 0:
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
'''Evaluate on chopped scenes'''
with torch.no_grad():
num_batches = len(testDataLoader)
total_correct = 0
total_seen = 0
loss_sum = 0
labelweights = np.zeros(NUM_CLASSES)
total_seen_class = [0 for _ in range(NUM_CLASSES)]
total_correct_class = [0 for _ in range(NUM_CLASSES)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
classifier = classifier.eval()
log_string('---- EPOCH %03d EVALUATION ----' % (global_epoch + 1))
for i, (points, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader), smoothing=0.9):
points = points.data.numpy()
points = torch.Tensor(points)
points, target = points.float().cuda(), target.long().cuda()
points = points.transpose(2, 1)
seg_pred, trans_feat = classifier(points)
pred_val = seg_pred.contiguous().cpu().data.numpy()
seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)
batch_label = target.cpu().data.numpy()
target = target.view(-1, 1)[:, 0]
loss = criterion(seg_pred, target, trans_feat, weights)
loss_sum += loss
pred_val = np.argmax(pred_val, 2)
correct = np.sum((pred_val == batch_label))
total_correct += correct
total_seen += (BATCH_SIZE * NUM_POINT)
tmp, _ = np.histogram(batch_label, range(NUM_CLASSES + 1))
labelweights += tmp
for l in range(NUM_CLASSES):
total_seen_class[l] += np.sum((batch_label == l))
total_correct_class[l] += np.sum((pred_val == l) & (batch_label == l))
total_iou_deno_class[l] += np.sum(((pred_val == l) | (batch_label == l)))
labelweights = labelweights.astype(np.float32) / np.sum(labelweights.astype(np.float32))
mIoU = np.mean(np.array(total_correct_class) / (np.array(total_iou_deno_class, dtype=np.float) + 1e-6))
log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
log_string('eval point avg class IoU: %f' % (mIoU))
log_string('eval point accuracy: %f' % (total_correct / float(total_seen)))
log_string('eval point avg class acc: %f' % (
np.mean(np.array(total_correct_class) / (np.array(total_seen_class, dtype=np.float) + 1e-6))))
iou_per_class_str = '------- IoU --------\n'
for l in range(NUM_CLASSES):
iou_per_class_str += 'class %s weight: %.3f, IoU: %.3f \n' % (
seg_label_to_cat[l] + ' ' * (14 - len(seg_label_to_cat[l])), labelweights[l - 1],
total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
log_string('Eval mean loss: %f' % (loss_sum / num_batches))
log_string('Eval accuracy: %f' % (total_correct / float(total_seen)))
if mIoU >= best_iou:
best_iou = mIoU
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/best_model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'class_avg_iou': mIoU,
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
log_string('Best mIoU: %f' % best_iou)
global_epoch += 1
if __name__ == '__main__':
args = parse_args()
main(args)test_semseg.py
import argparse
import os
from data_utils.S3DISDataLoader import ScannetDatasetWholeScene
from data_utils.indoor3d_util import g_label2color
import torch
import logging
from pathlib import Path
import sys
import importlib
from tqdm import tqdm
import provider
import numpy as np
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))
classes = ['ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door', 'table', 'chair', 'sofa', 'bookcase',
'board', 'clutter']
class2label = {cls: i for i, cls in enumerate(classes)}
seg_classes = class2label
seg_label_to_cat = {}
for i, cat in enumerate(seg_classes.keys()):
seg_label_to_cat[i] = cat
def parse_args():
'''PARAMETERS'''
parser = argparse.ArgumentParser('Model')
parser.add_argument('--batch_size', type=int, default=32, help='batch size in testing [default: 32]')
parser.add_argument('--gpu', type=str, default='0', help='specify gpu device')
parser.add_argument('--num_point', type=int, default=4096, help='point number [default: 4096]')
parser.add_argument('--log_dir', type=str, required=True, help='experiment root')
parser.add_argument('--visual', action='store_true', default=False, help='visualize result [default: False]')
parser.add_argument('--test_area', type=int, default=5, help='area for testing, option: 1-6 [default: 5]')
parser.add_argument('--num_votes', type=int, default=3, help='aggregate segmentation scores with voting [default: 5]')
return parser.parse_args()
def add_vote(vote_label_pool, point_idx, pred_label, weight):
B = pred_label.shape[0]
N = pred_label.shape[1]
for b in range(B):
for n in range(N):
if weight[b, n] != 0 and not np.isinf(weight[b, n]):
vote_label_pool[int(point_idx[b, n]), int(pred_label[b, n])] += 1
return vote_label_pool
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
experiment_dir = 'log/sem_seg/' + args.log_dir
visual_dir = experiment_dir + '/visual/'
visual_dir = Path(visual_dir)
visual_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir)
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
NUM_CLASSES = 13
BATCH_SIZE = args.batch_size
NUM_POINT = args.num_point
root = 'data/s3dis/stanford_indoor3d/'
TEST_DATASET_WHOLE_SCENE = ScannetDatasetWholeScene(root, split='test', test_area=args.test_area, block_points=NUM_POINT)
log_string("The number of test data is: %d" % len(TEST_DATASET_WHOLE_SCENE))
'''MODEL LOADING'''
model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0]
MODEL = importlib.import_module(model_name)
classifier = MODEL.get_model(NUM_CLASSES).cuda()
checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth')
classifier.load_state_dict(checkpoint['model_state_dict'])
classifier = classifier.eval()
with torch.no_grad():
scene_id = TEST_DATASET_WHOLE_SCENE.file_list
scene_id = [x[:-4] for x in scene_id]
num_batches = len(TEST_DATASET_WHOLE_SCENE)
total_seen_class = [0 for _ in range(NUM_CLASSES)]
total_correct_class = [0 for _ in range(NUM_CLASSES)]
total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
log_string('---- EVALUATION WHOLE SCENE----')
for batch_idx in range(num_batches):
print("Inference [%d/%d] %s ..." % (batch_idx + 1, num_batches, scene_id[batch_idx]))
total_seen_class_tmp = [0 for _ in range(NUM_CLASSES)]
total_correct_class_tmp = [0 for _ in range(NUM_CLASSES)]
total_iou_deno_class_tmp = [0 for _ in range(NUM_CLASSES)]
if args.visual:
fout = open(os.path.join(visual_dir, scene_id[batch_idx] + '_pred.obj'), 'w')
fout_gt = open(os.path.join(visual_dir, scene_id[batch_idx] + '_gt.obj'), 'w')
whole_scene_data = TEST_DATASET_WHOLE_SCENE.scene_points_list[batch_idx]
whole_scene_label = TEST_DATASET_WHOLE_SCENE.semantic_labels_list[batch_idx]
vote_label_pool = np.zeros((whole_scene_label.shape[0], NUM_CLASSES))
for _ in tqdm(range(args.num_votes), total=args.num_votes):
scene_data, scene_label, scene_smpw, scene_point_index = TEST_DATASET_WHOLE_SCENE[batch_idx]
num_blocks = scene_data.shape[0]
s_batch_num = (num_blocks + BATCH_SIZE - 1) // BATCH_SIZE
batch_data = np.zeros((BATCH_SIZE, NUM_POINT, 9))
batch_label = np.zeros((BATCH_SIZE, NUM_POINT))
batch_point_index = np.zeros((BATCH_SIZE, NUM_POINT))
batch_smpw = np.zeros((BATCH_SIZE, NUM_POINT))
for sbatch in range(s_batch_num):
start_idx = sbatch * BATCH_SIZE
end_idx = min((sbatch + 1) * BATCH_SIZE, num_blocks)
real_batch_size = end_idx - start_idx
batch_data[0:real_batch_size, ...] = scene_data[start_idx:end_idx, ...]
batch_label[0:real_batch_size, ...] = scene_label[start_idx:end_idx, ...]
batch_point_index[0:real_batch_size, ...] = scene_point_index[start_idx:end_idx, ...]
batch_smpw[0:real_batch_size, ...] = scene_smpw[start_idx:end_idx, ...]
batch_data[:, :, 3:6] /= 1.0
torch_data = torch.Tensor(batch_data)
torch_data = torch_data.float().cuda()
torch_data = torch_data.transpose(2, 1)
seg_pred, _ = classifier(torch_data)
batch_pred_label = seg_pred.contiguous().cpu().data.max(2)[1].numpy()
vote_label_pool = add_vote(vote_label_pool, batch_point_index[0:real_batch_size, ...],
batch_pred_label[0:real_batch_size, ...],
batch_smpw[0:real_batch_size, ...])
pred_label = np.argmax(vote_label_pool, 1)
for l in range(NUM_CLASSES):
total_seen_class_tmp[l] += np.sum((whole_scene_label == l))
total_correct_class_tmp[l] += np.sum((pred_label == l) & (whole_scene_label == l))
total_iou_deno_class_tmp[l] += np.sum(((pred_label == l) | (whole_scene_label == l)))
total_seen_class[l] += total_seen_class_tmp[l]
total_correct_class[l] += total_correct_class_tmp[l]
total_iou_deno_class[l] += total_iou_deno_class_tmp[l]
iou_map = np.array(total_correct_class_tmp) / (np.array(total_iou_deno_class_tmp, dtype=np.float) + 1e-6)
print(iou_map)
arr = np.array(total_seen_class_tmp)
tmp_iou = np.mean(iou_map[arr != 0])
log_string('Mean IoU of %s: %.4f' % (scene_id[batch_idx], tmp_iou))
print('----------------------------')
filename = os.path.join(visual_dir, scene_id[batch_idx] + '.txt')
with open(filename, 'w') as pl_save:
for i in pred_label:
pl_save.write(str(int(i)) + '\n')
pl_save.close()
for i in range(whole_scene_label.shape[0]):
color = g_label2color[pred_label[i]]
color_gt = g_label2color[whole_scene_label[i]]
if args.visual:
fout.write('v %f %f %f %d %d %d\n' % (
whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color[0], color[1],
color[2]))
fout_gt.write(
'v %f %f %f %d %d %d\n' % (
whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color_gt[0],
color_gt[1], color_gt[2]))
if args.visual:
fout.close()
fout_gt.close()
IoU = np.array(total_correct_class) / (np.array(total_iou_deno_class, dtype=np.float) + 1e-6)
iou_per_class_str = '------- IoU --------\n'
for l in range(NUM_CLASSES):
iou_per_class_str += 'class %s, IoU: %.3f \n' % (
seg_label_to_cat[l] + ' ' * (14 - len(seg_label_to_cat[l])),
total_correct_class[l] / float(total_iou_deno_class[l]))
log_string(iou_per_class_str)
log_string('eval point avg class IoU: %f' % np.mean(IoU))
log_string('eval whole scene point avg class acc: %f' % (
np.mean(np.array(total_correct_class) / (np.array(total_seen_class, dtype=np.float) + 1e-6))))
log_string('eval whole scene point accuracy: %f' % (
np.sum(total_correct_class) / float(np.sum(total_seen_class) + 1e-6)))
print("Done!")
if __name__ == '__main__':
args = parse_args()
main(args)
oko