339 lines
14 KiB
Python
339 lines
14 KiB
Python
import numpy as np
|
|
import open3d as o3d
|
|
from scipy.spatial.transform import Rotation as R
|
|
|
|
|
|
def get_rotation(right_vector, up_vector, look_at_vector, z_front=True):
|
|
rotation = np.eye(3)
|
|
rotation[:, 0] = right_vector
|
|
rotation[:, 1] = up_vector
|
|
rotation[:, 2] = look_at_vector
|
|
if not z_front:
|
|
rotation = rotation @ R.from_euler('y', 180, degrees=True).as_matrix()
|
|
return rotation
|
|
|
|
|
|
class TrajDescription:
|
|
RRot: str = "r_rot"
|
|
LRot: str = "l_rot"
|
|
UpRot: str = "up_rot"
|
|
DownRot: str = "down_rot"
|
|
Front: str = "front"
|
|
Back: str = "back"
|
|
Left: str = "left"
|
|
Right: str = "right"
|
|
Up: str = "up"
|
|
Down: str = "down"
|
|
|
|
def __init__(self, start_pt, look_at_pt, description_list):
|
|
self.z_front = False
|
|
self.up_vector = np.array([0., 0., -1.])
|
|
self.init_pose = self.get_pose(start_pt, look_at_pt, self.up_vector)
|
|
self.description_list = description_list
|
|
self.poses = self.generate_poses()
|
|
|
|
def get_pose(self, start_pt, look_at_pt, up_vector):
|
|
look_at_vector = look_at_pt - start_pt
|
|
look_at_vector = look_at_vector / np.linalg.norm(look_at_vector)
|
|
up_vector = up_vector
|
|
right_vector = np.cross(look_at_vector, up_vector)
|
|
right_vector = right_vector / np.linalg.norm(right_vector)
|
|
up_vector = np.cross(look_at_vector, right_vector)
|
|
up_vector = up_vector / np.linalg.norm(up_vector)
|
|
|
|
pose = np.eye(4)
|
|
pose[:3, :3] = get_rotation(right_vector, up_vector, look_at_vector, self.z_front)
|
|
pose[:3, 3] = start_pt
|
|
return pose
|
|
|
|
def generate_poses(self):
|
|
poses = [self.init_pose]
|
|
current_pose = self.init_pose.copy()
|
|
|
|
for description in self.description_list:
|
|
pose_type = description["type"]
|
|
value = description["value"]
|
|
|
|
if pose_type == self.RRot:
|
|
rotation_matrix = R.from_euler('y', -value, degrees=True).as_matrix()
|
|
new_pose = current_pose.copy()
|
|
new_pose[:3, :3] = current_pose[:3, :3] @ rotation_matrix
|
|
poses.append(new_pose)
|
|
poses.append(current_pose)
|
|
|
|
elif pose_type == self.LRot:
|
|
rotation_matrix = R.from_euler('y', value, degrees=True).as_matrix()
|
|
new_pose = current_pose.copy()
|
|
new_pose[:3, :3] = current_pose[:3, :3] @ rotation_matrix
|
|
poses.append(new_pose)
|
|
poses.append(current_pose)
|
|
|
|
elif pose_type == self.UpRot:
|
|
rotation_matrix = R.from_euler('x', -value, degrees=True).as_matrix()
|
|
new_pose = current_pose.copy()
|
|
new_pose[:3, :3] = current_pose[:3, :3] @ rotation_matrix
|
|
poses.append(new_pose)
|
|
poses.append(current_pose)
|
|
|
|
elif pose_type == self.DownRot:
|
|
rotation_matrix = R.from_euler('x', value, degrees=True).as_matrix()
|
|
new_pose = current_pose.copy()
|
|
new_pose[:3, :3] = current_pose[:3, :3] @ rotation_matrix
|
|
poses.append(new_pose)
|
|
poses.append(current_pose)
|
|
|
|
elif pose_type == self.Front:
|
|
new_pose = current_pose.copy()
|
|
new_pose[1, 3] = current_pose[1, 3] + value
|
|
poses.append(new_pose)
|
|
current_pose = new_pose
|
|
|
|
elif pose_type == self.Back:
|
|
new_pose = current_pose.copy()
|
|
new_pose[1, 3] = current_pose[1, 3] - value
|
|
poses.append(new_pose)
|
|
current_pose = new_pose
|
|
|
|
elif pose_type == self.Left:
|
|
new_pose = current_pose.copy()
|
|
new_pose[0, 3] = current_pose[0, 3] - value
|
|
poses.append(new_pose)
|
|
current_pose = new_pose
|
|
|
|
elif pose_type == self.Right:
|
|
new_pose = current_pose.copy()
|
|
new_pose[0, 3] = current_pose[0, 3] + value
|
|
poses.append(new_pose)
|
|
current_pose = new_pose
|
|
|
|
elif pose_type == self.Up:
|
|
new_pose = current_pose.copy()
|
|
new_pose[2, 3] = current_pose[2, 3] + value
|
|
poses.append(new_pose)
|
|
current_pose = new_pose
|
|
|
|
elif pose_type == self.Down:
|
|
new_pose = current_pose.copy()
|
|
new_pose[2, 3] = current_pose[2, 3] - value
|
|
poses.append(new_pose)
|
|
current_pose = new_pose
|
|
|
|
else:
|
|
raise ValueError(f"Invalid description: {description}")
|
|
|
|
return poses
|
|
|
|
|
|
class TrajGenerator:
|
|
def __init__(self, view_num: int, look_at_pts: np.ndarray, up_vectors: np.ndarray, start_pts: np.ndarray,
|
|
reverse_order=False, z_front=True):
|
|
"""Generates a circular trajectory around a specified point.
|
|
Args:
|
|
view_num (int): Number of views/poses to generate.
|
|
look_at_pts (np.ndarray): The point(s) to look at, shape (N, 3).
|
|
up_vectors (np.ndarray): The up direction(s), shape (N, 3).
|
|
start_pts (np.ndarray): The starting point(s) of the trajectory, shape (N, 3).
|
|
reverse_order (bool): If True, reverse the order of generated poses.
|
|
z_front (bool): If True, the camera's Z-axis points forward; otherwise, it points backward.
|
|
"""
|
|
assert look_at_pts.shape[0] == up_vectors.shape[0] == start_pts.shape[
|
|
0], "Input arrays must have the same number of points."
|
|
self.view_num = view_num
|
|
self.look_at_pts = look_at_pts
|
|
self.up_vectors = up_vectors
|
|
self.start_pts = start_pts
|
|
self.z_front = z_front
|
|
self.poses = self.generate_poses(reverse_order)
|
|
self.interpolated_poses = self.interpolate_poses(self.poses, num_interpolations=self.view_num, method='slerp')
|
|
|
|
def generate_poses(self, reverse_order=False):
|
|
poses = []
|
|
length = self.look_at_pts.shape[0]
|
|
for i in range(length):
|
|
order = i if not reverse_order else length - i - 1
|
|
pose = self.get_pose(self.start_pts[order], self.look_at_pts[order], self.up_vectors[order])
|
|
poses.append(pose)
|
|
return poses
|
|
|
|
def get_pose(self, start_pt, look_at_pt, up_vector, with_noise=False):
|
|
look_at_vector = look_at_pt - start_pt
|
|
look_at_vector = look_at_vector / np.linalg.norm(look_at_vector)
|
|
up_vector = up_vector
|
|
if with_noise:
|
|
noise = np.random.uniform(0, 0.1, up_vector.shape)
|
|
right_vector = np.cross(look_at_vector, up_vector + noise)
|
|
else:
|
|
right_vector = np.cross(look_at_vector, up_vector)
|
|
right_vector = right_vector / np.linalg.norm(right_vector)
|
|
up_vector = np.cross(look_at_vector, right_vector)
|
|
up_vector = up_vector / np.linalg.norm(up_vector)
|
|
|
|
pose = np.eye(4)
|
|
pose[:3, :3] = self.get_rotation(right_vector, up_vector, look_at_vector, self.z_front)
|
|
pose[:3, 3] = start_pt
|
|
return pose
|
|
|
|
def get_rotation(self, right_vector, up_vector, look_at_vector, z_front=True):
|
|
rotation = np.eye(3)
|
|
rotation[:, 0] = right_vector
|
|
rotation[:, 1] = up_vector
|
|
rotation[:, 2] = look_at_vector
|
|
if not z_front:
|
|
rotation = rotation @ R.from_euler('y', 180, degrees=True).as_matrix()
|
|
return rotation
|
|
|
|
def interpolate_poses(self, poses, num_interpolations, method='slerp'):
|
|
assert method in ['slerp', 'linear'], "Method must be 'slerp' or 'linear'."
|
|
return poses
|
|
|
|
|
|
class CircularTrajGenerator(TrajGenerator):
|
|
def __init__(self, view_num, look_at_pt, top_down_dist, top_down_direrction, radius, reverse_order=False,
|
|
z_front=True):
|
|
start_pt = look_at_pt + np.array([radius, 0, top_down_dist])
|
|
look_at_pts = np.tile(look_at_pt, (view_num, 1))
|
|
up_vectors = np.tile(top_down_direrction, (view_num, 1))
|
|
start_pts = np.zeros((view_num, 3))
|
|
for i in range(view_num):
|
|
angle = np.pi * 2 * i / view_num
|
|
start_point_x = look_at_pt[0] + radius * np.cos(angle)
|
|
start_point_y = look_at_pt[1] + radius * np.sin(angle)
|
|
start_point_z = look_at_pt[2] + top_down_dist
|
|
start_pts[i] = np.array([start_point_x, start_point_y, start_point_z])
|
|
super().__init__(view_num, look_at_pts, up_vectors, start_pts, reverse_order, z_front)
|
|
|
|
|
|
class TrajVisualizer:
|
|
def __init__(self, poses: list) -> None:
|
|
"""Visualizes SE3 poses."""
|
|
self.poses = poses
|
|
|
|
def visualize(self):
|
|
geometries = []
|
|
for pose in self.poses:
|
|
mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
|
|
mesh_frame.transform(pose)
|
|
geometries.append(mesh_frame)
|
|
o3d.visualization.draw_geometries(geometries)
|
|
|
|
|
|
def generate_camera_data_from_poses(poses, name="unknown"):
|
|
camera_data = {
|
|
"name": name,
|
|
"width": 640,
|
|
"height": 480,
|
|
"focus_length": 13,
|
|
"horizontal_aperture": 15.2908,
|
|
"vertical_aperture": 15.2908,
|
|
"is_local": False,
|
|
"trajectory": []
|
|
}
|
|
|
|
for pose in poses:
|
|
position = pose[:3, 3].tolist()
|
|
rotation_matrix = pose[:3, :3]
|
|
quaternion_xyzw = R.from_matrix(rotation_matrix).as_quat().tolist()
|
|
quaternion_wxyz = [quaternion_xyzw[3], quaternion_xyzw[0], quaternion_xyzw[1], quaternion_xyzw[2]]
|
|
|
|
camera_data["trajectory"].append({
|
|
"position": position,
|
|
"quaternion": quaternion_wxyz
|
|
})
|
|
|
|
return camera_data
|
|
|
|
|
|
def generate_circular_trajectory(top_down_dist=0.6, radius=0.3, look_at_pt=(0, 0, 0)):
|
|
VIEW_NUM = 100
|
|
LOOK_AT_PT = np.array(look_at_pt)
|
|
TOP_DOWN_DIST = [top_down_dist]
|
|
TOP_DOWN_DIRECTION = np.array([0., 0., -1.])
|
|
RADIUS = [radius]
|
|
|
|
pose_generator = CircularTrajGenerator(VIEW_NUM, LOOK_AT_PT, TOP_DOWN_DIST[0], TOP_DOWN_DIRECTION, RADIUS[0],
|
|
reverse_order=False, z_front=False)
|
|
return pose_generator.poses
|
|
|
|
|
|
def generate_trajectory_from_description(start_pt, look_at_pt, description_list):
|
|
traj_description = TrajDescription(start_pt, look_at_pt, description_list)
|
|
return traj_description.poses
|
|
|
|
|
|
def generate_front_left_trajectory():
|
|
VIEW_NUM = 100
|
|
LOOK_AT_PT = np.array([-0.717877984046936, -2.1181700229644775, 0.7807716727256775])
|
|
TOP_DOWN_DIST = [0.6]
|
|
TOP_DOWN_DIRECTION = np.array([0., 0., -1.])
|
|
RADIUS = [0.3]
|
|
|
|
start_pt = LOOK_AT_PT + np.array([RADIUS[0], 0, TOP_DOWN_DIST[0]])
|
|
look_at_pts = np.tile(LOOK_AT_PT, (VIEW_NUM, 1))
|
|
up_vectors = np.tile(TOP_DOWN_DIRECTION, (VIEW_NUM, 1))
|
|
start_pts = np.zeros((VIEW_NUM, 3))
|
|
for i in range(VIEW_NUM):
|
|
angle = np.pi * 2 * i / VIEW_NUM + np.pi / 4
|
|
start_point_x = LOOK_AT_PT[0] + RADIUS[0] * np.cos(angle)
|
|
start_point_y = LOOK_AT_PT[1] + RADIUS[0] * np.sin(angle)
|
|
start_point_z = LOOK_AT_PT[2] + TOP_DOWN_DIST[0]
|
|
start_pts[i] = np.array([start_point_x, start_point_y, start_point_z])
|
|
|
|
pose_generator = TrajGenerator(VIEW_NUM, look_at_pts, up_vectors, start_pts, reverse_order=False, z_front=False)
|
|
return pose_generator.poses
|
|
|
|
|
|
def main():
|
|
workspace_size = np.array([0.816116, 0.675734, 0.3])
|
|
workspace_position = np.array([-0.717877984046936, -2.1181700229644775, 0.7807716727256775])
|
|
desc_1 = [
|
|
# {"type": TrajDescription.LRot, "value": 20},
|
|
# {"type": TrajDescription.RRot, "value": 20},
|
|
# {"type": TrajDescription.DownRot, "value": 20},
|
|
# {"type": TrajDescription.UpRot, "value": 20},
|
|
{"type": TrajDescription.Left, "value": workspace_size[0]},
|
|
{"type": TrajDescription.Front, "value": workspace_size[1]},
|
|
# {"type": TrajDescription.Back, "value": workspace_size[1]},
|
|
# {"type": TrajDescription.Right, "value": workspace_size[0]}
|
|
]
|
|
desc_2 = [
|
|
{"type": TrajDescription.LRot, "value": 20},
|
|
{"type": TrajDescription.RRot, "value": 20},
|
|
{"type": TrajDescription.DownRot, "value": 20},
|
|
{"type": TrajDescription.UpRot, "value": 20},
|
|
{"type": TrajDescription.Right, "value": workspace_size[0]},
|
|
{"type": TrajDescription.Front, "value": workspace_size[1]},
|
|
{"type": TrajDescription.Back, "value": workspace_size[1]},
|
|
{"type": TrajDescription.Left, "value": workspace_size[0]}
|
|
]
|
|
desc_3 = [
|
|
{"type": TrajDescription.DownRot, "value": 20},
|
|
{"type": TrajDescription.UpRot, "value": 20},
|
|
{"type": TrajDescription.Left, "value": workspace_size[0]},
|
|
{"type": TrajDescription.Front, "value": workspace_size[1]},
|
|
{"type": TrajDescription.RRot, "value": 40},
|
|
{"type": TrajDescription.LRot, "value": 40},
|
|
{"type": TrajDescription.Back, "value": workspace_size[1]},
|
|
{"type": TrajDescription.Right, "value": workspace_size[0]}
|
|
]
|
|
desc_4 = [
|
|
{"type": TrajDescription.DownRot, "value": 20},
|
|
{"type": TrajDescription.UpRot, "value": 20},
|
|
{"type": TrajDescription.Right, "value": workspace_size[0]},
|
|
{"type": TrajDescription.Front, "value": workspace_size[1]},
|
|
{"type": TrajDescription.LRot, "value": 40},
|
|
{"type": TrajDescription.RRot, "value": 40},
|
|
{"type": TrajDescription.Back, "value": workspace_size[1]},
|
|
{"type": TrajDescription.Left, "value": workspace_size[0]}
|
|
]
|
|
|
|
start_pt = workspace_position - np.array([0, workspace_size[1] / 2, -0.3])
|
|
poses_1 = generate_trajectory_from_description(start_pt, workspace_position, desc_1)
|
|
poses_2 = generate_trajectory_from_description(start_pt, workspace_position, desc_2)
|
|
poses_3 = generate_trajectory_from_description(start_pt, workspace_position, desc_3)
|
|
poses_4 = generate_trajectory_from_description(start_pt, workspace_position, desc_4)
|
|
visualizer = TrajVisualizer(poses_1)
|
|
visualizer.visualize()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |