OpenMotionLab/MotionGPT

1import numpy as np2import trimesh3 4from pyrender import (OffscreenRenderer, PerspectiveCamera, DirectionalLight,5                      SpotLight, Mesh, Node, Scene)6 7 8def test_offscreen_renderer(tmpdir):9 10    # Fuze trimesh11    fuze_trimesh = trimesh.load('examples/models/fuze.obj')12    fuze_mesh = Mesh.from_trimesh(fuze_trimesh)13 14    # Drill trimesh15    drill_trimesh = trimesh.load('examples/models/drill.obj')16    drill_mesh = Mesh.from_trimesh(drill_trimesh)17    drill_pose = np.eye(4)18    drill_pose[0,3] = 0.119    drill_pose[2,3] = -np.min(drill_trimesh.vertices[:,2])20 21    # Wood trimesh22    wood_trimesh = trimesh.load('examples/models/wood.obj')23    wood_mesh = Mesh.from_trimesh(wood_trimesh)24 25    # Water bottle trimesh26    bottle_gltf = trimesh.load('examples/models/WaterBottle.glb')27    bottle_trimesh = bottle_gltf.geometry[list(bottle_gltf.geometry.keys())[0]]28    bottle_mesh = Mesh.from_trimesh(bottle_trimesh)29    bottle_pose = np.array([30        [1.0, 0.0, 0.0, 0.1],31        [0.0, 0.0, -1.0, -0.16],32        [0.0, 1.0, 0.0, 0.13],33        [0.0, 0.0, 0.0, 1.0],34    ])35 36    boxv_trimesh = trimesh.creation.box(extents=0.1 * np.ones(3))37    boxv_vertex_colors = np.random.uniform(size=(boxv_trimesh.vertices.shape))38    boxv_trimesh.visual.vertex_colors = boxv_vertex_colors39    boxv_mesh = Mesh.from_trimesh(boxv_trimesh, smooth=False)40    boxf_trimesh = trimesh.creation.box(extents=0.1 * np.ones(3))41    boxf_face_colors = np.random.uniform(size=boxf_trimesh.faces.shape)42    boxf_trimesh.visual.face_colors = boxf_face_colors43    # Instanced44    poses = np.tile(np.eye(4), (2,1,1))45    poses[0,:3,3] = np.array([-0.1, -0.10, 0.05])46    poses[1,:3,3] = np.array([-0.15, -0.10, 0.05])47    boxf_mesh = Mesh.from_trimesh(boxf_trimesh, poses=poses, smooth=False)48 49    points = trimesh.creation.icosphere(radius=0.05).vertices50    point_colors = np.random.uniform(size=points.shape)51    points_mesh = Mesh.from_points(points, colors=point_colors)52 53    direc_l = DirectionalLight(color=np.ones(3), intensity=1.0)54    spot_l = SpotLight(color=np.ones(3), intensity=10.0,55                       innerConeAngle=np.pi / 16, outerConeAngle=np.pi / 6)56 57    cam = PerspectiveCamera(yfov=(np.pi / 3.0))58    cam_pose = np.array([59        [0.0, -np.sqrt(2) / 2, np.sqrt(2) / 2, 0.5],60        [1.0, 0.0, 0.0, 0.0],61        [0.0, np.sqrt(2) / 2, np.sqrt(2) / 2, 0.4],62        [0.0, 0.0, 0.0, 1.0]63    ])64 65    scene = Scene(ambient_light=np.array([0.02, 0.02, 0.02]))66 67    fuze_node = Node(mesh=fuze_mesh, translation=np.array([68        0.1, 0.15, -np.min(fuze_trimesh.vertices[:,2])69    ]))70    scene.add_node(fuze_node)71    boxv_node = Node(mesh=boxv_mesh, translation=np.array([-0.1, 0.10, 0.05]))72    scene.add_node(boxv_node)73    boxf_node = Node(mesh=boxf_mesh)74    scene.add_node(boxf_node)75 76    _ = scene.add(drill_mesh, pose=drill_pose)77    _ = scene.add(bottle_mesh, pose=bottle_pose)78    _ = scene.add(wood_mesh)79    _ = scene.add(direc_l, pose=cam_pose)80    _ = scene.add(spot_l, pose=cam_pose)81    _ = scene.add(points_mesh)82 83    _ = scene.add(cam, pose=cam_pose)84 85    r = OffscreenRenderer(viewport_width=640, viewport_height=480)86    color, depth = r.render(scene)87 88    assert color.shape == (480, 640, 3)89    assert depth.shape == (480, 640)90    assert np.max(depth.data) > 0.0591    assert np.count_nonzero(depth.data) > (0.2 * depth.size)92    r.delete()93