my_utils/render.py

# import os, sys
# sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
# import trimesh
# import pyrender
# import numpy as np
# # import open3d as o3d
# from copy import deepcopy
# os.environ['PYOPENGL_PLATFORM'] = 'egl'
# from my_utils.refs import semantic_color_ref, graph_color_ref, joint_color_ref

# def get_rotation_axis_angle(k, theta):
#     '''
#     Rotation matrix converter from axis-angle using Rodrigues' rotation formula

#     Args:
#         k (np.ndarray): 3D unit vector representing the axis to rotate about.
#         theta (float): Angle to rotate with in radians.

#     Returns:
#         R (np.ndarray): 3x3 rotation matrix.
#     '''
#     if np.linalg.norm(k) == 0.:
#         return np.eye(3)
#     k = k / np.linalg.norm(k)
#     kx, ky, kz = k[0], k[1], k[2]
#     cos, sin = np.cos(theta), np.sin(theta)
#     R = np.zeros((3, 3), dtype=np.float32)
#     R[0, 0] = cos + (kx**2) * (1 - cos)
#     R[0, 1] = kx * ky * (1 - cos) - kz * sin
#     R[0, 2] = kx * kz * (1 - cos) + ky * sin
#     R[1, 0] = kx * ky * (1 - cos) + kz * sin
#     R[1, 1] = cos + (ky**2) * (1 - cos)
#     R[1, 2] = ky * kz * (1 - cos) - kx * sin
#     R[2, 0] = kx * kz * (1 - cos) - ky * sin
#     R[2, 1] = ky * kz * (1 - cos) + kx * sin
#     R[2, 2] = cos + (kz**2) * (1 - cos)
#     return R

# def rescale_axis(jtype, axis_d, axis_o, box_center):
#     '''
#     Function to rescale the axis for rendering
    
#     Args:
#     - jtype (int): joint type
#     - axis_d (np.array): axis direction
#     - axis_o (np.array): axis origin
#     - box_center (np.array): bounding box center

#     Returns:
#     - center (np.array): rescaled axis origin
#     - axis_d (np.array): rescaled axis direction
#     '''
#     if jtype == 0 or jtype == 1:
#         return [0., 0., 0.], [0., 0., 0.]
#     if jtype == 3 or jtype == 4:
#         center = box_center
#     else:
#         center = axis_o + np.dot(axis_d, box_center-axis_o) * axis_d
#     return center.tolist(), axis_d.tolist()

# # def get_axis_mesh(k, axis_o, bbox_center, joint_type):
# #     '''
# #     Function to get the axis mesh

# #     Args:
# #     - k (np.array): axis direction
# #     - center (np.array): axis origin
# #     - bbox_center (np.array): bounding box center
# #     - joint_type (int): joint type
# #     '''
# #     if joint_type == 0 or joint_type == 1 or np.linalg.norm(k) == 0. :
# #         return None
    
# #     k = k / np.linalg.norm(k)

# #     if joint_type == 3 or joint_type == 4: # prismatic or screw
# #         axis_o = bbox_center
# #     else: # revolute or continuous
# #         axis_o = axis_o + np.dot(k, bbox_center-axis_o) * k
# #     axis = o3d.geometry.TriangleMesh.create_arrow(cylinder_radius=0.015, cone_radius=0.03, cylinder_height=1.0, cone_height=0.08)
# #     arrow = np.array([0., 0., 1.], dtype=np.float32)
# #     n = np.cross(arrow, k) 
# #     rad = np.arccos(np.dot(arrow, k))
# #     R_arrow = get_rotation_axis_angle(n, rad)
# #     axis.rotate(R_arrow, center=(0, 0, 0))
# #     axis.translate(axis_o[:3])
# #     axis.compute_vertex_normals()
# #     vertices = np.asarray(axis.vertices)
# #     faces = np.asarray(axis.triangles)
# #     trimesh_axis = trimesh.Trimesh(vertices=vertices, faces=faces)
# #     # trimesh_axis.visual.vertex_colors = np.array([0, 0, 0, 1.0], dtype=np.float32)
# #     trimesh_axis.visual.vertex_colors = np.repeat(np.array([0, 0, 0, 1.0]), vertices.shape[0], axis=0)
# #     return trimesh_axis

# def get_camera_pose(eye, look_at, up):
#     """
#     Compute the 4x4 transformation matrix for a camera pose.
    
