gtsam/python/gtsam/examples/ShonanAveragingCLI.py

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())