Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import sys | |
| import numpy as np | |
| import quaternion | |
| sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| from objects.dict_utils import get_base_part_idx | |
| def transform_all_parts(part_vertices, obj_dict, joint_state, | |
| rotation_fix_range=True, dry_run=True): | |
| """ | |
| Transform all parts of the object according to the joint state\n | |
| - part_vertices: vertices of the object in rest pose in the form:\n | |
| - [K_parts, N_vertices, 3]\n | |
| - obj_dict: the object dictionary\n | |
| - joint_state: the joint state in the range of [0, 1]\n | |
| - rotation_fix_range (optional): whether to fix the rotation range to 90 degrees for revolute joints\n | |
| - dry_run (optional): if True, only return the transformation matrices without changing the vertices\n | |
| Return:\n | |
| - part_transformations: records of the transformations applied to the parts\n | |
| """ | |
| part_transformations = [[] for _ in range(len(obj_dict["diffuse_tree"]))] | |
| if joint_state == 0.0: | |
| return part_transformations | |
| # Get a visit order of the parts such that children parts are visited before parents | |
| part_visit_order = [] | |
| base_idx = get_base_part_idx(obj_dict) | |
| indices_to_visit = [base_idx] | |
| while len(indices_to_visit) > 0: # Breadth-first traversal | |
| current_idx = indices_to_visit.pop(0) | |
| part_visit_order.append(current_idx) | |
| # if current_idx == 9: | |
| # import ipdb | |
| # ipdb.set_trace() | |
| indices_to_visit += obj_dict["diffuse_tree"][current_idx]["children"] | |
| part_visit_order.reverse() | |
| # Transform the parts in the visit order - children first, then parents | |
| for i in part_visit_order: | |
| part = obj_dict["diffuse_tree"][i] | |
| joint = part["joint"] | |
| children_idxs = part["children"] | |
| # Store the transformation used to transform the part and its children | |
| applied_tramsformation_matrix = np.eye(4, dtype=np.float32) | |
| applied_rotation_axis_origin = np.array([np.nan, np.nan, np.nan], dtype=np.float32) | |
| applied_transformation_type = "none" | |
| if joint["type"] == "prismatic": | |
| # Translate the part and its children | |
| translation = np.array(joint["axis"]["direction"], dtype=np.float32) * joint["range"][1] * joint_state | |
| if not dry_run: | |
| part_vertices[[i] + children_idxs] += translation | |
| # Store the transformation used | |
| applied_tramsformation_matrix[:3, 3] = translation | |
| applied_transformation_type = "translation" | |
| elif joint["type"] == "revolute" or joint["type"] == "continuous": | |
| if joint["type"] == "revolute": | |
| if not rotation_fix_range: | |
| # Use the full range as specified in the object file | |
| rotation_radian = np.radians(joint["range"][1] * joint_state) | |
| else: | |
| # Fix the rotation range to 90 degrees | |
| rotation_range_sign = np.sign(joint["range"][1]) | |
| rotation_radian = np.radians(rotation_range_sign * 90 * joint_state) | |
| else: | |
| rotation_radian = np.radians(360 * joint_state) | |
| # Prepare the rotation matrix via axis-angle representation and quaternion | |
| rotation_axis_origin = np.array(joint["axis"]["origin"], dtype=np.float32) | |
| rotation_axis_direction = np.array(joint["axis"]["direction"], dtype=np.float32) / np.linalg.norm(joint["axis"]["direction"]) | |
| rotation_matrix = quaternion.as_rotation_matrix(quaternion.from_rotation_vector(rotation_radian * rotation_axis_direction)) | |
| if not dry_run: | |
| # Rotate the part and its children | |
| vertices_to_rotate = (part_vertices[[i] + children_idxs] - rotation_axis_origin) | |
| part_vertices[[i] + children_idxs] = np.matmul(rotation_matrix, vertices_to_rotate.transpose([0, 2, 1])).transpose([0, 2, 1]) + rotation_axis_origin | |
| # Store the transformation used | |
| applied_tramsformation_matrix[:3, :3] = rotation_matrix | |
| applied_rotation_axis_origin = rotation_axis_origin | |
| applied_transformation_type = "rotation" | |
| # Record the transformation used | |
| if not applied_transformation_type == "none": | |
| record = { | |
| "type": applied_transformation_type, | |
| "matrix": applied_tramsformation_matrix, | |
| "rotation_axis_origin": applied_rotation_axis_origin | |
| } | |
| for idx in [i] + children_idxs: | |
| part_transformations[idx].append(record) | |
| return part_transformations |