#     Parameters:
#         eye (np.ndarray): 3D position of the camera.
#         look_at (np.ndarray): 3D point the camera is looking at.
#         up (np.ndarray): Up vector.
        
#     Returns:
#         pose (np.ndarray): 4x4 transformation matrix representing the camera pose.
#     """
#     # Compute the forward, right, and new up vectors
#     forward = (look_at - eye)
#     forward = forward / np.linalg.norm(forward)
    
#     right = np.cross(forward, up)
#     right = right / np.linalg.norm(right)
    
#     new_up = np.cross(right, forward)
#     new_up = new_up / np.linalg.norm(new_up)
    
#     # Create rotation matrix
#     pose = np.eye(4)
#     pose[0:3, 0] = right
#     pose[0:3, 1] = new_up
#     pose[0:3, 2] = -forward  # Negative because the camera looks along the negative Z axis in its local coordinate
#     pose[0:3, 3] = eye
    
#     return pose

# def get_rotation_axis_angle_box(axis, angle):
#     axis = axis / np.linalg.norm(axis)
#     return trimesh.transformations.rotation_matrix(angle, axis)

# def get_colored_box(center, size, jtype=None, jrange=None, axis_d=None, axis_o=None):
#     '''
#     Create a solid color box and its animated state if joint info is provided

#     Args:
#         center (np.array): box center (3,)
#         size (np.array): box size (3,)
#         color (list or array): RGBA color, e.g. [255, 0, 0, 255]
#         jtype (int): joint type (2=rot, 3=slide, 4=screw, 5=continuous)
#         jrange (list): joint motion range
#         axis_d (np.array): axis direction (3,)
#         axis_o (np.array): axis origin (3,)

#     Returns:
#         box: trimesh.Trimesh at rest
#         box_anim: trimesh.Trimesh after transformation
#     '''
#     size = np.clip(size, a_min=0.005, a_max=3.0)
#     center = np.clip(center, a_min=-3.0, a_max=3.0)

#     # Rest state box
#     box = trimesh.creation.box(extents=size)
#     box.apply_translation(center)

#     # Animated state (deepcopy + transform)
#     box_anim = deepcopy(box)

#     if jtype is not None:
#         if jtype == 2:  # revolute
#             theta = np.deg2rad(jrange[1])
#             T = trimesh.transformations.translation_matrix(axis_o)
#             R_3 = get_rotation_axis_angle(axis_d, theta)
#             R = np.eye(4, dtype=np.float32)
#             R[:3, :3] = R_3
#             T_inv = trimesh.transformations.translation_matrix(-axis_o)
#             box_anim.apply_transform(T @ R @ T_inv)

#         elif jtype == 3:  # prismatic
#             dist = float(jrange[1])
#             T = trimesh.transformations.translation_matrix(axis_d * dist)
#             box_anim.apply_transform(T)

#         elif jtype == 4:  # screw
#             theta = np.pi / 4
#             dist = float(jrange[1])
#             T1 = trimesh.transformations.translation_matrix(-axis_o)
#             R = get_rotation_axis_angle(axis_d, theta)
#             T2 = trimesh.transformations.translation_matrix(axis_o + axis_d * dist)
#             box_anim.apply_transform(T1 @ R @ T2)

#         elif jtype == 5:  # continuous
#             theta = np.pi / 4
#             T = trimesh.transformations.translation_matrix(-axis_o)
#             R_3 = get_rotation_axis_angle(axis_d, theta)
#             R = np.eye(4, dtype=np.float32)
#             R[:3, :3] = R_3
#             T_inv = trimesh.transformations.translation_matrix(axis_o)
#             box_anim.apply_transform(T @ R @ T_inv)

#     return box, box_anim

# # def get_bbox_mesh_pair(center, size, radius=0.01, jtype=None, jrange=None, axis_d=None, axis_o=None):
# #     '''
# #     Function to get the bounding box mesh pair

# #     Args:
# #     - center (np.array): bounding box center
# #     - size (np.array): bounding box size
# #     - radius (float): radius of the cylinder
# #     - jtype (int): joint type
# #     - jrange (list): joint range
# #     - axis_d (np.array): axis direction
# #     - axis_o (np.array): axis origin

# #     Returns:
# #     - trimesh_box (trimesh object): trimesh object for the bbox at resting state
# #     - trimesh_box_anim (trimesh object): trimesh object for the bbox at opening state
# #     '''

# #     size = np.clip(size, a_max=3, a_min=0.005)
# #     center = np.clip(center, a_max=3, a_min=-3)

