diff --git a/PyTorch Model/model_functions.py b/PyTorch Model/model_functions.py
index 9cca817c898477221322bb93a94b44dda4273cbb..e8a8af6a9988aeba91d4f3753fd0f12aa9d25c97 100644
--- a/PyTorch Model/model_functions.py
+++ b/PyTorch Model/model_functions.py
@@ -12,42 +12,70 @@ from torch.optim import lr_scheduler
from data_loader import get_dataloader
import models
-from loss import GeoPoseLoss
+from losses import GeoPoseLoss
from pose_utils import *
+def save_checkpoint(epoch, model, optimizer, criterion): # for sx and sq
+ filename = os.path.join('posenet.pth.tar'.format(epoch))
+ checkpoint_dict = \
+ {'epoch': epoch, 'model_state_dict': model.state_dict(),
+ 'optim_state_dict': optimizer.state_dict(),
+ 'criterion_state_dict': criterion.state_dict()}
+ torch.save(checkpoint_dict, filename)
+
+
+def load_checkpoint(model, criterion, optimizer, filename):
+ if os.path.isfile(filename):
+ checkpoint = torch.load(filename)
+ model.load_state_dict(checkpoint['model_state_dict'])
+ optimizer.load_state_dict(checkpoint['optim_state_dict'])
+ criterion.load_state_dict(checkpoint['criterion_state_dict'])
+ epoch = checkpoint['epoch']
+ print("Checkpoint loaded from epoch: ", epoch)
+ return epoch
+ else:
+ print("No checkpoint found at", filename)
+ return 0 # for start epoch
+
+
def train_model(config):
cudnn.benchmark = True
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using device: ", device)
- # TODO set seed
- # TODO print out config
+ torch.manual_seed(7)
- # DATA
+ ### DATA
# dataset_path = './datasets/KingsCollege'
data_name = config.dataset_path.split('/')[-1]
model_save_path = 'models_%s' % data_name
tb_save_path = 'runs_%s' % data_name
+ print(config)
+
train_loader = get_dataloader(dataset_path=config.dataset_path, mode='train',
model=config.model, batch_size=config.batch_size)
val_loader = get_dataloader(dataset_path=config.dataset_path, mode='val',
model=config.model, batch_size=config.batch_size)
+ print(f"Loading data from: {config.dataset_path}")
+ print(f"No. of Training samples: {len(train_loader)*config.batch_size}; {len(train_loader)} batches of {config.batch_size}")
+ print(f"No. of Validation samples: {len(val_loader)*config.batch_size}; {len(val_loader)} batches of {config.batch_size}")
- # MODEL
+ ### MODEL
if config.model == 'googlenet':
model = models.GoogleNet(fixed_weight=config.fixed_weight, dropout_rate=config.dropout_rate)
else:
- model = models.ResNet(fixed_weight=config.fixed_weight, dropout_rate=config.dropout_rate)
+ # model = models.ResNet(fixed_weight=config.fixed_weight, dropout_rate=config.dropout_rate)
+ model = models.PoseNet()
model.to(device)
- if config.pretrained_model: # TODO either none or the path and epoch to start from
- model_path = '/%s_net.pth' % config.pretrained_model
+ if config.pretrained_model:
+ model_path = config.pretrained_model
model.load_state_dict(torch.load(model_path))
- print('Load pretrained network: ', model_path)
+ print('Loading pretrained network from: ', model_path)
- # LOSS AND OPTIMISER
- criterion = GeoPoseLoss(config.learn_beta)
+ ### LOSS AND OPTIMISER
+ criterion = GeoPoseLoss(learn_beta=config.learn_beta)
criterion.to(device)
if config.learn_beta:
@@ -63,7 +91,7 @@ def train_model(config):
# LR is decayed by the value of gamma
scheduler = lr_scheduler.StepLR(optimizer, step_size=config.num_epochs_decay, gamma=0.1)
- # TRAINING
+ ### TRAINING
num_epochs = config.num_epochs
if not os.path.exists(model_save_path):
@@ -72,13 +100,23 @@ def train_model(config):
# tensorboard
if not os.path.exists(tb_save_path):
os.makedirs(tb_save_path)
+ with open(tb_save_path + '/config.txt', mode="w+") as f:
+ f.write(str(config))
writer = SummaryWriter(log_dir=tb_save_path)
start_time = time.time()
n_iter = 0 # total no of batches over all epochs, useful for logger
+ n_val_iter = 0
+ if config.pretrained_model:
+ n_iter = int(config.pretrained_epoch) * len(train_loader) # provided batch size is same
+ n_val_iter = int(config.pretrained_epoch) * len(val_loader)
start_epoch = 0
if config.pretrained_model:
- start_epoch = int(config.pretrained_model)
+ start_epoch = int(config.pretrained_epoch)
+
+ # IF CHECKPOINT PRESENT, LOAD: - THIS IS ONLY FOR CONDOR
+ # start_epoch = load_checkpoint(model, criterion, optimizer, 'posenet.pth.tar')
+ # n_iter = int(start_epoch) * len(train_loader)
for epoch in range(start_epoch, num_epochs):
print('Epoch {}/{}'.format(epoch + 1, num_epochs))
@@ -109,13 +147,12 @@ def train_model(config):
# print("ACT POSE", pos_true, "ACT ORI", ori_true)
ori_out = F.normalize(ori_out, p=2, dim=1)
- ori_true = F.normalize(ori_true, p=2, dim=1)
+ ori_true = F.normalize(ori_true, p=2, dim=1) # doesn't make any difference
loss, loss_pos_print, loss_ori_print = criterion(pos_out, ori_out, pos_true, ori_true)
# print("LOSS", loss.item())
loss_print = loss.item()
-
- # loss = loss_pos + beta * loss_ori
+ del inputs, poses
# TODO use config.log_step
error_train.append(loss_print)
@@ -160,26 +197,32 @@ def train_model(config):
loss, loss_pos_print, loss_ori_print = criterion(pos_out, ori_out, pos_true, ori_true)
loss_print = loss.item()
+ del inputs, poses
error_val.append(loss_print)
- # TODO log val loss
+ writer.add_scalar('loss/overall_val_loss', loss_print, n_val_iter)
+ n_val_iter += 1
+
print('batch#{} val Loss: total loss {:.3f} / pos loss {:.3f} / ori loss {:.3f}'.format(i, loss_print,
loss_pos_print,
loss_ori_print))
# END OF EPOCH
- error_train_loss = np.median(error_train)
- error_val_loss = np.median(error_val)
+ error_train_avg = np.mean(error_train)
+ error_val_avg = np.mean(error_val)
- print('Overall Train and Validation loss for epoch {} / {}'.format(error_train_loss, error_val_loss))
+ print('Overall Train and Validation loss for epoch {} / {}'.format(error_train_avg, error_val_avg))
print('=' * 40)
print('=' * 40)
- writer.add_scalars('loss/trainval', {'train': error_train_loss, 'val': error_val_loss}, epoch+1)
+ writer.add_scalars('loss/trainval', {'train': np.median(error_train_avg), 'val': np.median(error_val_avg)}, epoch+1)
if (epoch + 1) % config.model_save_step == 0:
- # save_filename = model_save_path + '/%s_net.pth' % epoch
- save_filename = model_save_path + '/posenet.pth'
+ # save_checkpoint(epoch=epoch, model=model.cpu(), optimizer=optimizer, criterion=criterion.cpu())
+ # if torch.cuda.is_available():
+ # model.to(device)
+ # criterion.to(device)
+ save_filename = model_save_path + '/posenet_{}.pth'.format(epoch+1)
torch.save(model.cpu().state_dict(), save_filename)
if torch.cuda.is_available():
model.to(device)
@@ -210,19 +253,27 @@ def test_model(config):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Using device: ", device)
- # LOAD THE MODEL
+ # DATA
data_name = config.dataset_path.split('/')[-1]
model_save_path = 'models_%s' % data_name
tb_save_path = 'runs_%s' % data_name
- # f = open(tb_save_path + '/test_result.csv', 'w+')
+ f = open(tb_save_path + '/test_result.csv', 'w+')
+
+ test_loader = get_dataloader(dataset_path=config.dataset_path, mode='test') # BATCH SIZE IS 1
+ print(f"Loading data from: {config.dataset_path}")
+ print(f"No. of Test samples: {len(test_loader)}")
+ # LOAD THE MODEL
if config.pretrained_model:
model_save_path = config.pretrained_model
+ else:
+ print("No model specified") # TODO load best
if config.model == 'googlenet':
- model = models.GoogleNet(fixed_weight=config.fixed_weight, dropout_rate=config.dropout_rate)
+ model = models.GoogleNet()
else:
- model = models.ResNet(fixed_weight=config.fixed_weight, dropout_rate=config.dropout_rate)
+ # model = models.ResNet()
+ model = models.PoseNet()
model.to(device)
print('Loading pretrained model: ', model_save_path)
@@ -230,17 +281,16 @@ def test_model(config):
model.eval()
+ # arrays to store individual losses
+ pos_loss = []
+ ori_loss = []
total_pos_loss = 0
total_ori_loss = 0
- # arrays to store individual losses
- pos_loss_arr = []
- ori_loss_arr = []
+ # for plotting later
true_pose_list = []
estim_pose_list = []
- test_loader = get_dataloader(dataset_path=config.dataset_path, mode='test') # BATCH SIZE IS 1
for i, (inputs, poses) in enumerate(test_loader):
- print(i)
inputs = inputs.to(device)
pos_out, ori_out = model(inputs)
@@ -249,49 +299,53 @@ def test_model(config):
# Then, detach the tensor from the computation graph using detach()
pos_out = pos_out.squeeze(0).detach().cpu().numpy()
ori_out = F.normalize(ori_out, p=2, dim=1)
- ori_out = quat_to_euler(ori_out.squeeze(0).detach().cpu().numpy())
+ # ori_out = quat_to_euler(ori_out.squeeze(0).detach().cpu().numpy())
+ ori_out = ori_out.squeeze(0).detach().cpu().numpy()
print('pos out', pos_out)
print('ori_out', ori_out)
pos_true = poses[:, :3].squeeze(0).numpy()
ori_true = poses[:, 3:].squeeze(0).numpy()
-
- ori_true = quat_to_euler(ori_true)
+ # ori_true = quat_to_euler(ori_true)
print('pos true', pos_true)
print('ori true', ori_true)
- # l2 norm
- loss_pos_print = array_dist(pos_out, pos_true)
- loss_ori_print = array_dist(ori_out, ori_true)
- true_pose_list.append(np.hstack((pos_true, ori_true)))
- estim_pose_list.append(np.hstack((pos_out, ori_out)))
+ # l2 distance
+ loss_pos_print = position_dist(pos_out, pos_true)
+ loss_ori_print = rotation_dist(ori_out, ori_true)
+
+ pos_loss.append(loss_pos_print)
+ ori_loss.append(loss_ori_print)
total_pos_loss += loss_pos_print
total_ori_loss += loss_ori_print
- pos_loss_arr.append(loss_pos_print)
- ori_loss_arr.append(loss_ori_print)
- print('{}th Error: pos error {:.3f} / ori error {:.3f}'.format(i, loss_pos_print, loss_ori_print))
+ true_pose_list.append(np.hstack((pos_true, ori_true)))
+ estim_pose_list.append(np.hstack((pos_out, ori_out)))
- position_error = np.median(pos_loss_arr)
- rotation_error = np.median(ori_loss_arr)
+ print('batch#{} error: pos error {:.3f} / ori error {:.3f}'.format(i, loss_pos_print, loss_ori_print))
+ f.write('batch#{} error: pos error {:.3f} / ori error {:.3f} \n'.format(i, loss_pos_print, loss_ori_print))
+ # END OF EPOCH
print('=' * 20)
- print('Overall median pose error {:.3f}m / {:.3f}rad'.format(position_error, rotation_error))
- print('Overall average pose error {:.3f}m / {:.3f}rad'.format(np.mean(pos_loss_arr), np.mean(ori_loss_arr)))
- # f.close()
+ print('Overall median pose error {:.3f}m / {:.3f}o'.format(np.median(pos_loss), np.median(ori_loss)))
+ f.write('Overall median pose error {:.3f}m / {:.3f}o \n'.format(np.median(pos_loss), np.median(ori_loss)))
+ print('Overall average pose error {:.3f}m / {:.3f}o'.format(np.mean(pos_loss), np.mean(ori_loss)))
+ f.write('Overall average pose error {:.3f}m / {:.3f}o \n'.format(np.mean(pos_loss), np.mean(ori_loss)))
f_true = tb_save_path + '/pose_true.csv'
f_estim = tb_save_path + '/pose_estim.csv'
np.savetxt(f_true, true_pose_list, delimiter=',')
np.savetxt(f_estim, estim_pose_list, delimiter=',')
+ f.close()
+
if __name__ == '__main__':
# train_model(save_name='KingsCollege-ResNet')
# test_model(save_name='KingsCollege-ResNet')
# train_model(save_name='KingsCollege-GNet')
- train_model(save_name='KingsCollege-ResNet-V')
+ # train_model(save_name='KingsCollege-ResNet-V')
# test_model(save_name='KingsCollege-ResNet-V')
-
+ pass
diff --git a/PyTorch Model/models.py b/PyTorch Model/models.py
index f324adecdec1ad605d3dbd6aa9c91544eea3e71e..7f8bc4919b3d90d8f2e5cd33a650100f5007985a 100644
--- a/PyTorch Model/models.py
+++ b/PyTorch Model/models.py
@@ -6,6 +6,48 @@ from torchvision import models
from torchsummary import summary
+# https://github.com/NVlabs/geomapnet/blob/master/models/posenet.py
+class PoseNet(nn.Module):
+ def __init__(self, droprate=0.5, pretrained=True):
+ super(PoseNet, self).__init__()
+ self.droprate = droprate
+
+ # replace the last FC layer in feature extractor
+ self.feature_extractor = models.resnet34(weights="IMAGENET1K_V1")
+ self.feature_extractor.avgpool = nn.AdaptiveAvgPool2d(1)
+ in_dim = self.feature_extractor.fc.in_features
+ self.feature_extractor.fc = nn.Linear(in_dim, 2048)
+
+ self.fc_xyz = nn.Linear(2048, 3)
+ self.fc_wpqr = nn.Linear(2048, 4)
+
+ # initialize
+ if pretrained:
+ init_modules = [self.feature_extractor.fc, self.fc_xyz, self.fc_wpqr]
+ else:
+ init_modules = self.modules()
+
+ for m in init_modules:
+ if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+ nn.init.kaiming_normal_(m.weight.data)
+ if m.bias is not None:
+ nn.init.constant_(m.bias.data, 0)
+
+ def forward(self, x):
+ x = self.feature_extractor(x)
+ # 1 x 2048
+ x = F.relu(x)
+ if self.droprate > 0:
+ x = F.dropout(x, p=self.droprate)
+
+ xyz = self.fc_xyz(x)
+ wpqr = self.fc_wpqr(x)
+ return xyz, wpqr
+ # seems to be no random crop, no colour jitter;
+ # log of quaternion
+ # 300 epochs
+
+
class ResNet(nn.Module):
def __init__(self, fixed_weight=False, dropout_rate=0.0):
super(ResNet, self).__init__()
@@ -38,7 +80,7 @@ class ResNet(nn.Module):
def forward(self, x):
x = self.base_model(x)
# print("AFTER RESNET", x.shape)
- x = x.view(x.size(0), -1) # flatten
+ x = x.view(x.size(0), -1) # flatten TODO global average pool
# print("AFTER VIEW", x.shape)
x = self.fc_last(x)
x = F.relu(x)
@@ -111,127 +153,12 @@ class GoogleNet(nn.Module):
return pos, ori
-class NetVLAD(nn.Module):
- """NetVLAD layer implementation"""
-
- def __init__(self, num_clusters=64, dim=128, alpha=100.0,
- normalize_input=True):
- """
- Args:
- num_clusters : int
- The number of clusters
- dim : int
- Dimension of descriptors
- alpha : float
- Parameter of initialization. Larger value is harder assignment.
- normalize_input : bool
- If true, descriptor-wise L2 normalization is applied to input.
- """
- super(NetVLAD, self).__init__()
- self.num_clusters = num_clusters
- self.dim = dim
- self.alpha = alpha
- self.normalize_input = normalize_input
- self.conv = nn.Conv2d(dim, num_clusters, kernel_size=(1, 1), bias=True)
- self.centroids = nn.Parameter(torch.rand(num_clusters, dim)) # cluster centres
- self._init_params()
-
- def _init_params(self):
- # manually initialising rather than PyTorch assigning random values to weights and 0s to bias
- self.conv.weight = nn.Parameter(
- (2.0 * self.alpha * self.centroids).unsqueeze(-1).unsqueeze(-1)
- )
- self.conv.bias = nn.Parameter(
- - self.alpha * self.centroids.norm(dim=1)
- )
-
- def forward(self, x):
- N, C = x.shape[:2] # N->batch size, C->channels
-
- if self.normalize_input:
- x = F.normalize(x, p=2, dim=1) # across descriptor dim
-
- # soft-assignment
- soft_assign = self.conv(x).view(N, self.num_clusters, -1)
- soft_assign = F.softmax(soft_assign, dim=1)
-
- x_flatten = x.view(N, C, -1)
-
- # calculate residuals to each clusters
- residual = x_flatten.expand(self.num_clusters, -1, -1, -1).permute(1, 0, 2, 3) - \
- self.centroids.expand(x_flatten.size(-1), -1, -1).permute(1, 2, 0).unsqueeze(0)
- residual *= soft_assign.unsqueeze(2)
- vlad = residual.sum(dim=-1)
-
- vlad = F.normalize(vlad, p=2, dim=2) # intra-normalization
- vlad = vlad.view(x.size(0), -1) # flatten
- vlad = F.normalize(vlad, p=2, dim=1) # L2 normalize
-
- return vlad
-
-
-class ResNetVLAD(nn.Module):
- def __init__(self, fixed_weight=False):
- super(ResNetVLAD, self).__init__()
-
- # Discard layers at the end of base network
- encoder = models.resnet34(weights="IMAGENET1K_V1")
- self.base_model = nn.Sequential(
- encoder.conv1,
- encoder.bn1,
- encoder.relu,
- encoder.maxpool,
- encoder.layer1,
- encoder.layer2,
- encoder.layer3,
- encoder.layer4
- )
-
- if fixed_weight:
- for param in self.base_model.parameters():
- param.requires_grad = False
-
- dim = list(self.base_model.parameters())[-1].shape[0] # last channels (512)
- self.net_vlad = NetVLAD(num_clusters=32, dim=dim, alpha=1.0)
-
- self.fc_last = nn.Linear(16384, 2048, bias=True)
- self.fc_position = nn.Linear(2048, 3, bias=True)
- self.fc_rotation = nn.Linear(2048, 4, bias=True)
-
- init_modules = [self.fc_last, self.fc_position, self.fc_rotation]
-
- for module in init_modules:
- if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
- nn.init.kaiming_normal_(module.weight)
- if module.bias is not None:
- nn.init.constant_(module.bias, 0)
-
- def forward(self, x):
- # print("FORWARD PASS STARTING", x.shape)
- x = self.base_model(x)
- # print("AFTER RESNET", x.shape)
- embedded_x = self.net_vlad(x)
- # print("AFTER NETVLAD", embedded_x.shape)
-
- x = self.fc_last(embedded_x)
- x = F.relu(x)
-
- # from that space of 2048, 3 for pos and 4 for ori are sampled
- position = self.fc_position(x)
- rotation = self.fc_rotation(x)
-
- return position, rotation
-
-
if __name__ == '__main__':
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# model = GoogleNet()
- # model.to("cuda:0")
+ # model.to(device)
# summary(model, (3, 299, 299))
- # model2 = ResNet()
- # model2.to("cuda:0")
- # summary(model2, (3, 224, 224))
-
- model3 = ResNetVLAD()
- model3.to("cuda:0")
- summary(model3, (3, 224, 224))
+ model2 = PoseNet()
+ model2.to(device)
+ summary(model2, (3, 224, 224))
diff --git a/PyTorch Model/plot_ori.py b/PyTorch Model/plot_ori.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c7a2e14bc9c2e881fe260167d89d68c11f1a77b
--- /dev/null
+++ b/PyTorch Model/plot_ori.py
@@ -0,0 +1,66 @@
+# Uncomment for 6-plot 2D image
+# import matplotlib
+# matplotlib.use('Agg')
+
+import pandas as pd
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+from scipy.spatial.transform import Rotation
+
+PATH = 'runs_TurtlebotTourCVSSP/run 3 - moar data'
+pose_true = pd.read_csv(PATH + '/pose_true.csv')
+pose_estim = pd.read_csv(PATH + '/pose_estim.csv')
+
+# fig = plt.figure(figsize=(15, 10))
+N = 25 # start from which test sample
+
+# read ori values
+for i in range(6):
+ true_quaternion = pose_true.iloc[N, 3:7].values
+ predicted_quaternion = pose_estim.iloc[N, 3:7].values
+ N = N+1
+
+ # Convert quaternions to rotation matrices (when multiplied with a vector produced the rotated vector
+ true_rotation = Rotation.from_quat(true_quaternion).as_matrix()
+ print(true_rotation)
+ predicted_rotation = Rotation.from_quat(predicted_quaternion).as_matrix()
+
+ # Extract x, y, z axes from rotation matrices
+ true_x_axis = true_rotation[:, 0]
+ true_y_axis = true_rotation[:, 1]
+ true_z_axis = true_rotation[:, 2] # forward-facing or nose direction
+
+ predicted_x_axis = predicted_rotation[:, 0]
+ predicted_y_axis = predicted_rotation[:, 1]
+ predicted_z_axis = predicted_rotation[:, 2]
+
+ # Create a 3D plot
+ fig = plt.figure()
+ ax = fig.add_subplot(111, projection='3d')
+ # ax = fig.add_subplot(2, 3, i+1, projection='3d')
+
+ # Plot the true orientation (red axes)
+ ax.quiver(0, 0, 0, true_x_axis[0], true_x_axis[1], true_x_axis[2], color='r', linestyle='dashed')
+ ax.quiver(0, 0, 0, true_y_axis[0], true_y_axis[1], true_y_axis[2], color='r', linestyle='dotted')
+ ax.quiver(0, 0, 0, true_z_axis[0], true_z_axis[1], true_z_axis[2], color='r')
+
+ # Plot the predicted orientation (green axes)
+ ax.quiver(0, 0, 0, predicted_x_axis[0], predicted_x_axis[1], predicted_x_axis[2], color='g', linestyle='dashed')
+ ax.quiver(0, 0, 0, predicted_y_axis[0], predicted_y_axis[1], predicted_y_axis[2], color='g', linestyle='dotted')
+ ax.quiver(0, 0, 0, predicted_z_axis[0], predicted_z_axis[1], predicted_z_axis[2], color='g')
+
+ # Set plot limits and labels
+ ax.set_xlim([-1, 1])
+ ax.set_ylim([-1, 1])
+ ax.set_zlim([-1, 1])
+ ax.set_xlabel('X')
+ ax.set_ylabel('Y')
+ ax.set_zlabel('Z')
+ # ax.set_title('True vs Predicted Quaternion Orientations')
+
+ # ax.set_title(f'Subplot {i + 1}')
+ # plt.tight_layout() # prevent overlapping titles and labels
+
+ # Show the plot
+ plt.show()
+ # plt.savefig(PATH + '/ori_subplots.png')
diff --git a/PyTorch Model/plot_poses.py b/PyTorch Model/plot_poses.py
new file mode 100644
index 0000000000000000000000000000000000000000..febb3d82ebb3eb35b9f91e26ba2b71723dfb4685
--- /dev/null
+++ b/PyTorch Model/plot_poses.py
@@ -0,0 +1,32 @@
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+
+PATH = 'runs_TurtlebotTourCVSSP/run 3 - moar data'
+pose_true = pd.read_csv(PATH + '/pose_true.csv')
+pose_estim = pd.read_csv(PATH + '/pose_estim.csv')
+
+# read ALL pos values
+N = len(pose_true)
+position_true = pose_true.iloc[:N, 0:3].values
+position_estim = pose_estim.iloc[:N, 0:3].values
+
+fig = plt.figure()
+ax = fig.add_subplot(111, projection='3d') # Create a 3D subplot
+ax.set_xlabel('x')
+ax.set_ylabel('y')
+ax.set_zlabel('z')
+
+ax.scatter(position_true[:N, 0], position_true[:N, 1], position_true[:N, 2], c='r', marker='o', label='truth')
+ax.scatter(position_estim[:N, 0], position_estim[:N, 1], position_estim[:N, 2], c='b', marker='o', label='estimation')
+
+# # connect the true and estim points with green lines
+# for i in range(position_true.__len__()):
+# position_set = np.vstack((position_true[i,:], position_estim[i,:]))
+# ax.plot(position_set[:,0], position_set[:,1], position_set[:,2], color='green', linewidth=0.5)
+
+ax.legend()
+
+# plt.axis('scaled')
+plt.show()
diff --git a/PyTorch Model/pose_utils.py b/PyTorch Model/pose_utils.py
index 4161e53bb458a00ece9c819f032b689be46ecd6b..ac3126ab41ca94b6e32280d9287ee19a811a8694 100644
--- a/PyTorch Model/pose_utils.py
+++ b/PyTorch Model/pose_utils.py
@@ -2,7 +2,7 @@ import torch
import numpy as np
-def quat_to_euler(q, is_degree=False):
+def quat_to_euler(q, is_degree=True):
w, x, y, z = q[0], q[1], q[2], q[3]
t0 = +2.0 * (w * x + y * z)
@@ -26,21 +26,26 @@ def quat_to_euler(q, is_degree=False):
return np.array([roll, pitch, yaw])
-def array_dist(pred, target):
- return np.linalg.norm(pred - target, 2)
+def position_dist(pred, target):
+ return np.linalg.norm(pred-target, ord=2)
-def position_dist(pred, target):
- return np.linalg.norm(pred-target, 2)
+def rotation_dist(pred, target): # angle-axis
+
+ # Calculate quaternion difference
+ # scalar dot product by element-wise multiplying each element, then adding
+ quaternion_difference = np.dot(target, pred)
+ # Calculate rotation error in radians
+ # arc cosine (inverse cosine) is often used in trigonometric calculations, and in this context,
+ # it is used to compute the rotation angle based on the quaternion difference.
+ alpha = 2 * np.arccos(np.abs(quaternion_difference))
-def rotation_dist(pred, target):
- pred = quat_to_euler(pred)
- target = quat_to_euler(target)
+ # Convert radians to degrees
+ return alpha * (180.0 / np.pi)
- return np.linalg.norm(pred-target, 2)
-# only for bayesian
+# only for bayesian posenet
def fit_gaussian(pose_quat):
# pose_quat = pose_quat.detach().cpu().numpy()