320 lines
14 KiB
Python
320 lines
14 KiB
Python
import functools
|
|
|
|
import numpy as np
|
|
from solver.kpam import mp_terms, term_spec
|
|
from solver.kpam.optimization_problem import OptimizationProblemkPAM
|
|
from solver.kpam.transformations import affine_matrix_from_points
|
|
|
|
|
|
class OptimizationBuilderkPAM(object):
|
|
def __init__(self, specification):
|
|
self._specification = specification
|
|
|
|
def _check_and_get_keypoint_idx(self, keypoint_name, input_idx=-1):
|
|
# type: (str, int) -> int
|
|
"""
|
|
Given the keypoint name, return the index for that keypoint.
|
|
If input_idx is valid, also verify two index match each other
|
|
:param keypoint_name:
|
|
:param input_idx:
|
|
:return:
|
|
"""
|
|
assert keypoint_name in self._specification.keypoint_name2idx
|
|
idx = self._specification.keypoint_name2idx[keypoint_name]
|
|
|
|
if input_idx >= 0:
|
|
assert idx == input_idx
|
|
return idx
|
|
|
|
def build_optimization(self, keypoint_observation): # type: (np.ndarray) -> OptimizationProblemkPAM
|
|
"""
|
|
Given the observed keypoint, build a mathematical program whose
|
|
solution is the transform that map keypoint_observation to target
|
|
:param keypoint_observation: (N, 3) keypoint location
|
|
:return:
|
|
"""
|
|
# Basic check
|
|
assert keypoint_observation.shape[1] == 3
|
|
assert keypoint_observation.shape[0] == len(self._specification.keypoint_name2idx)
|
|
|
|
# Empty initialization of optimization problem
|
|
problem = OptimizationProblemkPAM()
|
|
problem.T_init = np.eye(4)
|
|
problem.mp = None
|
|
problem.xyzrpy = None
|
|
problem.has_solution = False
|
|
|
|
# Build the initial transformation
|
|
problem.T_init = self._build_initial_transformation(keypoint_observation)
|
|
|
|
# Build the mathematical_program with decision variable
|
|
problem.build_empty_mp()
|
|
|
|
# Process the cost terms
|
|
for cost in self._specification.cost_list:
|
|
self._process_cost_term(problem, keypoint_observation, cost)
|
|
|
|
# Process the constraint terms
|
|
for constraint in self._specification.constraint_list:
|
|
self._process_constraint_term(problem, keypoint_observation, constraint)
|
|
|
|
return problem
|
|
|
|
def _build_initial_transformation(self, keypoint_observation): # type: (np.ndarray) -> np.ndarray
|
|
# The container for all keypoint that has nominal target
|
|
from_keypoint = []
|
|
to_keypoint = []
|
|
|
|
# Get all the keypoint that has nominal target
|
|
keypoint_name2idx = self._specification.keypoint_name2idx
|
|
for keypoint_name in keypoint_name2idx:
|
|
exist_target, location = self._specification.get_nominal_target_position(keypoint_name)
|
|
if exist_target:
|
|
to_keypoint.append(location)
|
|
idx = keypoint_name2idx[keypoint_name]
|
|
from_keypoint.append(
|
|
[
|
|
keypoint_observation[idx, 0],
|
|
keypoint_observation[idx, 1],
|
|
keypoint_observation[idx, 2],
|
|
]
|
|
)
|
|
|
|
# Depends on whether we have enough keypoints
|
|
assert len(from_keypoint) == len(to_keypoint)
|
|
if len(from_keypoint) < 3:
|
|
return np.eye(4)
|
|
|
|
# There exists enough keypoint, let's compute using karbas algorithm
|
|
# The np array are (3, N) format of keypoints
|
|
from_keypoint_np = np.zeros(shape=(3, len(from_keypoint)))
|
|
to_keypoint_np = np.zeros_like(from_keypoint_np)
|
|
for i in range(len(from_keypoint)):
|
|
for j in range(3):
|
|
from_keypoint_np[j, i] = from_keypoint[i][j]
|
|
to_keypoint_np[j, i] = to_keypoint[i][j]
|
|
T = affine_matrix_from_points(from_keypoint_np, to_keypoint_np, shear=False, scale=False, usesvd=True)
|
|
return T
|
|
|
|
def _process_cost_term(self, problem, keypoint_observation, cost):
|
|
# type: (OptimizationProblemkPAM, np.ndarray, term_spec.OptimizationTermSpec) -> None
|
|
"""
|
|
Global dispatcher for the cost term
|
|
:param problem:
|
|
:param cost:
|
|
:return:
|
|
"""
|
|
if isinstance(cost, term_spec.Point2PointCostL2Spec):
|
|
self._process_keypoint_l2_cost_term(problem, keypoint_observation, cost)
|
|
elif isinstance(cost, term_spec.Point2PlaneCostSpec):
|
|
self._process_point2plane_cost_term(problem, keypoint_observation, cost)
|
|
else:
|
|
raise RuntimeError("The cost term type is unknown/not implemented")
|
|
|
|
def _process_constraint_term(self, problem, keypoint_observation, constraint):
|
|
# type: (OptimizationProblemkPAM, np.ndarray, term_spec.OptimizationTermSpec) -> None
|
|
"""
|
|
Global dispatcher for constraint term
|
|
:param problem:
|
|
:param keypoint_observation:
|
|
:param constraint:
|
|
:return:
|
|
"""
|
|
if isinstance(constraint, term_spec.Point2PointConstraintSpec):
|
|
self._process_point2point_constraint_term(problem, keypoint_observation, constraint)
|
|
elif isinstance(constraint, term_spec.AxisAlignmentConstraintSpec):
|
|
self._process_axis_alignment_constraint_term(problem, constraint)
|
|
elif isinstance(constraint, term_spec.KeypointAxisAlignmentConstraintSpec):
|
|
self._process_keypoint_axisalign_constraint_term(problem, keypoint_observation, constraint)
|
|
elif isinstance(constraint, term_spec.KeypointAxisOrthogonalConstraintSpec):
|
|
self._process_keypoint_axisorthogonal_constraint_term(problem, keypoint_observation, constraint)
|
|
else:
|
|
raise RuntimeError("The constraint type is unknown/not implemented")
|
|
|
|
def _process_keypoint_l2_cost_term(self, problem, keypoint_observation, cost):
|
|
# type: (OptimizationProblemkPAM, np.ndarray, term_spec.Point2PointCostL2Spec) -> None
|
|
idx = self._check_and_get_keypoint_idx(cost.keypoint_name, cost.keypoint_idx)
|
|
|
|
# The current location
|
|
from_keypoints = keypoint_observation[idx, :]
|
|
from_keypoints = np.reshape(from_keypoints, (1, 3))
|
|
|
|
# The target location
|
|
to_keypoints = np.zeros(shape=(1, 3))
|
|
to_keypoints[0, 0] = cost.target_position[0]
|
|
to_keypoints[0, 1] = cost.target_position[1]
|
|
to_keypoints[0, 2] = cost.target_position[2]
|
|
|
|
# The cost function
|
|
# It is wired that the cost term cannot use symbolic expression
|
|
cost_func = functools.partial(
|
|
mp_terms.keypoint_l2_cost,
|
|
problem,
|
|
from_keypoints,
|
|
to_keypoints,
|
|
cost.penalty_weight,
|
|
)
|
|
problem.mp.AddCost(cost_func, problem.xyzrpy)
|
|
|
|
def _process_point2plane_cost_term(self, problem, keypoint_observation, cost):
|
|
# type: (OptimizationProblemkPAM, np.ndarray, term_spec.Point2PlaneCostSpec) -> None
|
|
idx = self._check_and_get_keypoint_idx(cost.keypoint_name, cost.keypoint_idx)
|
|
|
|
# The current location
|
|
from_keypoints = keypoint_observation[idx, :]
|
|
from_keypoints = np.reshape(from_keypoints, (1, 3))
|
|
|
|
# The target location
|
|
to_keypoints = np.zeros(shape=(1, 3))
|
|
to_keypoints[0, 0] = cost.target_position[0]
|
|
to_keypoints[0, 1] = cost.target_position[1]
|
|
to_keypoints[0, 2] = cost.target_position[2]
|
|
|
|
# The cost function
|
|
# It is wired that the cost term cannot use symbolic expression
|
|
cost_func = functools.partial(
|
|
mp_terms.point2plane_cost,
|
|
problem,
|
|
from_keypoints,
|
|
to_keypoints,
|
|
cost.penalty_weight,
|
|
cost.plane_normal,
|
|
)
|
|
problem.mp.AddCost(cost_func, problem.xyzrpy)
|
|
|
|
def _process_point2point_constraint_term(self, problem, keypoint_observation, constraint):
|
|
# type: (OptimizationProblemkPAM, np.ndarray, term_spec.Point2PointConstraintSpec) -> None
|
|
idx = self._check_and_get_keypoint_idx(constraint.keypoint_name, constraint.keypoint_idx)
|
|
|
|
# The current location
|
|
from_keypoint = keypoint_observation[idx, :]
|
|
transformed_point_symbolic_val = mp_terms.transform_point_symbolic(problem, from_keypoint, problem.xyzrpy)
|
|
|
|
# The target location
|
|
to_keypoints = np.asarray(constraint.target_position).copy()
|
|
to_keypoints_lb = to_keypoints - constraint.tolerance
|
|
to_keypoints_ub = to_keypoints + constraint.tolerance
|
|
for j in range(3):
|
|
problem.mp.AddConstraint(transformed_point_symbolic_val[j] >= to_keypoints_lb[j])
|
|
problem.mp.AddConstraint(transformed_point_symbolic_val[j] <= to_keypoints_ub[j])
|
|
|
|
@staticmethod
|
|
def _process_axis_alignment_constraint_term(problem, constraint):
|
|
# type: (OptimizationProblemkPAM, np.ndarray, term_spec.AxisAlignmentConstraintSpec) -> None
|
|
assert constraint.tolerance < 1
|
|
assert constraint.tolerance > 0
|
|
|
|
# Check the axis
|
|
from_axis = np.asarray(constraint.from_axis, dtype=np.float32).copy() # type: np.ndarray
|
|
to_axis = np.asarray(constraint.target_axis, dtype=np.float32).copy() # type: np.ndarray
|
|
norm_from = np.linalg.norm(from_axis)
|
|
norm_to = np.linalg.norm(to_axis)
|
|
if norm_from < 1e-4 or norm_to < 1e-4:
|
|
print("Warning: the axis is zero, the constraint is not added to optimization.")
|
|
return
|
|
|
|
# Do normalization
|
|
from_axis *= 1.0 / norm_from
|
|
to_axis *= 1.0 / norm_to
|
|
|
|
# Build the term
|
|
vector_dot_symbolic_val = mp_terms.vector_dot_symbolic(problem, from_axis, to_axis, problem.xyzrpy)
|
|
|
|
# Add to mp
|
|
problem.mp.AddConstraint(
|
|
vector_dot_symbolic_val >= constraint.target_inner_product - constraint.tolerance
|
|
) # - constraint.target_inner_product # ← One space after #
|
|
problem.mp.AddConstraint(
|
|
vector_dot_symbolic_val <= constraint.target_inner_product + constraint.tolerance
|
|
) # - constraint.target_inner_product # ← One space after #
|
|
|
|
def _process_keypoint_axisalign_constraint_term(self, problem, keypoint_observation, constraint):
|
|
# type: (OptimizationProblemkPAM, np.ndarray, term_spec.KeypointAxisAlignmentConstraintSpec) -> None
|
|
from_idx = self._check_and_get_keypoint_idx(
|
|
constraint.axis_from_keypoint_name, constraint.axis_from_keypoint_idx
|
|
)
|
|
to_idx = self._check_and_get_keypoint_idx(constraint.axis_to_keypoint_name, constraint.axis_to_keypoint_idx)
|
|
|
|
# Compute the axis
|
|
from_point = keypoint_observation[from_idx, :]
|
|
to_point = keypoint_observation[to_idx, :]
|
|
vector = to_point - from_point
|
|
norm_vec = np.linalg.norm(vector)
|
|
if norm_vec < 1e-4:
|
|
print("Warning: the axis is zero, the constraint is not added to optimization.")
|
|
return
|
|
vector *= 1.0 / norm_vec
|
|
|
|
# Build axis alignment constraint
|
|
axis_constraint = term_spec.AxisAlignmentConstraintSpec()
|
|
axis_constraint.from_axis = [vector[0], vector[1], vector[2]]
|
|
axis_constraint.target_axis = constraint.target_axis
|
|
axis_constraint.tolerance = constraint.tolerance
|
|
axis_constraint.target_inner_product = constraint.target_inner_product
|
|
|
|
# OK
|
|
self._process_axis_alignment_constraint_term(problem, axis_constraint)
|
|
|
|
@staticmethod
|
|
def _process_axis_orthogonal_constraint_term(problem, constraint):
|
|
# make the axis orthogonal as the constraint
|
|
# type: (OptimizationProblemkPAM, np.ndarray, term_spec.AxisAlignmentConstraintSpec) -> None
|
|
assert constraint.tolerance < 1
|
|
assert constraint.tolerance > 0
|
|
|
|
# Check the axis
|
|
from_axis = np.asarray(constraint.from_axis, dtype=np.float32).copy() # type: np.ndarray
|
|
to_axis = np.asarray(constraint.target_axis, dtype=np.float32).copy() # type: np.ndarray
|
|
norm_from = np.linalg.norm(from_axis)
|
|
norm_to = np.linalg.norm(to_axis)
|
|
if norm_from < 1e-4 or norm_to < 1e-4:
|
|
print("Warning: the axis is zero, the constraint is not added to optimization.")
|
|
return
|
|
|
|
# Do normalization
|
|
from_axis *= 1.0 / norm_from
|
|
to_axis *= 1.0 / norm_to
|
|
|
|
# Build the term
|
|
vector_dot_symbolic_val = mp_terms.vector_dot_symbolic(problem, from_axis, to_axis, problem.xyzrpy)
|
|
|
|
# Add to mp
|
|
problem.mp.AddConstraint(
|
|
vector_dot_symbolic_val - constraint.target_inner_product <= constraint.tolerance
|
|
) # 1.0 -
|
|
problem.mp.AddConstraint(
|
|
vector_dot_symbolic_val - constraint.target_inner_product >= -constraint.tolerance
|
|
) # 1.0 -
|
|
|
|
def _process_keypoint_axisorthogonal_constraint_term(self, problem, keypoint_observation, constraint):
|
|
# type: (OptimizationProblemkPAM, np.ndarray, term_spec.KeypointAxisAlignmentConstraintSpec) -> None
|
|
from_idx = self._check_and_get_keypoint_idx(
|
|
constraint.axis_from_keypoint_name, constraint.axis_from_keypoint_idx
|
|
)
|
|
to_idx = self._check_and_get_keypoint_idx(constraint.axis_to_keypoint_name, constraint.axis_to_keypoint_idx)
|
|
|
|
# Compute the axis
|
|
from_point = keypoint_observation[from_idx, :]
|
|
to_point = keypoint_observation[to_idx, :]
|
|
vector = to_point - from_point
|
|
norm_vec = np.linalg.norm(vector)
|
|
if norm_vec < 1e-4:
|
|
print("Warning: the axis is zero, the constraint is not added to optimization.")
|
|
return
|
|
vector *= 1.0 / norm_vec
|
|
|
|
# Build axis alignment constraint
|
|
axis_constraint = term_spec.AxisAlignmentConstraintSpec()
|
|
axis_constraint.from_axis = [vector[0], vector[1], vector[2]]
|
|
axis_constraint.target_axis = constraint.target_axis
|
|
axis_constraint.tolerance = constraint.tolerance
|
|
axis_constraint.target_inner_product = constraint.target_inner_product
|
|
|
|
# OK
|
|
self._process_axis_orthogonal_constraint_term(problem, axis_constraint)
|
|
|
|
|
|
def test_builder():
|
|
"""Legacy test entry; no longer used."""
|
|
raise RuntimeError("test_builder is deprecated and should not be called.")
|