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