# #     line_box = o3d.geometry.TriangleMesh()
# #     z_cylinder = o3d.geometry.TriangleMesh.create_cylinder(radius=radius, height=size[2])
# #     y_cylinder = o3d.geometry.TriangleMesh.create_cylinder(radius=radius, height=size[1])
# #     R_y = get_rotation_axis_angle(np.array([1., 0., 0.], dtype=np.float32), np.pi / 2)
# #     y_cylinder.rotate(R_y, center=(0, 0, 0))
# #     x_cylinder = o3d.geometry.TriangleMesh.create_cylinder(radius=radius, height=size[0])
# #     R_x = get_rotation_axis_angle(np.array([0., 1., 0.], dtype=np.float32), np.pi / 2)
# #     x_cylinder.rotate(R_x, center=(0, 0, 0))
    
    
# #     z1 = deepcopy(z_cylinder)
# #     z1.translate(np.array([-size[0] / 2, size[1] / 2, 0.], dtype=np.float32))
# #     line_box += z1.translate(center[:3])
# #     z2 = deepcopy(z_cylinder)
# #     z2.translate(np.array([size[0] / 2, size[1] / 2, 0.], dtype=np.float32))
# #     line_box += z2.translate(center[:3])
# #     z3 = deepcopy(z_cylinder)
# #     z3.translate(np.array([-size[0] / 2, -size[1] / 2, 0.], dtype=np.float32))
# #     line_box += z3.translate(center[:3])
# #     z4 = deepcopy(z_cylinder)
# #     z4.translate(np.array([size[0] / 2, -size[1] / 2, 0.], dtype=np.float32))
# #     line_box += z4.translate(center[:3])
    
# #     y1 = deepcopy(y_cylinder)
# #     y1.translate(np.array([-size[0] / 2, 0., size[2] / 2], dtype=np.float32))
# #     line_box += y1.translate(center[:3])
# #     y2 = deepcopy(y_cylinder)
# #     y2.translate(np.array([size[0] / 2, 0., size[2] / 2], dtype=np.float32))
# #     line_box += y2.translate(center[:3])
# #     y3 = deepcopy(y_cylinder)
# #     y3.translate(np.array([-size[0] / 2, 0., -size[2] / 2], dtype=np.float32))
# #     line_box += y3.translate(center[:3])
# #     y4 = deepcopy(y_cylinder)
# #     y4.translate(np.array([size[0] / 2, 0., -size[2] / 2], dtype=np.float32))
# #     line_box += y4.translate(center[:3])
    
# #     x1 = deepcopy(x_cylinder)
# #     x1.translate(np.array([0., -size[1] / 2, size[2] / 2], dtype=np.float32))
# #     line_box += x1.translate(center[:3])
# #     x2 = deepcopy(x_cylinder)
# #     x2.translate(np.array([0., size[1] / 2, size[2] / 2], dtype=np.float32))
# #     line_box += x2.translate(center[:3])
# #     x3 = deepcopy(x_cylinder)
# #     x3.translate(np.array([0., -size[1] / 2, -size[2] / 2], dtype=np.float32))
# #     line_box += x3.translate(center[:3])
# #     x4 = deepcopy(x_cylinder)
# #     x4.translate(np.array([0., size[1] / 2, -size[2] / 2]))
# #     line_box += x4.translate(center[:3])

# #     # transform
# #     line_box_anim = deepcopy(line_box)
# #     if jtype == 2: # revolute
# #         theta = np.deg2rad(jrange[1])
# #         line_box_anim.translate(-axis_o)
# #         R = get_rotation_axis_angle(axis_d, theta)
# #         line_box_anim.rotate(R, center=(0, 0, 0))
# #         line_box_anim.translate(axis_o)
# #     elif jtype == 3: # prismatic
# #         dist = np.array(jrange[1], dtype=np.float32)
# #         line_box_anim.translate(axis_d * dist)
# #     elif jtype == 4: # screw
# #         dist = np.array(jrange[1], dtype=np.float32)
# #         theta = 0.25 * np.pi
# #         R = get_rotation_axis_angle(axis_d, theta)
# #         line_box_anim.translate(-axis_o)
# #         line_box_anim.rotate(R, center=(0, 0, 0))
# #         line_box_anim.translate(axis_o)
# #         line_box_anim.translate(axis_d * dist)
# #     elif jtype == 5: # continuous
# #         theta = 0.25 * np.pi
# #         R = get_rotation_axis_angle(axis_d, theta)
# #         line_box_anim.translate(-axis_o)
# #         line_box_anim.rotate(R, center=(0, 0, 0))
# #         line_box_anim.translate(axis_o)
    
