Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,791 Bytes
c28dddb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
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 |