表达旋转：四元数
原理简析
详细可以看SLAM14讲关于四元数的部分，如果看不懂，可以看相关资料中的虚数视频。
四元数在点云中是一个很好的表示旋转的方式。
一个三维空间点用四元数表示如下
P = [ 0 , x , y , z ] = [ 0 , v ] P = [0, x, y, z] = [0, v] P = [ 0 , x , y , z ] = [ 0 , v ] 三维空间的点是个纯虚四元数（实部为0），三个虚部与空间中的三个轴对应。
怎么用四元数表示一个旋转呢？如下所示
q = [ cos θ 2 ， n sin θ 2 ] q = \left [ \cos \frac{\theta }{2}， n\sin \frac{\theta }{2} \right ] q = [ cos 2 θ ， n sin 2 θ ] 左边的cos是实部，也就是说，实部存储的是旋转角度
右边的sin是虚部，n表示的是旋转轴，旋转轴有三个维度，所以虚部有三个数，每个数都要分别乘以一个θ/2。也就是说，虚部存放的是旋转轴，表示物体绕着哪个轴进行旋转
那么，一个点绕着一个轴进行旋转，旋转后的点怎么计算出来呢？
p ′ = q p q − 1 p' = qpq^{-1} p ′ = qp q − 1 其中，p'是结果，通过原点p左乘旋转四元数，右乘旋转四元数的转置得到，结果就是新的点p的四元数数值
一个单位四元数（即模长为1的四元数）可以表示三维空间中的任意旋转，
用Numpy实现四元数运算
Copy import numpy as np
x = np.radians(90) # 绕x轴进行90度运算
x_vector = np.array([1, 0, 0]) # 定义旋转轴，x轴的单位向量
def quaternion_multiply(quaternion1, quaternion0): # 四元数相乘函数
w0, x0, y0, z0 = quaternion0
w1, x1, y1, z1 = quaternion1
return np.array([-x1 * x0 - y1 * y0 - z1 * z0 + w1 * w0,
x1 * w0 + y1 * z0 - z1 * y0 + w1 * x0,
-x1 * z0 + y1 * w0 + z1 * x0 + w1 * y0,
x1 * y0 - y1 * x0 + z1 * w0 + w1 * z0], dtype=np.float64)
def quaternion_contrary(quaternion): # 四元数取逆函数
w, x, y, z = quaternion
q_star = np.array([w, -x, -y, -z]) # 共轭
q_abs_pow2 = np.sum(np.power(quaternion, 2)) # 模长的平方
return q_star / q_abs_pow2
p = np.array([0, 0, 2, 0]) # 定义一个点，在y轴上 (0, 2, 0) 如果此点绕x轴旋转，会变成 (0, 0, 2)
q_x = np.hstack((np.cos([x / 2]), np.sin(x / 2) * x_vector)) # 旋转四元数 (cosx/2, nsinx/2)
print(q_x)
print(quaternion_multiply(quaternion_multiply(q_x, p), quaternion_contrary(q_x)))python 最终输出的结果为
Copy [0.70710678 0.70710678 0. 0. ]
[0.0000000e+00 0.0000000e+00 4.4408921e-16 2.0000000e+00] 第一行，这是一个四元数，表示绕x轴旋转90度。这个四元数的实部为cos(旋转角/2)，虚部的第一个元素为1*sin(旋转角/2)，其余元素为0，符合预期，虚部的三个元素表示旋转轴(1, 0, 0)的方向乘以sin(旋转角度/2)。
第二行，表示结果的四元数表示，第一维度为0，也就是实部为0，表示是一个点。三个虚部表示点的坐标，其中第二个虚部非常小，可以看作是0，所以得到的坐标为(0, 0, 2)，符合预期。
解决思路
学习旋转角度，也就是分别绕着x轴、y轴、z轴旋转的角度，跟欧拉角差不多，但是用四元数进行计算，能够避免欧拉角的奇异性问题。
损失函数：两个四元数a和b，a乘以b的逆，能得到一个新的四元数，新的四元数表示了一个绕着新的轴进行的一次a到b的相对旋转，也就是说损失是拿这个新的四元数的实部的旋转角，旋转角越小表示两个旋转四元数越接近
网络结构
源码
四元数计算类 Quaternion
Copy import torch
class Quaternion:
def __init__(self, w, x, y, z):
self.value = torch.tensor([w, x, y, z])
def multiply(self, quaternion): # 与另一个四元数相乘
w0, x0, y0, z0 = self.value
w1, x1, y1, z1 = quaternion
return torch.tensor([-x1 * x0 - y1 * y0 - z1 * z0 + w1 * w0,
x1 * w0 + y1 * z0 - z1 * y0 + w1 * x0,
-x1 * z0 + y1 * w0 + z1 * x0 + w1 * y0,
x1 * y0 - y1 * x0 + z1 * w0 + w1 * z0])
def conjugate(self): # 取共轭数
w, x, y, z = self.value
q_star = torch.tensor([w, -x, -y, -z]) # 共轭
return q_star
def contrary(self): # 取转置
q_star = self.conjugate()
q_abs_pow2 = torch.sum(self.value**2) # 模长的平方
return q_star / q_abs_pow2
def getRotationMatrix(self): # 从四元数转旋转矩阵
q0, q1, q2, q3 = self.value
return torch.tensor([[1 - 2*q2*q2 - 2*q3*q3, 2*q1*q2 + 2*q0*q3, 2*q1*q3 - 2*q0*q2],
[2*q1*q2 - 2*q0*q3, 1-2*q1*q1 - 2*q3*q3, 2*q2*q3 + 2*q0*q1],
[2*q1*q3 + 2*q0*q2, 2*q2*q3 - 2*q0*q1, 1-2*q1*q1-2*q2*q2]])
def getValue(self):
return self.value
@staticmethod
def from_tensor(array): # 从tensor转成四元数类
array = array.squeeze()
return Quaternion(array[0], array[1], array[2], array[3])
@staticmethod
def get_rotation_quaternion(axis, angle): # 获得旋转四元数。0:x轴 1:y轴 2:z轴, 旋转角度
angle = torch.tensor(angle, dtype=torch.float32) * (torch.pi / 180)
R_vector = None
if axis == 0:
R_vector = torch.tensor([1, 0, 0])
elif axis == 1:
R_vector = torch.tensor([0, 1, 0])
elif axis == 2:
R_vector = torch.tensor([0, 0, 1])
return Quaternion.from_tensor(torch.cat((torch.cos(angle / 2), torch.sin(angle / 2) * R_vector)))
数据预处理 Dataset
Copy import random
from torch.utils.data import dataset
import open3d as o3d
import os
import numpy as np
import Quaternion as quat
import torch
def random_unit_quaternion(transform_range):
# 随机生成旋转角（0 到 180 度），并将其转换为弧度
angle = torch.rand(1) * torch.pi * 2
# 计算实部
w = torch.cos(angle / 2)
# 随机生成虚部
x, y, z = 2 * torch.rand(3) - 1
v_norm = torch.sqrt(x ** 2 + y ** 2 + z ** 2)
# 单位化虚部向量
x, y, z = x / v_norm, y / v_norm, z / v_norm
# 计算旋转轴的缩放系数
s = torch.sin(angle / 2)
# 根据实部和虚部组合成单位四元数
q = torch.tensor([w, x * s, y * s, z * s])
if w < 0:
q = -q # 始终保持第一个数是正数
# 将四元数转换为旋转矩阵
quat_obj = quat.Quaternion.from_tensor(q)
R = quat_obj.getRotationMatrix()
t = 2 * transform_range * torch.rand(3) - transform_range
return quat_obj, R, t
class LoadData(dataset.Dataset):
def __init__(self, points_count, std_model_paths, t_range,label_is_rotation):
super(LoadData, self).__init__()
folder_path = std_model_paths
self.file_paths = []
self.points_count = points_count
self.count = 0
self.pcds_array = []
self.label_is_rotation = label_is_rotation
self.t_range = t_range
for root, dirs, files in os.walk(folder_path):
for file in files:
if file == '.DS_Store':
continue
self.count = self.count + 1
self.pcds_array.append(torch.tensor(np.asarray(o3d.io.read_point_cloud(os.path.join(root, file)).points), dtype=torch.float32))
print(os.path.join(root, file))
def __getitem__(self, index):
with torch.no_grad():
pcd_std = self.pcds_array[index % self.count] # 标准点云
random_index = random.randint(0, self.count-1) # 生成一个0到self.count之间的随机整数
pcd_aligned = self.pcds_array[random_index] # 被对其点云
q, R, t = random_unit_quaternion(self.t_range) # 获取随机旋转四元数
sampled_indices_std = torch.randperm(pcd_std.shape[0])[:self.points_count]
pcd_std = pcd_std[sampled_indices_std] # 从模型中随机采样n个点
sampled_indices_aligned = torch.randperm(pcd_aligned.shape[0])[:self.points_count]
pcd_aligned = pcd_aligned[sampled_indices_aligned] # 从模型中随机采样n个点
points_after = torch.matmul(pcd_aligned - t, R.T) # 将采样到的点运用 随机旋转四元数
if self.label_is_rotation:
label = q.getValue()
else:
label = t
return points_after, pcd_std, label # 返回随机旋转后的点云 以及对应的四元数
def __len__(self):
return 500 # self.count
随机生成单位四元数（random_unit_quaternion）： 这个函数随机生成一个单位四元数表示的旋转，并将其转换为旋转矩阵。同时，生成一个随机平移向量。这个函数的主要作用是为点云数据应用随机的旋转和平移。
点云数据集类（LoadData）： 这个类继承自torch.utils.data.Dataset，用于加载点云数据并进行预处理。主要功能如下：
在__init__方法中，读取指定文件夹下的所有点云文件，并将其转换为Tensor格式存储。
在__getitem__方法中，对点云数据进行以下处理： a. 随机选择两个点云（标准点云和被对齐点云）。 b. 随机生成旋转四元数和平移向量。 c. 对点云数据进行随机采样，选取指定数量的点。 d. 将被对齐点云应用随机旋转和平移，得到旋转后的点云。 e. 返回旋转后的点云、标准点云和对应的四元数或平移向量（根据label_is_rotation参数决定）。
在__len__方法中，返回数据集的长度，这里设置为固定值500。
定义网络 Net
Copy import torch
import torch.nn as nn
import numpy as np
class PointNet(nn.Module):
def __init__(self):
super(PointNet, self).__init__()
self.conv1 = torch.nn.Conv1d(3, 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)
def forward(self, x):
x = torch.relu(self.bn1(self.conv1(x))) # 1D卷积层 3 -> 64、归一化、Relu
x = torch.relu(self.bn2(self.conv2(x))) # 1D卷积层 64->128、归一化、Relu
x = self.bn3(self.conv3(x)) # 1D卷积层 128->1024、归一化
x = torch.max(x, 2, keepdim=True)[0] # 最大池化 （batch_size, 1024）
x = x.view(-1, 1024)
return x
class RotationNet(nn.Module):
def __init__(self):
super(RotationNet, self).__init__()
self.pointnet = PointNet()
self.fc1 = nn.Linear(2048, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 4)
self.bn1 = nn.BatchNorm1d(128)
self.bn2 = nn.BatchNorm1d(64)
def forward(self, x1, x2):
# 计算x1和x2的质心
centroid_x1 = torch.mean(x1, dim=2, keepdim=True) # (32, 3, 1)
centroid_x2 = torch.mean(x2, dim=2, keepdim=True) # (32, 3, 1)
# 将x1和x2归一化到质心位置
x1 = x1 - centroid_x1
x2 = x2 - centroid_x2
# 特征提取
x1 = self.pointnet(x1)
x2 = self.pointnet(x2)
x = torch.cat((x1, x2), dim=1)
x = torch.relu(self.bn1(self.fc1(x)))
x = torch.relu(self.bn2(self.fc2(x)))
x = self.fc3(x).float()
x = torch.nn.functional.normalize(x, p=2, dim=1)
# 确保四元数第一个分量为正值
mask = x[:, 0] < 0
x[mask] = -x[mask]
return x
class TranslationNet(nn.Module):
def __init__(self):
super(TranslationNet, self).__init__()
self.pointnet = PointNet()
self.fc1 = nn.Linear(2048, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 3)
self.bn1 = nn.BatchNorm1d(128)
self.bn2 = nn.BatchNorm1d(64)
def forward(self, x1, x2):
x1 = self.pointnet(x1)
x2 = self.pointnet(x2)
x = torch.cat((x1, x2), dim=1)
x = torch.relu(self.bn1(self.fc1(x)))
x = torch.relu(self.bn2(self.fc2(x)))
x = self.fc3(x).float()
return x PointNet（点云特征提取网络）:
在forward方法中，输入的点云数据依次经过三个1D卷积层和批归一化层，然后通过ReLU激活函数。
最后，对结果进行最大池化操作，得到大小为(batch_size, 1024)的特征表示。
RotationNet（旋转预测网络）：
包含一个PointNet实例用于特征提取，以及三个全连接层和相应的批归一化层。
在forward方法中，首先计算输入点云的质心，然后将点云归一化到质心位置。
使用PointNet提取点云特征，然后将两个点云的特征沿维度1拼接。
拼接后的特征依次通过全连接层、批归一化层和ReLU激活函数。
最后，经过一层全连接层并进行L2范数归一化，得到大小为(batch_size, 4)的四元数表示。
TranslationNet（平移预测网络）：
包含一个PointNet实例用于特征提取，以及三个全连接层和相应的批归一化层。
在forward方法中，使用PointNet提取点云特征，然后将两个点云的特征沿维度1拼接。
拼接后的特征依次通过全连接层、批归一化层和ReLU激活函数。
最后，经过一层全连接层，得到大小为(batch_size, 3)的平移向量表示。
四元数损失定义 QuatLoss
Copy import torch
import torch.nn as nn
import numpy as np
class QuatLoss(nn.Module):
def __init__(self):
super(QuatLoss, self).__init__()
def forward(self, y_hat, label):
w, x, y, z = label.split(1, dim=-1)
label_inv = torch.cat([w, -x, -y, -z], dim=-1) # label共轭
sa, xa, ya, za = y_hat.split(1, dim=-1)
sb, xb, yb, zb = label_inv.split(1, dim=-1)
loss_quat = torch.cat([
sa * sb - xa * xb - ya * yb - za * zb,
sa * xb + xa * sb + ya * zb - za * yb,
sa * yb - xa * zb + ya * sb + za * xb,
sa * zb + xa * yb - ya * xb + za * sb
], dim=-1)
# loss = torch.pow(torch.acos(loss_quat[:, 0]) * 2, 2)
cos_angle = loss_quat[:, 0] # 四元数点积（它们之间的夹角余弦值）
loss = 1 - torch.square(cos_angle) # 计算 1 减去余弦值的平方
return torch.mean(loss) 这段代码定义了一个四元数损失函数（QuatLoss），它继承自torch.nn.Module。这个损失函数的目的是计算两个四元数之间的角度差异，并使用这个差异作为预测旋转的损失。
在forward方法中：
将标签（label）四元数分解为实部（w）和虚部（x, y, z）。
将预测四元数（y_hat）分解为实部（sa）和虚部（xa, ya, za）。
将标签共轭四元数（label_inv）分解为实部（sb）和虚部（xb, yb, zb）。
提取乘积四元数中的实部（cos_angle），它表示输入四元数之间的夹角的余弦值。
这个损失函数的作用是衡量预测四元数（y_hat）与标签四元数（label）之间的角度差异。当预测四元数与标签四元数之间的角度相差较大时，损失值较大；当它们之间的角度相差较小时，损失值较小。这有助于神经网络在训练过程中学习到正确的旋转信息。
训练旋转四元数 TrainQuaternion
Copy import argparse
import time
import torch
from torch.optim.lr_scheduler import StepLR
import DataLoader
import Net
import QuatLoss
import TransformLoss
import torch
from torch.utils.data import dataloader
import os
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
parser = argparse.ArgumentParser()
parser.add_argument('--batchSize', type=int, default=32, help='input batch size')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=0)
parser.add_argument('--nepoch', type=int, default=10000, help='number of epochs to train for')
parser.add_argument('--outf', type=str, default='seg', help='output folder')
parser.add_argument('--lr', type=float, default=0.1, help='learning rate')
parser.add_argument('--dataset', type=str, default='StandardModel', help="dataset path")
parser.add_argument('--npoints', type=int, default=500, help="sample point counts")
parser.add_argument('--stdModelPath', type=str, default='StandardModel', help="standard model path")
opt = parser.parse_args()
# 保存和输出设置
save_epoch = 100 # 几个epoch保存一次
println_epoch = 10 # 几个epoch输出一次loss
# 实例化网络
net = Net.RotationNet()
net.to(device)
# net = torch.load('./Save/1679852443_SGD_QuatLoss_batchsize32_/1679852443_epoch_4700_loss_0.01935810090508312.module')
# 优化器
optimizer = torch.optim.SGD(params=net.parameters(), lr=opt.lr, momentum=0.3)
optimizer_name = 'SGD'
# 损失函数
criterion = QuatLoss.QuatLoss()
criterion_name = 'QuatLoss'
# 学习率衰减
scheduler = StepLR(optimizer, step_size=1000, gamma=0.95) # 学习率衰减
# DataLoader
train_dataset = DataLoader.LoadData(points_count=opt.npoints,
std_model_paths=opt.stdModelPath,
t_range=0.2,
label_is_rotation=True)
# t_range 是 随机出的平移矩阵的范围， label_is_rotation 指数据的label是旋转四元数还是平移向量
train_loader = dataloader.DataLoader(dataset=train_dataset, batch_size=opt.batchSize, shuffle=True, num_workers=opt.workers)
# 模型保存路径
name = str(int(time.time()))
save_path = os.path.join('RotationNetSave', name + '_' + optimizer_name + '_' + criterion_name + '_batchsize' + str(opt.batchSize) + '_lr_' + str(opt.lr))
if not os.path.exists(save_path):
os.makedirs(save_path)
# 记录损失值
loss_sum = 0
loss_count = 0
# 开始训练
for epoch in range(opt.nepoch):
count = 0
for points_align, points_std, label in train_loader:
count += 1
points_align = points_align.float().to(device)
points_std = points_std.float().to(device)
label = label.float().to(device)
x1 = points_align.float().transpose(1, 2)
x2 = points_std.float().transpose(1, 2)
optimizer.zero_grad()
y_hat = net(x1, x2).float()
loss = criterion(y_hat, label)
loss_sum += loss.item()
loss_count += 1
loss.backward()
optimizer.step()
scheduler.step() # 更新学习率
if epoch % save_epoch == 0:
torch.save(net, save_path + "/" + name + '_epoch_' + str(epoch) + '_loss_' + str(loss_sum / loss_count) +'.module')
if epoch % println_epoch == 0:
print("epoch ", epoch, " loss ", loss_sum / loss_count, " lr ", get_lr(optimizer))
loss_sum = 0
loss_count = 0
torch.save(net, save_path + "/" + name + '.module')
创建并配置平移网络模型、优化器、损失函数和学习率衰减策略。
加载预训练的RotationNet模型，用于将点云旋转到正确的位姿。
进行训练循环，每个epoch中：
使用RotationNet网络调整输入点云的位姿。
训练平移向量 TrainTranslateVector
Copy import argparse
import time
import torch
from torch.optim.lr_scheduler import StepLR
import DataLoader
import Net
import QuatLoss
import torch
from torch.utils.data import dataloader
import os
import Quaternion as quat
import TransformLoss
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
parser = argparse.ArgumentParser()
parser.add_argument('--batchSize', type=int, default=32, help='input batch size')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=0)
parser.add_argument('--nepoch', type=int, default=10000, help='number of epochs to train for')
parser.add_argument('--outf', type=str, default='seg', help='output folder')
parser.add_argument('--lr', type=float, default=0.1, help='learning rate')
parser.add_argument('--dataset', type=str, default='StandardModel', help="dataset path")
parser.add_argument('--npoints', type=int, default=500, help="sample point counts")
parser.add_argument('--stdModelPath', type=str, default='StandardModel', help="standard model path")
opt = parser.parse_args()
# 保存和输出设置
save_epoch = 100 # 几个epoch保存一次
println_epoch = 10 # 几个epoch输出一次loss
# 实例化平移网络
net = Net.TranslationNet()
net.to(device)
# 加载旋转网络
rotation_net = torch.load('./RotationNetSave/1680457393_SGD_QuatLoss_batchsize32_lr_0.1/1680457393_epoch_9900_loss_0.026239702384918927.module',
map_location=torch.device(device))
# 优化器
optimizer = torch.optim.SGD(params=net.parameters(), lr=opt.lr, momentum=0.3)
optimizer_name = 'SGD'
# 损失函数
# criterion = torch.nn.MSELoss()
# criterion_name = 'MSELoss'
criterion = TransformLoss.TransformLoss()
criterion_name = 'TransformLoss'
# 学习率衰减
scheduler = StepLR(optimizer, step_size=5000, gamma=0.95) # 学习率衰减
# DataLoader
train_dataset = DataLoader.LoadData(points_count=opt.npoints,
std_model_paths=opt.stdModelPath,
t_range=0.2,
label_is_rotation=False)
# t_range 是 随机出的平移矩阵的范围， label_is_rotation 指数据的label是旋转四元数还是平移向量
train_loader = dataloader.DataLoader(dataset=train_dataset,
batch_size=opt.batchSize, shuffle=True,
num_workers=opt.workers)
# 模型保存路径
name = str(int(time.time()))
save_path = os.path.join('TranslationNetSave', name + '_' + optimizer_name + '_' + criterion_name + '_batchsize' + str(opt.batchSize) + '_lr_' + str(opt.lr))
if not os.path.exists(save_path):
os.makedirs(save_path)
# 记录损失值
loss_sum = 0
loss_count = 0
# 开始训练
for epoch in range(opt.nepoch):
count = 0
for points_align, points_std, label in train_loader:
count += 1
points_align = points_align.float().to(device)
points_std = points_std.float().to(device)
label = label.float().to(device)
x1 = points_align.float().transpose(1, 2)
x2 = points_std.float().transpose(1, 2)
# 将x1和x2用RotationNet进行位姿调整
r = rotation_net(x1, x2).float() # 得到四元数
batch_size = r.shape[0] # 获取batchsize
for i in range(batch_size):
quat_obj = quat.Quaternion.from_tensor(r[i]) # 转四元数类
R = quat_obj.getRotationMatrix().to(device) # 得到R旋转矩阵
x1[i] = torch.matmul(x1[i].transpose(0, 1), R).transpose(0, 1) # 每一个点云都进行旋转
optimizer.zero_grad()
y_hat = net(x1, x2).float()
loss = criterion(x1, x2, y_hat)
loss_sum += loss.item()
loss_count += 1
loss.backward()
optimizer.step()
scheduler.step() # 更新学习率
if epoch % save_epoch == 0:
torch.save(net, save_path + "/" + name + '_epoch_' + str(epoch) + '_loss_' + str(loss_sum / loss_count) +'.module')
if epoch % println_epoch == 0:
print("epoch ", epoch, " loss ", loss_sum / loss_count, " lr ", get_lr(optimizer))
loss_sum = 0
loss_count = 0
torch.save(net, save_path + "/" + name + '.module')
导入所需库和模块，包括PyTorch、自定义的DataLoader、Net等。
定义get_lr()函数，用于获取优化器的当前学习率。
定义命令行参数，包括批次大小、工作进程数、训练周期数、输出文件夹等。
创建TranslationNet实例并将其加载到相应设备上。
测试代码 Test
Copy import open3d as o3d
import numpy as np
import torch
from torch.utils.data import dataloader
import DataLoader
import QuatLoss
import Quaternion as quat
import ICP
number_of_points = 500
rotation_net = torch.load('./RotationNetSave/1680457393_SGD_QuatLoss_batchsize32_lr_0.1/1680457393_epoch_9900_loss_0.026239702384918927.module',
map_location=torch.device('cpu'))
transform_net = torch.load('./TranslationNetSave/1680461748_SGD_MSELoss_batchsize32_lr_0.1/1680461748_epoch_13000_loss_0.0005985869252072007.module',
map_location=torch.device('cpu'))
test_dataset = DataLoader.LoadData(points_count=500, std_model_paths='StandardModel', t_range=0.2, label_is_rotation=False)
test_loader = dataloader.DataLoader(dataset=test_dataset, batch_size=32, shuffle=False)
criterion = QuatLoss.QuatLoss()
# 是否开启ICP的细配准
use_icp = True
for points_align, points_std, label in test_loader:
x1 = points_align.float().transpose(1, 2)
x2 = points_std.float().transpose(1, 2)
y_hat_rotation = rotation_net(x1, x2).float()
# loss = criterion(y_hat_rotation, label)
x1_single = x1[0].transpose(0, 1)
x2_single = x2[0].transpose(0, 1)
# 输入 需对其点云
pcd_before = o3d.geometry.PointCloud()
pcd_before.points = o3d.utility.Vector3dVector(x1_single)
pcd_before.paint_uniform_color([0, 1, 0]) # 指定显示为绿色
# 将输入点云进行预测值的旋转
r = rotation_net(x1, x2)
batch_size = r.shape[0]
for i in range(batch_size):
quat_obj = quat.Quaternion.from_tensor(r[i])
R = quat_obj.getRotationMatrix()
x1[i] = torch.matmul(x1[i].transpose(0, 1), R).transpose(0, 1)
# 进行平移
y_hat_transform = transform_net(x1, x2).float()
x1 = x1.transpose(1, 2)
x1 = x1 + y_hat_transform.unsqueeze(1)
x2 = x2.transpose(1, 2)
points_after = x1[0].detach()
if use_icp:
# 用ICP算法进行细调整
R, t = ICP.ICP(x1[0].detach().numpy(), x2[0].detach().numpy(), 80)
points_after_ICP = torch.matmul(x1[0], torch.from_numpy(R).float())
t = torch.from_numpy(t.reshape(1, 3)).float()
points_after_ICP += t
# 显示旋转后的模型
pcd_after = o3d.geometry.PointCloud()
pcd_after.points = o3d.utility.Vector3dVector(points_after.detach())
pcd_after.paint_uniform_color([1, 0, 0]) # 红色
# 显示std模型
pcd_std = o3d.geometry.PointCloud()
pcd_std.points = o3d.utility.Vector3dVector(x2_single)
pcd_std.paint_uniform_color([0, 0, 0]) # 黑色
if use_icp:
pcd_after.points = o3d.utility.Vector3dVector(points_after_ICP.detach())
o3d.visualization.draw_geometries([pcd_before, pcd_after, pcd_std]) # 显示对其后位置
else:
o3d.visualization.draw_geometries([pcd_before, pcd_after, pcd_std]) # 显示对其后位置
ython 这段代码实现了使用训练好的RotationNet和TranslationNet进行点云对齐。主要步骤如下：
导入所需库，包括Open3D、NumPy、PyTorch等。
加载训练好的RotationNet和TranslationNet模型。
设置是否使用ICP算法（迭代最近点算法）进行细对齐。
遍历测试数据集的批次，执行以下操作：
将输入点云和标准点云传入RotationNet，得到预测的旋转四元数。
将旋转后的点云和标准点云传入TranslationNet，得到预测的平移向量。
如果使用ICP算法，则对旋转和平移后的点云进行细对齐。
在这个脚本中，使用Open3D库进行点云的可视化和处理。通过预测的旋转四元数和平移向量，对输入点云进行旋转和平移，以实现与标准点云的对齐。如果启用了ICP算法，还会对旋转和平移后的点云进行细对齐。最后，将输入点云、对齐后的点云和标准点云可视化并显示。