126 lines
4.2 KiB
Python
126 lines
4.2 KiB
Python
"""
|
|
GTSAM Copyright 2010-2018, Georgia Tech Research Corporation,
|
|
Atlanta, Georgia 30332-0415
|
|
All Rights Reserved
|
|
Authors: Frank Dellaert, et al. (see THANKS for the full author list)
|
|
|
|
See LICENSE for the license information
|
|
|
|
Shonan Rotation Averaging CLI reads a *pose* graph, extracts the
|
|
rotation constraints, and runs the Shonan algorithm.
|
|
|
|
Author: Frank Dellaert
|
|
Date: August 2020
|
|
"""
|
|
# pylint: disable=invalid-name, E1101
|
|
|
|
import argparse
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
import gtsam
|
|
from gtsam.utils import plot
|
|
|
|
def estimate_poses_given_rot(factors: gtsam.BetweenFactorPose3s,
|
|
rotations: gtsam.Values,
|
|
d: int = 3):
|
|
""" Estimate Poses from measurements, given rotations. From SfmProblem in shonan.
|
|
|
|
Arguments:
|
|
factors -- data structure with many BetweenFactorPose3 factors
|
|
rotations {Values} -- Estimated rotations
|
|
|
|
Returns:
|
|
Values -- Estimated Poses
|
|
"""
|
|
|
|
I_d = np.eye(d)
|
|
|
|
def R(j):
|
|
return rotations.atRot3(j) if d == 3 else rotations.atRot2(j)
|
|
|
|
def pose(R, t):
|
|
return gtsam.Pose3(R, t) if d == 3 else gtsam.Pose2(R, t)
|
|
|
|
graph = gtsam.GaussianFactorGraph()
|
|
model = gtsam.noiseModel.Unit.Create(d)
|
|
|
|
# Add a factor anchoring t_0
|
|
graph.add(0, I_d, np.zeros((d,)), model)
|
|
|
|
# Add a factor saying t_j - t_i = Ri*t_ij for all edges (i,j)
|
|
for factor in factors:
|
|
keys = factor.keys()
|
|
i, j, Tij = keys[0], keys[1], factor.measured()
|
|
measured = R(i).rotate(Tij.translation())
|
|
graph.add(j, I_d, i, -I_d, measured, model)
|
|
|
|
# Solve linear system
|
|
translations = graph.optimize()
|
|
|
|
# Convert to Values.
|
|
result = gtsam.Values()
|
|
for j in range(rotations.size()):
|
|
tj = translations.at(j)
|
|
result.insert(j, pose(R(j), tj))
|
|
|
|
return result
|
|
|
|
|
|
def run(args):
|
|
"""Run Shonan averaging and then recover translations linearly before saving result."""
|
|
|
|
# Get input file
|
|
if args.input_file:
|
|
input_file = args.input_file
|
|
else:
|
|
if args.named_dataset == "":
|
|
raise ValueError(
|
|
"You must either specify a named dataset or an input file")
|
|
input_file = gtsam.findExampleDataFile(args.named_dataset)
|
|
|
|
if args.dimension == 2:
|
|
print("Running Shonan averaging for SO(2) on ", input_file)
|
|
shonan = gtsam.ShonanAveraging2(input_file)
|
|
if shonan.nrUnknowns() == 0:
|
|
raise ValueError("No 2D pose constraints found, try -d 3.")
|
|
initial = shonan.initializeRandomly()
|
|
rotations, _ = shonan.run(initial, 2, 10)
|
|
factors = gtsam.parse2DFactors(input_file)
|
|
elif args.dimension == 3:
|
|
print("Running Shonan averaging for SO(3) on ", input_file)
|
|
shonan = gtsam.ShonanAveraging3(input_file)
|
|
if shonan.nrUnknowns() == 0:
|
|
raise ValueError("No 3D pose constraints found, try -d 2.")
|
|
initial = shonan.initializeRandomly()
|
|
rotations, _ = shonan.run(initial, 3, 10)
|
|
factors = gtsam.parse3DFactors(input_file)
|
|
else:
|
|
raise ValueError("Can only run SO(2) or SO(3) averaging")
|
|
|
|
print("Recovering translations")
|
|
poses = estimate_poses_given_rot(factors, rotations, args.dimension)
|
|
|
|
print("Writing result to ", args.output_file)
|
|
gtsam.writeG2o(gtsam.NonlinearFactorGraph(), poses, args.output_file)
|
|
|
|
plot.plot_trajectory(1, poses, scale=0.2)
|
|
plot.set_axes_equal(1)
|
|
plt.show()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('-n', '--named_dataset', type=str, default="pose3example-grid",
|
|
help='Find and read frome example dataset file')
|
|
parser.add_argument('-i', '--input_file', type=str, default="",
|
|
help='Read pose constraints graph from the specified file')
|
|
parser.add_argument('-o', '--output_file', type=str, default="shonan.g2o",
|
|
help='Write solution to the specified file')
|
|
parser.add_argument('-d', '--dimension', type=int, default=3,
|
|
help='Optimize over 2D or 3D rotations')
|
|
parser.add_argument("-p", "--plot", action="store_true", default=True,
|
|
help="Plot result")
|
|
run(parser.parse_args())
|