# #     vertices = np.asarray(line_box.vertices)
# #     faces = np.asarray(line_box.triangles)
# #     trimesh_box = trimesh.Trimesh(vertices=vertices, faces=faces)
# #     trimesh_box.visual.vertex_colors = np.array([0.0, 1.0, 1.0, 1.0], dtype=np.float32)
    
# #     vertices_anim = np.asarray(line_box_anim.vertices)
# #     faces_anim = np.asarray(line_box_anim.triangles)
# #     trimesh_box_anim = trimesh.Trimesh(vertices=vertices_anim, faces=faces_anim)
# #     trimesh_box_anim.visual.vertex_colors = np.array([0.0, 1.0, 1.0, 1.0], dtype=np.float32)
    
# #     return trimesh_box, trimesh_box_anim


# def get_color_from_palette(palette, idx):
#     '''
#     Function to get the color from the palette

#     Args:
#     - palette (list): list of color reference
#     - idx (int): index of the color

#     Returns:
#     - color (np.array): color in the index of idx
#     '''
#     ref = palette[idx % len(palette)]
#     ref_list = [int(i) for i in ref[4:-1].split(',')]
#     if idx % len(palette) == 0:
#         ref_list.append(120)
#     else:
#         ref_list.append(255)
#     color = np.array([ref_list], dtype=np.float32) / 255.
#     return color



# def render_anim_parts(aabbs, axiss, resolution=256):
#     '''
#     Function to render the 3D bounding boxes and axes in the scene

#     Args:
#         aabbs: list of trimesh objects for the bounding box of each part
#         axiss: list of trimesh objects for the axis of each part
#         resolution: resolution of the rendered image

#     Returns:
#         color_img: rendered image
#     '''
#     n_parts = len(aabbs)
#     # build mesh for each 3D bounding box
#     scene = pyrender.Scene()
#     for i in range(n_parts):
#         scene.add(aabbs[i])
#         if axiss[i] is not None:
#             scene.add(axiss[i])

#     # Add light to the scene
#     scene.ambient_light = np.full(shape=3, fill_value=1.5, dtype=np.float32)
#     light = pyrender.DirectionalLight(color=np.ones(2), intensity=5.0)
    
#     # Add camera to the scene
#     pose = get_camera_pose(eye=np.array([1.5, 1.2, 4.5]), look_at=np.array([0, 0, 0]), up=np.array([0, 1, 0]))
#     camera = pyrender.PerspectiveCamera(yfov=np.pi / 5.0, aspectRatio=1.0)
#     scene.add(light, pose=pose)
#     scene.add(camera, pose=pose)

#     # Offscreen Rendering
#     offscreen_renderer = pyrender.OffscreenRenderer(resolution, resolution)

#     # Render the scene
#     color_img, _ = offscreen_renderer.render(scene)

#     # Cleanup
#     offscreen_renderer.delete()
#     scene.clear()
#     return color_img


# def draw_boxes_axiss_anim(aabbs_0, aabbs_1, axiss, mode='graph', resolution=256, types=None):
#     '''
#     Function to draw the 3D bounding boxes and axes of the two frames

