78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
import numpy as np
|
|
from pydrake.all import MathematicalProgram, Solve
|
|
from solver.kpam import SE3_utils
|
|
|
|
|
|
class OptimizationProblemkPAM(object):
|
|
# The property about initial transformation
|
|
# There is always an initial guess, either identity or from keypoint matching
|
|
T_init = np.ndarray(shape=(4, 4)) # type: np.ndarray
|
|
|
|
# The final mathematical program
|
|
mp = None # type: MathematicalProgram
|
|
xyzrpy = None # type: np.ndarray
|
|
|
|
# The solution to the program
|
|
has_solution = False
|
|
xyzrpy_sol = np.ndarray(shape=(6,))
|
|
# T_action.dot(observed_keypoint) = target_keypoint
|
|
T_action = np.ndarray(shape=(4, 4))
|
|
|
|
# Build empty problem
|
|
def build_empty_mp(self):
|
|
# Construct the problem
|
|
mp = MathematicalProgram()
|
|
xyz = mp.NewContinuousVariables(3, "xyz")
|
|
rpy = mp.NewContinuousVariables(3, "rpy")
|
|
xyz_rpy = np.concatenate((xyz, rpy))
|
|
mp.SetInitialGuessForAllVariables(np.zeros(6))
|
|
|
|
# Store the result to problem
|
|
self.mp = mp
|
|
self.xyzrpy = xyz_rpy
|
|
|
|
|
|
class OptimizationProblemkPAMJoint(object):
|
|
# The property about initial transformation
|
|
# There is always an initial guess, either identity or from keypoint matching
|
|
T_init = np.ndarray(shape=(4, 4)) # type: np.ndarray
|
|
|
|
# The final mathematical program
|
|
mp = None # type: MathematicalProgram
|
|
xyzrpy = None # type: np.ndarray
|
|
|
|
# The solution to the program
|
|
has_solution = False
|
|
xyzrpy_sol = np.ndarray(shape=(6,))
|
|
# T_action.dot(observed_keypoint) = target_keypoint
|
|
T_action = np.ndarray(shape=(4, 4))
|
|
|
|
# Build empty problem
|
|
def build_empty_mp(self):
|
|
# Construct the problem
|
|
mp = MathematicalProgram()
|
|
xyz = mp.NewContinuousVariables(3, "xyz")
|
|
rpy = mp.NewContinuousVariables(3, "rpy")
|
|
xyz_rpy = np.concatenate((xyz, rpy))
|
|
mp.SetInitialGuessForAllVariables(np.zeros(6))
|
|
|
|
# Store the result to problem
|
|
self.mp = mp
|
|
self.xyzrpy = xyz_rpy
|
|
|
|
|
|
def solve_kpam(problem): # type: (OptimizationProblemkPAM) -> bool
|
|
result = Solve(problem.mp) # mp.
|
|
if not result.is_success():
|
|
problem.has_solution = False
|
|
return False
|
|
|
|
# Save the result to problem
|
|
problem.xyzrpy_sol = result.get_x_val() # problem.mp.GetSolution(problem.xyzrpy)
|
|
T_eps = SE3_utils.xyzrpy_to_matrix(xyzrpy=problem.xyzrpy_sol)
|
|
problem.T_action = np.dot(problem.T_init, T_eps)
|
|
problem.has_solution = True
|
|
|
|
# OK
|
|
return True
|