源代码链接:
源代码中没有日志记录的部分,所以就修改了一部分。
run_nn_seg.py
def get_arguments():
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='shapenetpart') # 71.27, 73.95
parser.add_argument('--bz', type=int, default=1) # Freeze as 1
parser.add_argument('--points', type=int, default=1024)
parser.add_argument('--stages', type=int, default=4)
parser.add_argument('--dim', type=int, default=144)
parser.add_argument('--k', type=int, default=90)
parser.add_argument('--de_k', type=int, default=6) # propagate neighbors in decoder
parser.add_argument('--alpha', type=int, default=1000)
parser.add_argument('--beta', type=int, default=100)
parser.add_argument('--gamma', type=int, default=300) # Best as 300
args = parser.parse_args()
return args
这是一个使用argparse模块解析命令行参数的Python函数。
它定义了几个参数:
--dataset
:数据集名称,默认值为'shapenetpart'。该参数用于指定要使用的数据集。--bz
:批量大小(batch size),默认值为1。该参数用于指定在每次训练中使用的样本数量。--points
:点云中点的数量默认值为1024。该参数用于指定点云模型中的点数目。--stages
:编码器和解码器阶段数默认值为4。该参数用于指定神经网络的层数。--dim
:编码器和解码器中隐藏单元的维度,默认值为144。该参数表示编码器和解码器中包含的神经元数目。--k
:用于构建局部区域邻居的最近邻点数目,默认值为90。该参数用于确定对于每个点选择多少个最接近的邻居来计算其特征。--de_k
:解码器中用于传播邻居的最近邻点数目,默认值为6。该参数用于指定解码器在计算上采样后每个点周围的点的数量。--alpha
:重建损失的权重,默认值为1000。该参数用于调整Autoencoder的编码和解码器之间的平衡。--beta
:邻域损失的权重,默认值为100。该参数用于调整局部区域邻居的影响力。--gamma
:同一性正则化(identity regularization)的权重,默认值为300。该参数用于控制隐空间中点与重构点的距离,使得它们更容易对齐。
在该函数中,--dim
参数用于指定编码器和解码器中隐藏单元的维度。这些隐藏单元是神经网络中的节点,它们接收输入并产生输出。隐藏单元的数量和维度通常是神经网络的超参数,需要进行调整以优化模型的性能。
举例来说,如果编码器和解码器中隐藏单元的维度为100,那么每个隐藏单元都包含100个权重值,并且神经网络中会有100个这样的节点。当神经网络接收输入数据时,这些隐藏单元会学习提取数据的不同特征,并将这些特征压缩到一个低维度的表示中。这可以使得神经网络更加高效地学习并处理数据。当解码器接收到这些压缩后的表示时,它会将其解压回原始数据的形式。
@torch.no_grad()
是一个装饰器函数,用于标记一个代码块,以便在执行该代码块时禁用梯度计算,从而减少内存的使用和加速计算。在PyTorch中,所有的Tensor都可以跟踪其上发生的运算,并自动地构建计算图和计算梯度,这种机制称为自动求导(Autograd)。但有时候我们需要在不需要计算梯度的情况下对Tensor进行操作,比如测试模型或生成样本等。这时可以使用@torch.no_grad()
来临时关闭自动求导的功能,避免浪费大量内存。
def log_string(str):
logger.info(str)
print(str)
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) # 获取当前时间并转换为标准字符串(年-月-日-时-分)
# 创建文件夹
exp_dir = Path('./log/') # 使用 Path 类创建一个路径对象 exp_dir,指定日志文件存储的根目录为 './log/'
exp_dir.mkdir(exist_ok=True) # 目录存在正常返回,不存在创建
exp_dir = exp_dir.joinpath('nn_seg') # 在 exp_dir 变量所代表的目录路径下创建一个名为 'nn_seg' 的子目录
exp_dir.mkdir(exist_ok=True)
# exp_dir = exp_dir.joinpath(timestr) # 在 exp_dir 变量所代表的目录路径下创建一个以时间为名的子目录
# exp_dir.mkdir(exist_ok=True)
log_dir = exp_dir.joinpath('logs/') # 训练日志文件
log_dir.mkdir(exist_ok=True)
args = get_arguments()
logger = logging.getLogger("nn_seg") # 创建了一个名为 "nn_seg" 的日志记录器 logger
logger.setLevel(logging.INFO) # 设置了日志记录器 logger 的日志级别为 INFO,即只记录 INFO 级别及以上的日志信息。
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s') # 日志格式化器 设置日志记录的格式。 时间-记录器名称-日志级别-内容
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, "nn_seg")) # 文件处理器,用于将日志信息写入到文件中
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
print('==> Loading args..')
log_string('==> Loading args..')
log_string(args)
print(args)
上述代码就是将每次训练的情况记录到对应路径下的文件。
point_nn = Point_NN_Seg(input_points=args.points, num_stages=args.stages,
embed_dim=args.dim, k_neighbors=args.k, de_neighbors=args.de_k,
alpha=args.alpha, beta=args.beta).cuda()
point_nn.eval()
这段代码是在创建一个名为
point_nn
的神经网络模型,并将其移动到 GPU 上进行计算(.cuda()
)。然后通过调用.eval()
方法告诉模型现在是处于评估状态。具体来说,这段代码使用名为
Point_NN_Seg
的自定义神经网络类初始化point_nn
对象,并传递了一些参数作为输入。它会接收一个由点构成的输入,在多个阶段中使用 k 近邻和空洞卷积等技术提取特征,最终输出每个点所属的类别或分割结果。
train_loader = DataLoader(PartNormalDataset(npoints=args.points, split='trainval', normalize=False),
num_workers=8, batch_size=args.bz, shuffle=False, drop_last=False)
test_loader = DataLoader(PartNormalDataset(npoints=args.points, split='test', normalize=False),
num_workers=8, batch_size=args.bz, shuffle=False, drop_last=False)
这段代码使用了名为
PartNormalDataset
的数据集类来创建训练集和测试集的数据加载器。在创建数据加载器时,代码传递了一些参数。其中
npoints
参数指定了每个点云中的点的数量;
split
参数指定了数据集中要使用的数据子集(训练集或测试集);
normalize
参数指定是否将输入数据规范化。
num_workers
参数指定了使用的线程数,
batch_size
参数指定了每次读取的数据批次大小,
shuffle
参数指定是否对数据进行随机洗牌,
drop_last
参数指定是否在最后一个批次有不足batch_size
个样本时丢弃该批次。
num_part, num_shape = 50, 16
# We organize point-memory bank by 16 shape labels
feature_memory = [[] for i in range(num_shape)]
label_memory = [[] for i in range(num_shape)]
这段代码初始化了两个变量
num_part
和num_shape
分别为 50 和 16。接着创建了两个空列表feature_memory
和label_memory
,每个列表的长度都等于num_shape
。这两个列表的目的是按照 16 种形状标签来组织点记忆库。
feature_memory
和label_memory
的每个元素对应于一个特定的形状标签。
# 本for循环主要是将输入数据进行预处理后,通过非参数编码器和解码器将点云特征转换为特征向量,然后提取每个形状的部件原型并存储在内存中。
# 点云数据(points)、形状标签(shape_label)、部件标签(part_label)和法向量(norm_plt)
for points, shape_label, part_label, norm_plt in tqdm(train_loader):
# pre-process
points = points.float().cuda().permute(0, 2, 1)
shape_label = shape_label.long().cuda().squeeze(1)
part_label = part_label.long().cuda()
这段代码是对输入数据进行预处理的部分,针对每个批次的数据,将点云数据
points
、形状标签shape_label
和部件标签part_label
进行以下三个操作:1.使用
float()
函数将点云数据从默认的数据类型转换为32位浮点数。2.使用
.cuda()
方法将数据移动到GPU上进行加速计算。3.使用
.permute(0, 2, 1)
方法将点云数据中每个点的坐标从(x,y,z)
的顺序重新排列(x,z,y)
的顺序。其中,
.squeeze(1)
方法用于去除形状标签张量中维度为1的维度,因为该维度在后续使用时可能会导致问题。
# Pass through the Non-Parametric Encoder + Decoder
# 通过非参数化的编码器+解码器
point_features = point_nn(points)
# All 2048 point features in a shape
point_features = point_features.permute(0, 2, 1) # bz, 2048, c
# Extracting part prototypes for a shape
# 提取一个形状的零件原型
feature_memory_list = []
label_memory_list = []
这段代码将经过预处理的点云数据
points
通过一个非参数化编码器和解码器(point_nn
)进行转换,得到一个包含所有点在特征空间内的表示的特征矩阵point_features
。这个特征矩阵的维度是(batch_size, num_points, feature_dim)
,其中num_points
为点云数据中点的数量,feature_dim
为每个点在特征空间内的维度。接着,将特征矩阵的维度重新排列成
(batch_size, feature_dim, num_points)
,并用于后续提取每个形状的部件原型。此外,创建了两个空列表feature_memory_list
和label_memory_list
,它们将分别用于存储每个形状的部件原型和标签信息。
for i in range(num_part):
# Find the point indices for the part_label within a shape
part_mask = (part_label == i)
if torch.sum(part_mask) == 0:
continue
# Extract point features for the part_label
# 提取part_label的点特征
part_features = point_features[part_mask]
# Obtain part prototypes by average point features for the part_label
# 通过对part_label的平均点特征获得零件原型
part_features = part_features.mean(0).unsqueeze(0)
feature_memory_list.append(part_features)
label_memory_list.append(torch.tensor(i).unsqueeze(0))
这段代码循环遍历每个形状中的所有部件,其中
num_part
指示了该形状中包含的部件数量。对于每个部件,首先使用布尔操作(part_label == i)
构建一个掩码(mask)来表示该形状中属于该部件的点的索引。如果掩码所代表的部件没有任何点,则跳过该部件的处理。否则,从特征矩阵point_features
中提取在该掩码内的所有点的特征,将这些特征求平均值得到该部件的部件原型,并添加到feature_memory_list
列表中。同时,将该部件的标签i
作为one-hot向量添加到label_memory_list
列表中。通过循环迭代完成后,
feature_memory_list
列表和label_memory_list
列表分别存储了该形状中所有部件的部件原型和标签信息。
# Feature Memory: store prototypes indexed by the corresponding shape_label
# 特征内存:存储原型,并以相应的形状标签为索引。
feature_memory_list = torch.cat(feature_memory_list, dim=0)
feature_memory[int(shape_label)].append(feature_memory_list)
# Label Memory: store labels indexed by the corresponding shape_label
# 标签存储器:存储由相应的shape_label索引的标签。
label_memory_list = torch.cat(label_memory_list, dim=0)
label_memory_list = F.one_hot(label_memory_list, num_classes=num_part)
label_memory[int(shape_label)].append(label_memory_list)
feature_memory存储了由相应的shape_label索引的原型,而label_memory则存储了同样由shape_label索引的标签。
基于代码片段,feature_memory_list和label_memory_list是包含特征和标签的张量。这些张量使用PyTorch的cat()函数沿着第一个维度连接在一起,创建一个表示特定shape_label的所有特征或标签的单个张量。
然后,将结果张量添加到feature_memory或label_memory字典中的适当列表中,其中shape_label用作索引键。需要注意的是,在将标签添加到label_memory字典之前,使用PyTorch的F.one_hot()函数对其进行了one-hot编码。
torch.cat
是PyTorch函数,用于沿着给定维度连接张量。在这个情况下,feature_memory_list
包含了多个Tensor的列表,正在沿着0维进行拼接。拼接后的结果张量将与输入张量具有相同的维数,但是0维的大小等于每个输入张量的0维大小之和。例如,如果
feature_memory_list
包含三个大小为(3,5),(4,5)和(2,5)的二维张量,则调用torch.cat(feature_memory_list,dim = 0)
将导致一个大小为(9,5)的单个二维张量,其中前三行对应于第一个张量,接下来的四行对应于第二个张量,最后两行对应于第三个张量。通过这种方式连接张量,可以将多个张量合并为一个更大的张量,从而使模型能够更有效地处理数据。
for i in range(num_shape):
# Feature Memory 特征
feature_memory[i] = torch.cat(feature_memory[i], dim=0)
feature_memory[i] /= feature_memory[i].norm(dim=-1, keepdim=True)
feature_memory[i] = feature_memory[i].permute(1, 0)
# print("Feature Memory of the " + str(i) + "-th shape is", feature_memory[i].shape)
str1 = "Feature Memory of the " + str(i) + "-th shape is", feature_memory[i].shape
log_string(str1)
# Label Memory 标签
label_memory[i] = torch.cat(label_memory[i], dim=0).cuda().float()
首先将特征和标签数据分别存储在feature_memory和label_memory数组中。然后,对于每一个形状(即循环变量i所代表的形状),对其特征数据进行归一化处理,并将其转置。最后,对于每一个形状,将其对应的标签数据合并为一个张量,并将其传输到GPU上进行浮点数类型转换。
logits_list, label_list = [], []
for points, shape_label, part_label, norm_plt in tqdm(test_loader):
# pre-process
points = points.float().cuda().permute(0, 2, 1)
shape_label = shape_label.long().cuda().squeeze(1)
part_label = part_label.long().cuda()
# Pass through the Non-Parametric Encoder + Decoder
# 通过非参数化的编码器+解码器
point_features = point_nn(points)
point_features = point_features.permute(0, 2, 1).squeeze(0) # 2048, c
point_features /= point_features.norm(dim=-1, keepdim=True)
# Similarity Matching
# 相似性匹配
Sim = point_features @ feature_memory[int(shape_label)]
# Label Integrate
# 标识整合
logits = (-args.gamma * (1 - Sim)).exp() @ label_memory[int(shape_label)]
logits_list.append(logits.unsqueeze(0))
label_list.append(part_label)
这段代码是一个形状识别的测试过程。其中,test_loader是已经准备好的数据加载器,用于从测试集中读取数据。
首先,这段代码通过循环遍历test_loader中的每个数据样本,并将其分别存储在points、shape_label、part_label和norm_plt变量中。
- points表示点云数据,
- shape_label表示形状标签,
- part_label表示部件标签,
- norm_plt则是一个归一化参数。
接下来,将点云数据points通过point_nn进行非参数化编码器+解码器的处理,得到对应的点云特征point_features。然后,对point_features进行归一化处理,并计算其与feature_memory[int(shape_label)]之间的相似度Sim。
最后,将Sim传入公式 (-args.gamma * (1 - Sim)).exp() @ label_memory[int(shape_label)] 中,得到一个维度为(1, N)的输出logits,并将其放入logits_list列表中。同时,将part_label也放入另一个列表label_list中。该过程可以重复多次,直至test_loader中所有数据样本被遍历完成。
logits_list = torch.cat(logits_list, dim=0)
label_list = torch.cat(label_list, dim=0)
# Compute mIoU
iou = compute_overall_iou(logits_list, label_list)
miou = np.mean(iou) * 100
print(f"Point-NN's part segmentation mIoU: {miou:.2f}.")
log_string(f"Point-NN's part segmentation mIoU: {miou:.2f}.")
这段代码片段是使用平均交并比(mIoU)作为评估指标来评估点云分割模型的性能。
前两行将预测的logits列表和实际标签列表沿着指定的维度(dim=0)连接起来,跨越了整个评估数据集。这样做的结果是,logits_list和label_list变成了包含所有预测的logits和实际标签的张量,跨越了整个评估数据集。
接下来,调用compute_overall_iou()函数并将连接后的logits和标签作为输入,它返回每个类别的IOU值。然后将IOU值平均得到mIoU。
下面就是Point-NN的Part Segmentation网络构建
point_nn_seg.py
# Non-Parametric Network
class Point_NN_Seg(nn.Module):
def __init__(self, input_points=2048, num_stages=5, embed_dim=144,
k_neighbors=128, de_neighbors=6, beta=1000, alpha=100):
super().__init__()
# Non-Parametric Encoder and Decoder
self.EncNP = EncNP(input_points, num_stages, embed_dim, k_neighbors, alpha, beta)
self.DecNP = DecNP(num_stages, de_neighbors)
def forward(self, x):
# xyz: point coordinates
# x: point features
xyz = x.permute(0, 2, 1)
# Non-Parametric Encoder
xyz_list, x_list = self.EncNP(xyz, x)
# Non-Parametric Decoder
x = self.DecNP(xyz_list, x_list)
return x
整体来看其使用非参数化方法构建的神经网络类,用于处理点云数据。该网络包含了输入点的数量、阶段数、嵌入维度、K邻居数量、解码器邻居数量等参数。该网络的核心部分是非参数编码器和解码器。在前向传递过程中,该网络将点云坐标作为输入并转换为xyz形式,然后将其输入到非参数编码器中,编码器生成编码点列表和特征点列表。接着,这些列表将传递给非参数解码器,解码器会根据编码点列表和特征点列表生成点云数据。最终输出点云的特征表示。
def __init__(self, input_points=2048, num_stages=5, embed_dim=144,
k_neighbors=128, de_neighbors=6, beta=1000, alpha=100):
super().__init__()
# Non-Parametric Encoder and Decoder
self.EncNP = EncNP(input_points, num_stages, embed_dim, k_neighbors, alpha, beta)
self.DecNP = DecNP(num_stages, de_neighbors)
这是 `Point_NN_Seg` 类的构造函数,它定义了模型的各个参数和模块。其中,
- `input_points` 参数指定了输入点云中点的个数,
- `num_stages` 参数指定了非参数编码器中使用的阶段数量,
- `embed_dim` 参数指定了编码器生成嵌入向量的维度,
- `k_neighbors` 参数指定了编码器中每个点考虑的邻居数量,
- `de_neighbors` 参数指定了解码器中每个点考虑的邻居数量,
- `beta` 和 `alpha` 参数分别用于计算损失函数中的两个权重。
接下来,构造函数创建了两个模块 `EncNP` 和 `DecNP`,它们分别是非参数化编码器和解码器。
`EncNP` 模块采用 `input_points`、`num_stages`、`embed_dim`、`k_neighbors`、`alpha` 和 `beta` 参数进行初始化。
`DecNP` 模块采用 `num_stages` 和 `de_neighbors` 参数进行初始化。
def forward(self, x):
# xyz: point coordinates xyz:点的坐标
# x: point features x:点的特征
xyz = x.permute(0, 2, 1)
# Non-Parametric Encoder
xyz_list, x_list = self.EncNP(xyz, x)
# Non-Parametric Decoder
x = self.DecNP(xyz_list, x_list)
return x
前向传播函数。
输入参数x为点云数据,其中包含每个点的坐标xyz和特征x(???)。permute函数是为了方便后面处理。
接下来通过Non-Parametric Encoder对输入数据进行编码,得到xyz_list和x_list两个输出结果。这一过程可以理解为将输入数据映射到一个高维空间中,并提取其中的特征信息。具体实现方式需要查看模型的结构和参数设置。
然后将编码后的xyz_list和x_list输入到Non-Parametric Decoder中进行解码,得到重建后的点云数据x。这一过程可以理解为从高维空间中重建出原始数据,在重建过程中加入了额外的特征信息。
最后将重建后的点云数据x作为输出返回。
非参数化编码器
整体
# Non-Parametric Encoder
class EncNP(nn.Module):
def __init__(self, input_points, num_stages, embed_dim, k_neighbors, alpha, beta):
super().__init__()
self.input_points = input_points
self.num_stages = num_stages
self.embed_dim = embed_dim
self.alpha, self.beta = alpha, beta
# Raw-point Embedding 原始点嵌入
self.raw_point_embed = PosE_Initial(3, self.embed_dim, self.alpha, self.beta)
self.FPS_kNN_list = nn.ModuleList() # FPS, kNN
self.LGA_list = nn.ModuleList() # Local Geometry Aggregation 局部几何图形聚合
self.Pooling_list = nn.ModuleList() # Pooling
out_dim = self.embed_dim
group_num = self.input_points
# Multi-stage Hierarchy 多阶段的层次结构
for i in range(self.num_stages):
out_dim = out_dim * 2
group_num = group_num // 2
self.FPS_kNN_list.append(FPS_kNN(group_num, k_neighbors))
self.LGA_list.append(LGA(out_dim, self.alpha, self.beta))
self.Pooling_list.append(Pooling(out_dim))
def forward(self, xyz, x):
# Raw-point Embedding 原始点嵌入
x = self.raw_point_embed(x)
xyz_list = [xyz] # [B, N, 3]
x_list = [x] # [B, C, N]
# Multi-stage Hierarchy 多阶段的层次结构
for i in range(self.num_stages):
# FPS, kNN
xyz, lc_x, knn_xyz, knn_x = self.FPS_kNN_list[i](xyz, x.permute(0, 2, 1))
# Local Geometry Aggregation 局部几何图形聚合
knn_x_w = self.LGA_list[i](xyz, lc_x, knn_xyz, knn_x)
# Pooling
x = self.Pooling_list[i](knn_x_w)
xyz_list.append(xyz)
x_list.append(x)
return xyz_list, x_list
`EncNP`类以输入点数、处理阶段数量、嵌入维度、每个阶段要考虑的最近邻居数量以及两个超参数(`alpha`和`beta`)作为输入。这段代码定义了一个使用PyTorch库实现的非参数编码器,用于处理点云数据。编码器包含多个分层处理阶段,逐渐降低输入数据的维度,并在不同尺度上提取特征。
`forward`方法通过迭代指定数量的处理阶段来执行输入点云数据的编码。在每个阶段,它首先应用点嵌入操作,然后执行一系列操作,包括最远点采样(FPS)、k最近邻搜索、局部几何聚合(LGA)和池化。每个阶段的输出存储在列表中,并由该方法返回。
# Raw-point Embedding 原始点嵌入
self.raw_point_embed = PosE_Initial(3, self.embed_dim, self.alpha, self.beta)
# 定义了三个 nn.ModuleList() 对象,分别为 FPS_kNN_list、LGA_list 和 Pooling_list。
self.FPS_kNN_list = nn.ModuleList() # FPS, kNN
self.LGA_list = nn.ModuleList() # Local Geometry Aggregation 局部几何图形聚合
self.Pooling_list = nn.ModuleList() # Pooling
out_dim = self.embed_dim
group_num = self.input_points
第一层称为“raw_point_embed”,它是“PosE_Initial”类的一个实例。该层用于初始化每个原始点在输入数据中的嵌入。该层以原始输入数据的维数(对于3D点数据为3)、所需嵌入的大小和两个额外参数alpha和beta作为输入。
第二层是名为“FPS_kNN_list”的模块列表。该模块列表包含执行最远点采样(FPS)和k-最近邻(kNN)操作的图层,用于输入点。
第三层是另一个名为“LGA_list”的模块列表。该模块列表包含执行局部几何聚合(LGA)的图层,在每个点上基于其邻居聚合局部几何特征。
最后,“Pooling_list”模块列表包含对LGA_list输出执行池化的层。这些池化操作用于降低特征映射的维度,并为原始点数据创建最终嵌入。
# Multi-stage Hierarchy 多阶段的层次结构
for i in range(self.num_stages):
out_dim = out_dim * 2
group_num = group_num // 2
self.FPS_kNN_list.append(FPS_kNN(group_num, k_neighbors))
self.LGA_list.append(LGA(out_dim, self.alpha, self.beta))
self.Pooling_list.append(Pooling(out_dim))
这段代码是一个多阶段的层次结构,其中包含了三个不同的列表:FPS_kNN_list、LGA_list 和 Pooling_list。在每一轮循环中,代码会执行以下步骤:
- 将输出维度 out_dim 乘以 2,从而使得每轮输出的特征数量翻倍;
- 将组数 group_num 除以 2,从而将输入数据分成更小的子集;
- 向 FPS_kNN_list 添加一个新的 FPS_kNN 层,该层使用 k_neighbors 进行最近邻采样;
- 向 LGA_list 添加一个新的 LGA 层,该层使用 out_dim、alpha 和 beta 这些参数进行局部几何聚合;
- 向 Pooling_list 添加一个新的 Pooling 层,该层用于汇总局部特征。
def forward(self, xyz, x):
# Raw-point Embedding 原始点嵌入
x = self.raw_point_embed(x)
xyz_list = [xyz] # [B, N, 3]
x_list = [x] # [B, C, N]
# Multi-stage Hierarchy 多阶段的层次结构
for i in range(self.num_stages):
# FPS, kNN
xyz, lc_x, knn_xyz, knn_x = self.FPS_kNN_list[i](xyz, x.permute(0, 2, 1))
# Local Geometry Aggregation 局部几何图形聚合
knn_x_w = self.LGA_list[i](xyz, lc_x, knn_xyz, knn_x)
# Pooling
x = self.Pooling_list[i](knn_x_w)
xyz_list.append(xyz)
x_list.append(x)
return xyz_list, x_list
输入参数为
xyz
和x
。首先,对于输入的x
进行原始点嵌入操作,即将每个点的特征向量映射到一个低维向量表示。之后,通过多阶段的层次结构进行特征提取。具体来说,首先使用 FPS(最远点采样) 和 kNN(k近邻) 操作选择当前点云中固定数量的关键点和它们的 k 个最近邻点,并生成它们的坐标和特征表示。然后,使用局部几何图形聚合操作将这些点及其特征聚合为新的特征向量。接着,使用池化操作对这些新的特征向量进行降维,并在下一阶段的特征提取中重复这个过程。
最后,该函数返回了
xyz_list
和x_list
两个列表,分别包含了每个阶段中选择的关键点坐标和对应的特征表示。这些列表可用于后续任务的处理,如分类、分割等。
原始点嵌入的PosE
# PosE for Raw-point Embedding
class PosE_Initial(nn.Module):
def __init__(self, in_dim, out_dim, alpha, beta):
super().__init__()
self.in_dim = in_dim
self.out_dim = out_dim
self.alpha, self.beta = alpha, beta
def forward(self, xyz):
B, _, N = xyz.shape
feat_dim = self.out_dim // (self.in_dim * 2)
feat_range = torch.arange(feat_dim).float().cuda()
dim_embed = torch.pow(self.alpha, feat_range / feat_dim)
div_embed = torch.div(self.beta * xyz.unsqueeze(-1), dim_embed)
sin_embed = torch.sin(div_embed)
cos_embed = torch.cos(div_embed)
position_embed = torch.stack([sin_embed, cos_embed], dim=4).flatten(3)
position_embed = position_embed.permute(0, 1, 3, 2).reshape(B, self.out_dim, N)
return position_embed
`in_dim` 表示输入的维度,
`out_dim` 表示输出的维度,
`alpha` 和 `beta` 是两个超参数。
在 前向传播中,首先获取输入张量 `xyz` 的形状,将输出维度平均分成 `in_dim` 个部分,并将每个部分划分为 `[sin(embedding)]` 和 `[cos(embedding)]` 两个部分。这些嵌入向量组合起来,成为位置编码向量。
这个过程可以被看作是为原始点云数据引入一些位置信息,因为点云数据本身并没有固定的位置或坐标信息,而这些信息对于很多应用来说是很有用的。这种位置编码方法也可以应用于其他类型的序列数据,例如自然语言处理或时间序列数据。
def forward(self, xyz):
B, _, N = xyz.shape
feat_dim = self.out_dim // (self.in_dim * 2)
feat_range = torch.arange(feat_dim).float().cuda()
dim_embed = torch.pow(self.alpha, feat_range / feat_dim)
div_embed = torch.div(self.beta * xyz.unsqueeze(-1), dim_embed)
sin_embed = torch.sin(div_embed)
cos_embed = torch.cos(div_embed)
position_embed = torch.stack([sin_embed, cos_embed], dim=4).flatten(3)
position_embed = position_embed.permute(0, 1, 3, 2).reshape(B, self.out_dim, N)
return position_embed
接受张量 `xyz`,表示空间中点的坐标,然后返回一个包含位置嵌入的张量`position_embed`。
函数的前几行从 `xyz` 张量中提取批次大小 `B`、点的数量 `N` 和特征维度 `feat_dim`。输出维度计算为 `self.out_dim = self.in_dim * 2 * feat_dim`,其中 `self.in_dim` 是模型的参数之一。
接下来,代码通过将可学习的标量参数 `alpha` 提升到介于 0 到 `feat_dim` 之间的一系列等间距值的幂次,并将结果除以 `beta` 参数乘以 `xyz` 张量,构建了一个嵌入矩阵。这会生成形状为 `(B, 3, N, feat_dim)` 的张量 `div_embed`。
对 `div_embed` 应用逐元素的 `sin` 和 `cos` 函数,得到两个形状与 `div_embed` 相同的张量 `sin_embed` 和 `cos_embed`。将它们沿着最后一个轴连接起来,得到一个形状为 `(B, 3, N, feat_dim, 2)` 的张量。然后将此张量展平成形状为 `(B, 3, N, 2*feat_dim)` 的张量。
最后,将张量排列成形状为 `(B, 2*feat_dim, N, 3)`,然后重塑成形状为 `(B, self.out_dim, N)`,并返回该张量。
FPS + K-NN
# FPS + k-NN
class FPS_kNN(nn.Module):
def __init__(self, group_num, k_neighbors):
super().__init__()
self.group_num = group_num
self.k_neighbors = k_neighbors
def forward(self, xyz, x):
B, N, _ = xyz.shape
# FPS
fps_idx = pointnet2_utils.furthest_point_sample(xyz, self.group_num).long()
lc_xyz = index_points(xyz, fps_idx)
lc_x = index_points(x, fps_idx)
# kNN
knn_idx = knn_point(self.k_neighbors, xyz, lc_xyz)
knn_xyz = index_points(xyz, knn_idx)
knn_x = index_points(x, knn_idx)
return lc_xyz, lc_x, knn_xyz, knn_x
这段代码实现了一个基于 FPS(Furthest Point Sampling)和 k-NN(k-Nearest Neighbors)的子采样方法。该方法可以用于点云数据处理,其原理是先使用 FPS 方法选取一组关键点,然后在选取的关键点周围使用 k-NN 方法选择其他点。
具体来说,给定输入张量 `xyz`(形状为 [B, N, 3] 表示 N 个点的坐标)和特征张量 `x`(形状为[B, N, C]表示每个点对应的特征向量),该方法的输入参数 `group_num` 和 `k_neighbors` 分别表示 FPS 和 k-NN 的超参数。
在前向传播过程中,首先使用 FPS 方法选取 `group_num` 个关键点,并将这些关键点的坐标保存在 `lc_xyz` 张量中,将这些关键点对应的特征向量保存在 `lc_x` 中。然后,在这些关键点周围使用 k-NN 方法选择每个点的 `k_neighbors` 个最近邻,并将这些邻居点的坐标保存在 `knn_xyz` 张量中,将这些点对应的特征向量保存在 `knn_x` 中。
最终输出 `lc_xyz`、`lc_x`、`knn_xyz` 和 `knn_x` 四个张量,表示:选取的关键点的坐标和特征向量分别保存在 `lc_xyz` 和 `lc_x` 中,每个关键点周围的 `k_neighbors` 个最近邻的坐标和特征向量分别保存在 `knn_xyz` 和 `knn_x` 中。
这种子采样方法可以用于减少点云数据的规模,从而降低计算复杂度,在某些点云任务中取得较好效果。
def forward(self, xyz, x):
B, N, _ = xyz.shape
# FPS
fps_idx = pointnet2_utils.furthest_point_sample(xyz, self.group_num).long()
lc_xyz = index_points(xyz, fps_idx)
lc_x = index_points(x, fps_idx)
# kNN
knn_idx = knn_point(self.k_neighbors, xyz, lc_xyz)
knn_xyz = index_points(xyz, knn_idx)
knn_x = index_points(x, knn_idx)
return lc_xyz, lc_x, knn_xyz, knn_x
`furthest_point_sample` 函数,实现了 FPS(Furthest Point Sampling)算法。该算法是一种常用的点云子采样方法,其目的是从输入的点云数据中选取 k个最具代表性的点,作为后续处理的关键点。
具体来说,给定输入张量 `xyz`(形状为 [B, N, 3]表示N个点的坐标)和超参数 `group_num`,该函数的输出 `fps_idx` 是一个形状为[B, k]的整数张量,表示选取的k个关键点在原始点云中的下标。这里的 `k` 等于 `group_num`。
在实现过程中,该函数首先随机选择一个起始点作为第一个关键点,并计算其他所有点到该点的欧几里得距离。然后,在剩下的点中找出与已选点距离最远的点,并将其加入关键点集合。这个过程重复进行,直到选出 `k` 个关键点。
在返回结果之前,`fps_idx` 还需要进行类型转换,将浮点数类型的下标转换为整数类型。因此使用了 `.long()` 函数对结果进行类型转换。
`index_points()`的功能是根据给定的下标,在输入张量中提取对应的子集。
具体来说,`index_points()` 的输入包括一个原始张量 `points`(形状为[B, N, C]表示N个点的坐标或特征向量)、一个下标张量 `idx`(形状为[B, M]表示需要选取的点在原始张量中的下标)以及一个可选的布尔型参数 `keep_dim`(表示是否保留维度),输出则是选取的子集张量 `new_points`(形状为[B, M, C])。
在这里,`fps_idx` 是一个形状为[B, k]的整数张量,表示从原始点云数据中选取的 $k$ 个关键点的下标。因此,`lc_xyz` 和 `lc_x` 就分别是从输入的点云坐标和特征张量中选取这些关键点对应的子集,即形状分别为[B, k, 3]和[B, k, C]的张量。
这个过程相当于通过 FPS 算法从原始点云数据中选出了一部分关键点,并将这些点的坐标和特征向量保存到 `lc_xyz` 和 `lc_x` 中,以便进行后续操作(如 k-NN 子采样等)。
`knn_point()`实现了 k-NN(k-Nearest Neighbors)算法。该算法是一种常用的点云子采样方法,其目的是对于每个关键点,选取它周围最近的k个邻居点作为后续处理的重要对象。
具体来说,给定输入张量 `xyz`(形状为[B, N, 3]表示N个点的坐标)、选取的关键点坐标张量 `lc_xyz`(形状为[B, k, 3]表示k个关键点的坐标)以及超参数 `k_neighbors`,该函数的输出 `knn_idx` 是一个形状为[B, k, k\_neighbors]的整数张量,表示每个关键点周围最近的k\_neighbors个邻居点在原始点云中的下标。
在实现过程中,该函数首先计算出每个关键点到所有点的欧几里得距离,并根据距离从小到大排列。然后,选择与每个关键点距离最近的k\_neighbors个点,将它们的下标加入 `knn_idx` 中。(这个过程可以使用第三方库 Faiss 实现)
局部几何图形聚合
# Local Geometry Aggregation
class LGA(nn.Module):
def __init__(self, out_dim, alpha, beta):
super().__init__()
self.geo_extract = PosE_Geo(3, out_dim, alpha, beta)
def forward(self, lc_xyz, lc_x, knn_xyz, knn_x):
# Normalize x (features) and xyz (coordinates)
mean_x = lc_x.unsqueeze(dim=-2)
std_x = torch.std(knn_x - mean_x)
mean_xyz = lc_xyz.unsqueeze(dim=-2)
std_xyz = torch.std(knn_xyz - mean_xyz)
knn_x = (knn_x - mean_x) / (std_x + 1e-5)
knn_xyz = (knn_xyz - mean_xyz) / (std_xyz + 1e-5)
# Feature Expansion
B, G, K, C = knn_x.shape
knn_x = torch.cat([knn_x, lc_x.reshape(B, G, 1, -1).repeat(1, 1, K, 1)], dim=-1)
# Geometry Extraction
knn_xyz = knn_xyz.permute(0, 3, 1, 2)
knn_x = knn_x.permute(0, 3, 1, 2)
knn_x_w = self.geo_extract(knn_xyz, knn_x)
return knn_x_w
这段代码实现了 LGA(Local Geometry Aggregation)模块。给定两个点云,即局部中心点云 `lc_xyz` 和其邻域点云 `knn_xyz`,以及它们对应的特征矩阵 `lc_x` 和 `knn_x`,该模块通过以下步骤将邻域信息与中心点信息进行融合:
1. 将待融合的邻域特征矩阵 `knn_x` 进行标准化处理,使其特征均值为 0、方差为 1。
2. 将局部中心点特征矩阵 `lc_x` 复制并拼接到之前标准化后的邻域特征矩阵 `knn_x` 的最后一维上,得到新的特征矩阵。
3. 将局部中心点坐标 `lc_xyz` 和邻域点坐标 `knn_xyz` 作为输入,经过 PosE_Geo 模块进行几何信息提取,得到一个新的特征矩阵 `knn_x_w`。
4. 返回 `knn_x_w` 作为聚合后的邻域特征矩阵。
在这个过程中,PosE_Geo 模块通过注意力机制和多头机制,使用局部中心点和邻域点之间的相对位置和距离等几何信息,对邻域特征进行加权和聚合。通过这一过程,LGA 模块能够有效地融合局部中心点和其邻域的信息,提高点云的表示能力。
非参数化解码器
class DecNP(nn.Module):
def __init__(self, num_stages, de_neighbors):
super().__init__()
self.num_stages = num_stages
self.de_neighbors = de_neighbors
def propagate(self, xyz1, xyz2, points1, points2):
"""
Input:
xyz1: input points position data, [B, N, 3] 输入点的位置数据
xyz2: sampled input points position data, [B, S, 3] 采样的输入点位置数据
points1: input points data, [B, D', N] 输入点的数据
points2: input points data, [B, D'', S] 输入点的数据
Return:
new_points: upsampled points data, [B, D''', N] 采样后的点数据
"""
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[:, :, :self.de_neighbors], idx[:, :, :self.de_neighbors] # [B, N, 3]
dist_recip = 1.0 / (dists + 1e-8)
norm = torch.sum(dist_recip, dim=2, keepdim=True)
weight = dist_recip / norm
weight = weight.view(B, N, self.de_neighbors, 1)
index_points(xyz1, idx)
interpolated_points = torch.sum(index_points(points2, idx) * weight, 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)
return new_points
def forward(self, xyz_list, x_list):
xyz_list.reverse()
x_list.reverse()
x = x_list[0]
for i in range(self.num_stages):
# Propagate point features to neighbors
# 将点状特征传播给邻居
x = self.propagate(xyz_list[i + 1], xyz_list[i], x_list[i + 1], x)
return x
DecNP的目标是将稀疏点集作为输入,并输出具有学习特征的更密集的点云。
DecNP模块的输入包括
位置数据(xyz1)和点数据(points1)表示的稀疏点集
位置数据(xyz2)和点数据(points2)表示的更稀疏的点集。
然后,模块执行邻域传播步骤,其中它将从xyz2到xyz1传播点特征。
在传播步骤期间,模块首先使用由square_distance()定义的距离度量为xyz1中的每个点计算k个最近邻居。然后,它根据每个邻居与源点的距离计算权重,这些权重被归一化为所有k个邻居的距离之和。然后,这些权重用于插值与每个点的k个最近邻居相关联的特征向量。
最后,将插值特征与原始点特征连接起来,并通过一个密集层产生具有学习特征的输出点云。此过程重复多个阶段,以逐渐增加点云的密度。
def propagate(self, xyz1, xyz2, points1, points2):
"""
Input:
xyz1: input points position data, [B, N, 3] 输入点的位置数据
xyz2: sampled input points position data, [B, S, 3] 采样的输入点位置数据
points1: input points data, [B, D', N] 输入点的数据
points2: input points data, [B, D'', S] 输入点的数据
Return:
new_points: upsampled points data, [B, D''', N] 采样后的点数据
"""
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[:, :, :self.de_neighbors], idx[:, :, :self.de_neighbors] # [B, N, 3]
dist_recip = 1.0 / (dists + 1e-8)
norm = torch.sum(dist_recip, dim=2, keepdim=True)
weight = dist_recip / norm
weight = weight.view(B, N, self.de_neighbors, 1)
index_points(xyz1, idx)
interpolated_points = torch.sum(index_points(points2, idx) * weight, 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)
return new_points
这是一个用于点云上采样的方法,输入包含两个点云:xyz1 和 xyz2,以及它们对应的特征数据 points1 和 points2。其中,xyz1 是原始的点云,而 xyz2 是在其上进行采样得到的点云。函数的主要目的是将 points2 的特征插值到 xyz1 上,并返回新的点云特征数据 new_points。
具体来说,该函数首先通过计算 xyz1 和 xyz2 之间的距离来确定每个 xyz1 点最近的几个 xyz2 点的索引和权重。然后,根据权重插值得到新的点云特征 interpolated_points。最后,如果原始的点云特征 points1 不为空,则将 interpolated_points 和 points1 拼接在一起作为新的点云特征 new_points 返回;否则直接返回 interpolated_points。
详细解释:
这是一个名为`propagate`的方法,它接受四个输入:`xyz1`、`xyz2`、`points1`和`points2`。该函数的目的是对点云数据进行上采样。
以下是每个输入表示的含义:
- - `xyz1`:大小为`[B,N,3]`的张量,表示输入点的位置数据。
- - `xyz2`:大小为`[B,S,3]`的张量,表示上采样点的位置数据。
- - `points1`:大小为`[B,D',N]`的张量,表示输入点的特征数据。此张量是可选的,可以设置为`None`。
- - `points2`:大小为`[B,D'',S]`的张量,表示上采样点的特征数据。
步骤如下:
1. 对`points2`进行置换,使其具有维度`[B,S,D'']`。
2. 获取`xyz1`张量的形状以及`xyz2`中的点数。
3. 如果`S`等于1,则将`points2`重复一次以便与`xyz1`中的每个点相匹配。
4. 否则,使用`square_distance`函数(未在此处显示)计算`xyz1`和`xyz2`之间的平方距离。然后按最后一个维度对距离进行排序,并仅保留`xyz1`中每个点的最近`self.de_neighbors`个邻居。
5. 计算距离的倒数,并通过将它们除以所有倒数之和来归一化这些值。将结果的形状重塑为`[B,N,self.de_neighbors,1]`维度张量。
6. 使用`index_points`函数(未在此处显示)选择`xyz2`中最近邻居的位置,然后将这些位置与相应的权重相乘并求和,以获得插值点。
7. 如果`points1`不为None,则将其置换为`[B,N,D']`维度,将其与插值点沿最后一个维度连接起来,然后将组合后的张量重新排列为`[B,D''', N]`维度。否则,只需重新排列插值点张量,使其具有`[B,D''', N]`维度。
8. 返回上采样的点云数据。
def forward(self, xyz_list, x_list):
xyz_list.reverse()
x_list.reverse()
x = x_list[0]
for i in range(self.num_stages):
# Propagate point features to neighbors
# 将点状特征传播给邻居
x = self.propagate(xyz_list[i + 1], xyz_list[i], x_list[i + 1], x)
return x
输入的 xyz_list 和 x_list 反转,并且从反转后的列表中取出第一个元素 x 作为初始值。然后,它使用一个 for 循环进行多层传播,每次传播将当前层的点的坐标和特征与下一层的点的坐标和特征作为输入,通过 propagate 函数进行处理,并将处理后的结果作为下一层的输入。最终,该函数返回经过多层传播后的输出结果。
数据读取
train_loader = DataLoader(PartNormalDataset(npoints=args.points, split='trainval', normalize=False),
num_workers=8, batch_size=args.bz, shuffle=False, drop_last=False)
test_loader = DataLoader(PartNormalDataset(npoints=args.points, split='test', normalize=False),
num_workers=8, batch_size=args.bz, shuffle=False, drop_last=False)
PyTorch 中的数据加载器(DataLoader)所需的部分参数和属性。
- dataset(数据集):加载数据的数据集。
- batch_size (int, optional): 每批要加载多少个样本。(默认:1)。
- shuffle (bool, 可选): 设置为True',以便在每个周期重新洗刷数据。(默认: False)。num_workers (int, optional): 多少个子进程用于数据加载。装载。0表示数据将在主进程中加载。 (默认:0)
- drop_last (bool, optional): 设置为True,以放弃最后一个不完整的批次、如果数据集的大小不能被批次的大小所分割。如果False和数据集的大小不能被批次的大小所分割,那么最后一批将会更小。(默认: False)
PartNormalDataset(npoints=args.points, split='trainval', normalize=False)
一个用于处理3D点云数据的PyTorch数据集类。该类可以从数据集中加载具有指定数量点数的3D点云,并将其分为训练和验证集。normalize参数控制是否对点云进行归一化处理,如果
设置为True,则会对点云进行均值和方差归一化处理。
#打开文件
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()}
这段代码是PartNormalDataset类的构造函数中的一部分,用于读取类别名称和对应的编号。self.catfile指定了存储类别信息的文件路径,通过打开该文件并逐行读取来获取类别名称和编号。ls[0]表示类别编号,ls[1]表示类别名称。将这些键值对存储在一个字典self.cat中,并使用字典推导式将其转换为相同的形式(即去除空格)。这样,在加载点云数据时可以使用类别编号来索引对应的类别名称。
ls = line.strip().split()
去除首尾空白字符并按照空格符分割成一个字符串列表。strip()方法用于删除字符串首尾的空格和换行符等无意义字符,split()方法则使用空格符(默认分隔符)将字符串分割成多个子字符串,并返回一个字符串列表ls。这里假设每行文本中包含两个元素,因此可以使用ls[0]和ls[1]来访问类别编号和名称。
elf.cat[ls[0]] = ls[1]
这行代码的作用是将类别编号和名称存储在self.cat字典中。`ls[0]`表示类别编号,`ls[1]`表示类别名称,将它们作为键值对存储在self.cat字典中,即使用类别编号作为键,类别名称作为值。这样可以方便地通过类别编号来查找对应的类别名称。
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)])
使用了Python的内置库os和json来读取JSON文件。该文件包含训练集中所有点云文件的路径信息,是一个JSON格式的文件。具体地,该行代码通过os.path.join()方法拼接出完整的文件路径,并将其传递给open()函数以打开文件。第二个参数'r'表示以只读模式打开文件。
train_ids = set([str(d.split('/')[2]) for d in json.load(f)])
这段使用了列表推导式和set()函数来提取训练集中所有点云文件的文件名,并将它们存储在train_ids集合中。具体地,json.load(f)从打开的JSON文件对象f中读取JSON数据,并返回一个包含所有数据的列表。然后,列表推导式遍历该列表中的每个元素d,并将其传递给str.split('/')方法以切割出文件路径中的各个部分。由于文件名位于路径的第三个部分(索引为2),因此可以通过d.split('/')[2]获取到文件名。最后,将所有文件名转换为字符串类型并存储在train_ids集合中。set()函数用于创建一个集合对象,并自动去掉其中的重复元素。因此,train_ids集合中存储了训练集中所有点云文件的唯一文件名。
for item in self.cat:
self.meta[item] = []
dir_point = os.path.join(self.root, self.cat[item])
fns = sorted(os.listdir(dir_point))
# 从fns中筛选出训练集和验证集所需要的文件名。
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)
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.meta字典中。这个方法使用了os模块来获取指向数据集文件的路径,并根据split参数的值选择要加载的文件。根据选择的文件,将数据集文件的路径添加到适当的self.meta[item]列表中。
第二个for循环,它遍历了列表fns中的每个文件名fn,并进行了以下操作:
- 首先使用os.path.basename()函数获取fn的基础文件名(即去掉路径部分的文件名)。
- 然后使用os.path.splitext()函数将基础文件名分离出文件名和扩展名,得到一个元组,这里只取了第一个元素即文件名,赋值给变量token。
- 接着使用os.path.join()函数将dir_point和token拼接成完整的文件路径,并将其添加到self.meta[item]列表中。
换句话说,这段代码的作用是把每个文件名fn转化为对应的完整文件路径,并将这些完整文件路径添加到一个字典self.meta中对应item项的值的末尾。这通常用于数据集加载操作,方便后续处理。
self.datapath = []
for item in self.cat:
for fn in self.meta[item]:
self.datapath.append((item, fn))
创建了一个空列表
datapath
,然后对于cat
列表中的每个元素,在相应的meta
字典中循环遍历每个文件名,并将一个元组添加到datapath
中,该元组包含类别名称(来自cat
列表)和文件名(来自相应的meta
字典)。因此,最终
datapath
将是一个元组的列表,其中每个元组表示一个文件路径,由类别名称(来自cat
列表)和文件名(来自相应的meta
字典)组成。示例 meta = {'动物': ['狗.jpg', '猫.jpg', '鸟.jpg'],
'植物': ['树.jpg', '花.jpg']}那么执行这段代码后,`datapath` 列表将变成:
[('动物', '狗.jpg'), ('动物', '猫.jpg'), ('动物', '鸟.jpg'), ('植物', '树.jpg'), ('植物', '花.jpg')]其中每个元组表示一个文件路径,例如 `('动物', '狗.jpg')` 表示 `动物` 类别中的 `狗.jpg` 文件,以此类推。
self.classes = dict(zip(self.cat, range(len(self.cat))))
# 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]}
self.cache = {} # from index to (point_set, cls, seg) tuple
self.cache_size = 20000
self.classes = dict(zip(self.cat, range(len(self.cat))))
创建一个字典,将 self.cat 中的元素作为键(key),对应的值(value)是该元素在 self.cat 列表中的索引值。具体来说,range(len(self.cat)) 生成一个从 0 到 self.cat 的长度减一的整数序列,zip 函数将 self.cat 和这个整数序列一一对应起来,然后 dict 函数将它们转化为字典类型。这样做的目的常常是方便地通过类别名称查找其索引或者通过索引查找对应的类别名称。
这段代码实例化了一个类,该类具有以下属性:
- cat: 类别列表
- classes: 从类别到整数标签的字典映射
- seg_classes: 从类别到分割标签列表的字典映射
- cache: 从索引到(点集、类别、分割)元组的字典缓存
- cache_size: 缓存大小`zip()` 方法将 `cat` 列表中的元素与一系列递增数字(从 0 开始)相匹配,生成一个从类别字符串到整数标签的字典映射。例如,如果 `cat = ['Cat', 'Dog', 'Bird']`,则 `classes = {'Cat': 0, 'Dog': 1, 'Bird': 2}`。
`seg_classes` 字典映射将每个类别字符串与一个分割标签列表相关联。例如,'Earphone' 类别对应的分割标签为 [16, 17, 18]。
`cache` 是一个字典缓存,用于存储从索引到其对应点集、类别和分割标签的元组。这可以用于快速访问模型预处理的数据。
`cache_size` 属性定义了缓存的最大大小。当缓存超出此限制时,最早添加的条目将被删除。
if index in self.cache:
point_set, normal, seg, cls = 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)
point_set = data[:, 0:3]
normal = data[:, 3:6]
seg = data[:, -1].astype(np.int32)
if len(self.cache) < self.cache_size:
self.cache[index] = (point_set, normal, seg, cls)
如果索引在缓存中存在:
从缓存中获取点集、法向量、分割结果和类别信息,赋值给变量point_set、normal、seg和cls。
否则:
获取数据路径fn和类别cat。
根据类别获取对应的整数标签cls,并将其转换为numpy数组。
加载文件fn[1]中的数据,将其转换为浮点型,并赋值给变量data。
从data中提取点集,法向量和分割结果,分别赋值给变量point_set、normal和seg。
如果缓存大小没超过限制,则将点集、法向量、分割结果和类别信息保存到缓存中。
if self.normalize:
point_set = pc_normalize(point_set)
choice = np.random.choice(len(seg), self.npoints, replace=True)
这部分代码可能是在数据准备阶段做的一些操作:
如果self.normalize为True,则对point_set进行归一化处理,调用了pc_normalize函数。
从seg中随机选择npoints个元素,并将其索引保存到变量choice中。其中replace参数为True表示可以重复选择同一个元素。