#     Args:
#         aabbs_0: list of trimesh objects for the bounding box of each part in the resting state
#         aabbs_1: list of trimesh objects for the bounding box of each part in the open state
#         axiss: list of trimesh objects for the axis of each part
#         mode: 
#             'graph'     using palette corresponding to graph node, 
#             'jtype'     using palette corresponding to joint type, 
#             'semantic'  using palette corresponding to semantic label
#         resolution: resolution of the rendered image
#         types: ids corresponding to each joint type or semantic label, if mode is 'jtype' or 'semantic'
#     '''
#     n_parts = len(aabbs_0)
#     ren_aabbs_0 = []
#     ren_aabbs_1 = []
#     ren_axiss = []
#     if mode == 'graph':
#         palette = graph_color_ref
#         # Add meshes to the scene
#         for i in range(n_parts):
#             color = get_color_from_palette(palette, i)
#             aabb_0 = pyrender.Mesh.from_trimesh(aabbs_0[i], smooth=False)
#             aabb_0.primitives[0].color_0 = color.repeat(aabb_0.primitives[0].positions.shape[0], axis=0)
#             ren_aabbs_0.append(aabb_0)
#             aabb_1 = pyrender.Mesh.from_trimesh(aabbs_1[i], smooth=False)
#             aabb_1.primitives[0].color_0 = color.repeat(aabb_1.primitives[0].positions.shape[0], axis=0)
#             ren_aabbs_1.append(aabb_1)
#             if axiss[i] is not None:
#                 axis = pyrender.Mesh.from_trimesh(axiss[i], smooth=False)
#                 axis.primitives[0].color_0 = color.repeat(axis.primitives[0].positions.shape[0], axis=0)
#                 ren_axiss.append(axis)
#             else:
#                 ren_axiss.append(None)     
#     elif mode == 'jtype' or mode == 'semantic':
#         assert types is not None
#         palette = joint_color_ref if mode == 'jtype' else semantic_color_ref
#         # Add meshes to the scene
#         for i in range(n_parts):
#             color = get_color_from_palette(palette, types[i])
#             aabb_0 = pyrender.Mesh.from_trimesh(aabbs_0[i], smooth=False)
#             aabb_0.primitives[0].color_0 = color.repeat(aabb_0.primitives[0].positions.shape[0], axis=0)
#             ren_aabbs_0.append(aabb_0)
#             aabb_1 = pyrender.Mesh.from_trimesh(aabbs_1[i], smooth=False)
#             aabb_1.primitives[0].color_0 = color.repeat(aabb_1.primitives[0].positions.shape[0], axis=0)
#             ren_aabbs_1.append(aabb_1)

#             if axiss[i] is not None:
#                 axis = pyrender.Mesh.from_trimesh(axiss[i], smooth=False)
#                 ren_axiss.append(axis)
#             else:
#                 ren_axiss.append(None)
#     else:
#         raise ValueError('mode must be either graph or type')

#     img0 = render_anim_parts(ren_aabbs_0, ren_axiss, resolution=resolution)
#     img1 = render_anim_parts(ren_aabbs_1, ren_axiss, resolution=resolution)
#     return np.concatenate([img0, img1], axis=1)

# def prepare_meshes(info_dict):
#     """
#     Function to prepare the bbox and axis meshes for visualization

#     Args:
#     - info_dict (dict): output json containing the graph information
#     """
#     from my_utils.refs import joint_ref, sem_ref
#     tree = info_dict["diffuse_tree"]
#     bbox_0, bbox_1, axiss, labels, jtypes = [], [], [], [], []
#     root_id = 0
#     # get root id
#     for node in tree:
#         if node["parent"] == -1:
#             root_id = node["id"]
#     for node in tree:
#         # retrieve info
#         box_cen = np.array(node["aabb"]["center"], dtype=np.float32)
#         box_size = np.array(node["aabb"]["size"], dtype=np.float32)
#         axis_d = np.array(node["joint"]["axis"]["direction"], dtype=np.float32)
#         axis_o = np.array(node["joint"]["axis"]["origin"], dtype=np.float32)
#         jtype = joint_ref["fwd"][node["joint"]["type"]]
#         # construct meshes for bbox in two states (closed and fully open)
#         if node["id"] == root_id or node["parent"] == root_id:  # use the joint info directly
#             bb_0, bb_1 = get_colored_box(
#                 box_cen,
#                 box_size,
#                 jtype=jtype,
#                 jrange= node["joint"]["range"],
#                 axis_d=axis_d,
#                 axis_o=axis_o,
#             )
#         else: # use the parent joint info
#             parent_id = node["parent"]
#             bb_0, bb_1 = get_colored_box(
#                 box_cen,
#                 box_size,
#                 jtype=joint_ref["fwd"][tree[parent_id]["joint"]["type"]],
#                 jrange=tree[parent_id]["joint"]["range"],
#                 axis_d=np.array(tree[parent_id]["joint"]["axis"]["direction"], dtype=np.float32),
#                 axis_o=np.array(tree[parent_id]["joint"]["axis"]["origin"], dtype=np.float32),
#             )
#         # construct mesh for joint axis
#         axis_mesh = get_axis_mesh(axis_d, axis_o, box_cen, node["joint"]["type"])
#         # append
#         bbox_0.append(bb_0)
#         bbox_1.append(bb_1)
#         axiss.append(axis_mesh)
#         labels.append(sem_ref["fwd"][node["name"]])
#         jtypes.append(jtype)

#     return {
#         "bbox_0": bbox_0,
#         "bbox_1": bbox_1,
#         "axiss": axiss,
#         "labels": labels,
#         "jtypes": jtypes,
#     }