OpenMotionLab/MotionGPT

1"""Examples of using pyrender for viewing and offscreen rendering.2"""3import pyglet4pyglet.options['shadow_window'] = False5import os6import numpy as np7import trimesh8 9from pyrender import PerspectiveCamera,\10                     DirectionalLight, SpotLight, PointLight,\11                     MetallicRoughnessMaterial,\12                     Primitive, Mesh, Node, Scene,\13                     Viewer, OffscreenRenderer, RenderFlags14 15#==============================================================================16# Mesh creation17#==============================================================================18 19#------------------------------------------------------------------------------20# Creating textured meshes from trimeshes21#------------------------------------------------------------------------------22 23# Fuze trimesh24fuze_trimesh = trimesh.load('./models/fuze.obj')25fuze_mesh = Mesh.from_trimesh(fuze_trimesh)26 27# Drill trimesh28drill_trimesh = trimesh.load('./models/drill.obj')29drill_mesh = Mesh.from_trimesh(drill_trimesh)30drill_pose = np.eye(4)31drill_pose[0,3] = 0.132drill_pose[2,3] = -np.min(drill_trimesh.vertices[:,2])33 34# Wood trimesh35wood_trimesh = trimesh.load('./models/wood.obj')36wood_mesh = Mesh.from_trimesh(wood_trimesh)37 38# Water bottle trimesh39bottle_gltf = trimesh.load('./models/WaterBottle.glb')40bottle_trimesh = bottle_gltf.geometry[list(bottle_gltf.geometry.keys())[0]]41bottle_mesh = Mesh.from_trimesh(bottle_trimesh)42bottle_pose = np.array([43    [1.0, 0.0,  0.0, 0.1],44    [0.0, 0.0, -1.0, -0.16],45    [0.0, 1.0,  0.0, 0.13],46    [0.0, 0.0,  0.0, 1.0],47])48 49#------------------------------------------------------------------------------50# Creating meshes with per-vertex colors51#------------------------------------------------------------------------------52boxv_trimesh = trimesh.creation.box(extents=0.1*np.ones(3))53boxv_vertex_colors = np.random.uniform(size=(boxv_trimesh.vertices.shape))54boxv_trimesh.visual.vertex_colors = boxv_vertex_colors55boxv_mesh = Mesh.from_trimesh(boxv_trimesh, smooth=False)56 57#------------------------------------------------------------------------------58# Creating meshes with per-face colors59#------------------------------------------------------------------------------60boxf_trimesh = trimesh.creation.box(extents=0.1*np.ones(3))61boxf_face_colors = np.random.uniform(size=boxf_trimesh.faces.shape)62boxf_trimesh.visual.face_colors = boxf_face_colors63boxf_mesh = Mesh.from_trimesh(boxf_trimesh, smooth=False)64 65#------------------------------------------------------------------------------66# Creating meshes from point clouds67#------------------------------------------------------------------------------68points = trimesh.creation.icosphere(radius=0.05).vertices69point_colors = np.random.uniform(size=points.shape)70points_mesh = Mesh.from_points(points, colors=point_colors)71 72#==============================================================================73# Light creation74#==============================================================================75 76direc_l = DirectionalLight(color=np.ones(3), intensity=1.0)77spot_l = SpotLight(color=np.ones(3), intensity=10.0,78                   innerConeAngle=np.pi/16, outerConeAngle=np.pi/6)79point_l = PointLight(color=np.ones(3), intensity=10.0)80 81#==============================================================================82# Camera creation83#==============================================================================84 85cam = PerspectiveCamera(yfov=(np.pi / 3.0))86cam_pose = np.array([87    [0.0,  -np.sqrt(2)/2, np.sqrt(2)/2, 0.5],88    [1.0, 0.0,           0.0,           0.0],89    [0.0,  np.sqrt(2)/2,  np.sqrt(2)/2, 0.4],90    [0.0,  0.0,           0.0,          1.0]91])92 93#==============================================================================94# Scene creation95#==============================================================================96 97scene = Scene(ambient_light=np.array([0.02, 0.02, 0.02, 1.0]))98 99#==============================================================================100# Adding objects to the scene101#==============================================================================102 103#------------------------------------------------------------------------------104# By manually creating nodes105#------------------------------------------------------------------------------106fuze_node = Node(mesh=fuze_mesh, translation=np.array([0.1, 0.15, -np.min(fuze_trimesh.vertices[:,2])]))107scene.add_node(fuze_node)108boxv_node = Node(mesh=boxv_mesh, translation=np.array([-0.1, 0.10, 0.05]))109scene.add_node(boxv_node)110boxf_node = Node(mesh=boxf_mesh, translation=np.array([-0.1, -0.10, 0.05]))111scene.add_node(boxf_node)112 113#------------------------------------------------------------------------------114# By using the add() utility function115#------------------------------------------------------------------------------116drill_node = scene.add(drill_mesh, pose=drill_pose)117bottle_node = scene.add(bottle_mesh, pose=bottle_pose)118wood_node = scene.add(wood_mesh)119direc_l_node = scene.add(direc_l, pose=cam_pose)120spot_l_node = scene.add(spot_l, pose=cam_pose)121 122#==============================================================================123# Using the viewer with a default camera124#==============================================================================125 126v = Viewer(scene, shadows=True)127 128#==============================================================================129# Using the viewer with a pre-specified camera130#==============================================================================131cam_node = scene.add(cam, pose=cam_pose)132v = Viewer(scene, central_node=drill_node)133 134#==============================================================================135# Rendering offscreen from that camera136#==============================================================================137 138r = OffscreenRenderer(viewport_width=640*2, viewport_height=480*2)139color, depth = r.render(scene)140 141import matplotlib.pyplot as plt142plt.figure()143plt.imshow(color)144plt.show()145 146#==============================================================================147# Segmask rendering148#==============================================================================149 150nm = {node: 20*(i + 1) for i, node in enumerate(scene.mesh_nodes)}151seg = r.render(scene, RenderFlags.SEG, nm)[0]152plt.figure()153plt.imshow(seg)154plt.show()155 156r.delete()157 158