整体结构
class get_model(nn.Module):
def __init__(self, num_classes, normal_channel=False):
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
该模型包括了三个 Set Abstraction 层和三个 Feature Propagation 层。
在 Set Abstraction 层,输入为点云的坐标和属性信息,经过一系列减少点数、聚合特征信息的操作后,输出为更小数量的点云和更高维度的特征向量。
在 Feature Propagation 层,输入为上一层的点云和特征向量以及下一层的点云,通过一系列将特征向量传播到上一层点云中的操作,输出为更多的点云和更低维度的特征向量。
最后,经过一些全连接层和非线性激活函数,输出为每个点云对应的类别预测结果和最后一层的特征向量。整个模型结构可以看作是将点云逐渐抽象成高维特征向量,再通过反向传播得到各点对应的类别预测结果的一个过程。
PointNetSetAbstraction
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)
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]
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
模块用于执行点云下采样和特征提取。输入包括点云的位置数据
xyz
和属性数据points
,输出包括采样后的点云位置数据new_xyz
和对应的特征向量new_points_concat
。具体来说,这个模块首先根据指定的参数进行采样并分组操作,得到新的点云位置数据
new_xyz
和相应的特征数据new_points
。然后,经过一系列卷积层和批归一化层,将新的特征向量new_points
转换为更高维度的特征向量new_points_concat
。最后,对new_points_concat
执行max pooling操作,得到每个被采样点的全局最大池化特征向量,并返回new_xyz
和new_points
。
def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
npoint
: 点云采样的数量。radius
: 每个点的邻域半径。nsample
: 每个点邻域内的最大点数。in_channel
: 输入点云的通道数。mlp
: 每个局部特征提取层中,多层感知机(MLP)的输出维度序列。group_all
: 是否考虑整个点云作为一个局部区域。如果为 True,则所有点都被视为一个区域,否则每个点都有其自己的局部区域。
xyz = xyz.permute(0, 2, 1)
if points is not None:
points = points.permute(0, 2, 1)
将输入点云中的维度顺序从
[B, C, N]
调整为[B, N, C]
。
假设我们有一个大小为 [2, 3, 4]
的张量 xyz
,其含义是两个 点云样本,每个样本由 4 个点组成,每个点包括 3 个空间坐标。
经过变化,最终如下:
把每个点的自己信息放在一起了(个人理解)
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)
其中 `xyz` 表示点云中的坐标信息,`points` 表示每个点的属性信息。
如果 `self.group_all` 为真,则将所有点分为一个 group,并返回新的 `new_xyz` 和 `new_points`;
否则使用给定的参数进行采样和分组,并返回新的 `new_xyz` 和 `new_points`。具体来说,`npoint` 是采样后得到的点数,`radius` 是用于判断是否属于同一组的半径范围,`nsample` 是每个 group 中包含的最大点数。
# 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]
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)
这段代码是在进行点云数据处理,输入的数据包括采样后的点的位置信息和点的其他属性信息。其中,new_xyz是采样的点的位置信息,new_points是点的其他属性信息,且通过permute函数将new_points变为[B, C+D, nsample,npoint]的形式。
接着,代码使用一个for循环遍历多个卷积层,每个卷积层都包含一个卷积操作和一个批归一化操作。对于每个卷积层,在经过卷积操作后,新的点云信息会先经过批归一化后再进行ReLU激活函数的处理,最终得到新的点云信息。
对张量
new_points
进行沿着第二个维度(即列)取最大值的操作,并返回一个新的张量。对张量
new_xyz
进行将第二个和第三个维度(即列和深度)进行交换的操作,并返回一个新的张量。
sample_and_group_all
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_points
这是一个用于对点云数据进行采样和分组的函数。点云数据通常包括每个点的位置和属性信息。
输入参数:
- xyz:点云中每个点的空间坐标,大小为[B, N, 3],其中B表示批量大小,N表示点的数量,3表示空间坐标轴数。
- points:点云中每个点的属性信息,大小为[B, N, D],其中D表示属性维度数。
返回值:
- new_xyz:采样后的新点的位置数据,大小为[B, 1, 3],即在原来的所有点中随机抽取一个点作为新点的位置。
- new_points:采样并分组后的新点数据,大小为[B, 1, N, 3+D],其中第一维为批量大小,第二维为采样后得到的一个点,第三维为原始点的数量,第四维为每个点的坐标和属性信息。
sample_and_group
def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False):
"""
Input:
npoint:
radius:
nsample:
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_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
new_xyz = index_points(xyz, fps_idx)
idx = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
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
这段代码实现了一个点云采样和分组的函数,主要用于将大规模的点云数据转换成小规模的可用于训练神经网络的数据。
具体来说,函数
sample_and_group
接受以下参数:
npoint
: 采样后的点数。radius
: 在每个采样点周围搜索邻近点的半径。nsample
: 每个采样点要选多少个邻近点。xyz
: 输入点云的位置信息,形状为(B,N,3)
。points
: 输入点云的其它属性信息,形状为(B,N,D)
。returnfps
: 是否返回最远点采样(farthest point sampling)的索引,即fps_idx
。其中,
B
表示 batch size,即输入有多少个点云。N
表示每个点云中有多少个点。C
表示每个点的坐标维度,这里是 3。函数首先使用
farthest_point_sample
对输入点云进行最远点采样,得到npoint
个采样点的索引fps_idx
。然后,针对每个采样点,使用
query_ball_point
函数在以该点为中心、半径为radius
的球体内搜索nsample
个邻近点的索引,并将结果存储在idx
中。最后,使用
index_points
函数按照这些索引从输入点云中提取对应的位置和其它属性信息,并根据采样点将邻近点进行分组,得到形状为(B,npoint,nsample,C)
和(B,npoint,nsample,D)
的新张量grouped_xyz_norm
和new_points
。如果
returnfps
参数设置为 True,则函数还会返回最远点采样的结果fps_idx
和在每个采样点周围搜索到的邻近点的位置信息grouped_xyz
。
index_points
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_points
这是用于从输入点云数据中抽取特定索引的函数。其输入包括两个参数:points表示输入的点云数据,其维度为[B, N, C],即B个点云数据,每个点云数据有N个点,每个点的属性有C个;idx表示需要从输入点云数据中提取的点集合的索引,其维度为[B, S],即B个数据集,每个数据集中需要提取S个点的索引。
函数的输出为new_points,其维度为[B, S, C],即B个数据集,每个数据集中有S个点,每个点的属性有C个。
PointNetFeaturePropagation
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)
new_points = torch.cat([points1, interpolated_points], dim=-1)
else:
new_points = interpolated_points
new_points = new_points.permute(0, 2, 1)
for i, conv in enumerate(self.mlp_convs):
bn = self.mlp_bns[i]
new_points = F.relu(bn(conv(new_points)))
return new_points
该模块用于将低维特征点云上采样到高维特征点云。
输入参数:
- xyz1: 输入点云的位置信息,大小为 [B, C, N]
- xyz2: 采样后的点云的位置信息,大小为 [B, C, S]
- points1: 输入点云的特征信息,大小为 [B, D, N]
- points2: 采样后的点云的特征信息,大小为 [B, D, S]其中 B 表示 batch size,C 表示点的坐标数(可以是 2 或 3),N 和 S 分别表示输入点云和采样点云的点数,D 表示点的特征维度。
输出参数:
- new_points: 上采样后的新特征点云,大小为 [B, D', N],其中 D' 是输出的点云特征维度。具体实现流程如下:
1. 将输入点云和采样点云的位置信息转置为 [B, N, C] 和 [B, S, C]。
2. 计算输入点云中每个点到采样点云中最近的 3 个点之间的距离,并按照距离从小到大排序,得到距离和对应的索引值矩阵 dists 和 idx,大小为 [B, N, 3]。
3. 根据距离计算每个输入点云上的点对采样点云上最近 3 个点的插值权重,并根据权重进行线性插值,得到新的特征点云 interpolated_points,大小为 [B, N, D]。
4. 将输入点云的特征信息和新的特征点云拼接在一起,并按照维度转置为 [B, N, D'],其中 D' 是输出的点云特征维度。
5. 对新的特征点云进行若干个一维卷积和批归一化操作,得到最终上采样后的新特征点云 new_points。
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
首先,将变量xyz1和xyz2的维度顺序从(0, 1, 2)改为(0, 2, 1),即将第二个和第三个维度交换,这样做的目的是为了在后续计算中方便地进行矩阵计算。
然后,将变量points2的维度顺序也从(0, 1, 2)改为(0, 2, 1)。
接着,获取变量xyz1的形状,并将结果分别赋值给B、N和C三个变量。其中B代表batch size(批大小),N代表点云中点的数量,C代表每个点的坐标数(通常为3)。
最后,获取变量xyz2的形状,并将结果分别赋值给下划线、S和下划线三个变量。其中S代表另一个点云的点的数量。
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)
这是一个用于点云配准的代码段。它会根据输入的两组点云xyz1和xyz2,通过计算它们之间的距离,生成一组插值后的点云interpolated_points。
其中,如果S等于1,即只有一组点云,则直接将第二组点云points2重复N次得到interpolated_points;否则,先计算出每个点与距离最近的3个点之间的距离和排序后的索引,然后根据距离倒数和归一化权重,对第二组点云进行加权平均得到插值后的点云interpolated_points。
这个过程可以用于将两组不完全重叠的点云进行配准,使它们在空间上更加一致,方便后续处理。
if points1 is not None:
points1 = points1.permute(0, 2, 1)
new_points = torch.cat([points1, interpolated_points], dim=-1)
else:
new_points = interpolated_points
这段代码会根据输入的点云数据points1和插值后得到的点云数据interpolated_points,将它们拼接成新的点云数据new_points。
如果原始点云数据points1不为空,则首先需要对其进行维度调整,将维度从[B,N,C]调整为[B,C,N]。然后使用torch.cat函数将两组点云数据在最后一个维度上进行拼接,形成新的点云数据。拼接后的点云数据new_points的维度为[B, C, 2N]。
如果原始点云数据points1为空,则直接将插值后得到的点云数据作为新的点云数据,即new_points = interpolated_points。此时,new_points的维度为[B, C, N]。
new_points = new_points.permute(0, 2, 1)
for i, conv in enumerate(self.mlp_convs):
bn = self.mlp_bns[i]
new_points = F.relu(bn(conv(new_points)))
这段代码首先对新生成的点云数据new_points进行维度调整,将维度从[B, C, 2N]调整为[B, 2N, C],以适应接下来的卷积操作。
然后,通过循环遍历self.mlp_convs中的卷积层和对应的批归一化层self.mlp_bns,在每层卷积之后都添加ReLU激活函数,并使用批归一化进行特征缩放和偏移。
经过循环处理后,最终得到的new_points是经过多层卷积、批归一化和ReLU激活函数处理后得到的特征表示。
square_distance
def square_distance(src, dst):
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
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]
计算两组点云之间每个点之间欧式距离的函数