152 lines
5.0 KiB
Python
152 lines
5.0 KiB
Python
import numpy as np
|
|
|
|
class PoseUtil:
|
|
ROTATION = 1
|
|
TRANSLATION = 2
|
|
SCALE = 3
|
|
|
|
@staticmethod
|
|
def get_uniform_translation(trans_m_min, trans_m_max, trans_unit, debug=False):
|
|
if isinstance(trans_m_min, list):
|
|
x_min, y_min, z_min = trans_m_min
|
|
x_max, y_max, z_max = trans_m_max
|
|
else:
|
|
x_min, y_min, z_min = trans_m_min, trans_m_min, trans_m_min
|
|
x_max, y_max, z_max = trans_m_max, trans_m_max, trans_m_max
|
|
|
|
x = np.random.uniform(x_min, x_max)
|
|
y = np.random.uniform(y_min, y_max)
|
|
z = np.random.uniform(z_min, z_max)
|
|
translation = np.array([x, y, z])
|
|
if trans_unit == "cm":
|
|
translation = translation / 100
|
|
if debug:
|
|
print("uniform translation:", translation)
|
|
return translation
|
|
|
|
@staticmethod
|
|
def get_uniform_rotation(rot_degree_min=0, rot_degree_max=180, debug=False):
|
|
axis = np.random.randn(3)
|
|
axis /= np.linalg.norm(axis)
|
|
theta = np.random.uniform(
|
|
rot_degree_min / 180 * np.pi, rot_degree_max / 180 * np.pi
|
|
)
|
|
|
|
K = np.array(
|
|
[[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0]]
|
|
)
|
|
R = np.eye(3) + np.sin(theta) * K + (1 - np.cos(theta)) * (K @ K)
|
|
if debug:
|
|
print("uniform rotation:", theta * 180 / np.pi)
|
|
return R
|
|
|
|
@staticmethod
|
|
def get_uniform_pose(
|
|
trans_min, trans_max, rot_min=0, rot_max=180, trans_unit="cm", debug=False
|
|
):
|
|
translation = PoseUtil.get_uniform_translation(
|
|
trans_min, trans_max, trans_unit, debug
|
|
)
|
|
rotation = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
|
|
pose = np.eye(4)
|
|
pose[:3, :3] = rotation
|
|
pose[:3, 3] = translation
|
|
return pose
|
|
|
|
@staticmethod
|
|
def get_n_uniform_pose(
|
|
trans_min,
|
|
trans_max,
|
|
rot_min=0,
|
|
rot_max=180,
|
|
n=1,
|
|
trans_unit="cm",
|
|
fix=None,
|
|
contain_canonical=True,
|
|
debug=False,
|
|
):
|
|
if fix == PoseUtil.ROTATION:
|
|
translations = np.zeros((n, 3))
|
|
for i in range(n):
|
|
translations[i] = PoseUtil.get_uniform_translation(
|
|
trans_min, trans_max, trans_unit, debug
|
|
)
|
|
if contain_canonical:
|
|
translations[0] = np.zeros(3)
|
|
rotations = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
|
|
elif fix == PoseUtil.TRANSLATION:
|
|
rotations = np.zeros((n, 3, 3))
|
|
for i in range(n):
|
|
rotations[i] = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
|
|
if contain_canonical:
|
|
rotations[0] = np.eye(3)
|
|
translations = PoseUtil.get_uniform_translation(
|
|
trans_min, trans_max, trans_unit, debug
|
|
)
|
|
else:
|
|
translations = np.zeros((n, 3))
|
|
rotations = np.zeros((n, 3, 3))
|
|
for i in range(n):
|
|
translations[i] = PoseUtil.get_uniform_translation(
|
|
trans_min, trans_max, trans_unit, debug
|
|
)
|
|
for i in range(n):
|
|
rotations[i] = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
|
|
if contain_canonical:
|
|
translations[0] = np.zeros(3)
|
|
rotations[0] = np.eye(3)
|
|
|
|
pose = np.eye(4, 4, k=0)[np.newaxis, :].repeat(n, axis=0)
|
|
pose[:, :3, :3] = rotations
|
|
pose[:, :3, 3] = translations
|
|
|
|
return pose
|
|
|
|
@staticmethod
|
|
def get_n_uniform_pose_batch(
|
|
trans_min,
|
|
trans_max,
|
|
rot_min=0,
|
|
rot_max=180,
|
|
n=1,
|
|
batch_size=1,
|
|
trans_unit="cm",
|
|
fix=None,
|
|
contain_canonical=False,
|
|
debug=False,
|
|
):
|
|
|
|
batch_poses = []
|
|
for i in range(batch_size):
|
|
pose = PoseUtil.get_n_uniform_pose(
|
|
trans_min,
|
|
trans_max,
|
|
rot_min,
|
|
rot_max,
|
|
n,
|
|
trans_unit,
|
|
fix,
|
|
contain_canonical,
|
|
debug,
|
|
)
|
|
batch_poses.append(pose)
|
|
pose_batch = np.stack(batch_poses, axis=0)
|
|
return pose_batch
|
|
|
|
@staticmethod
|
|
def get_uniform_scale(scale_min, scale_max, debug=False):
|
|
if isinstance(scale_min, list):
|
|
x_min, y_min, z_min = scale_min
|
|
x_max, y_max, z_max = scale_max
|
|
else:
|
|
x_min, y_min, z_min = scale_min, scale_min, scale_min
|
|
x_max, y_max, z_max = scale_max, scale_max, scale_max
|
|
|
|
x = np.random.uniform(x_min, x_max)
|
|
y = np.random.uniform(y_min, y_max)
|
|
z = np.random.uniform(z_min, z_max)
|
|
scale = np.array([x, y, z])
|
|
if debug:
|
|
print("uniform scale:", scale)
|
|
return scale
|