OpenMotionLab/MotionGPT

1import numpy as np2import pytest3import trimesh4 5from pyrender import (Mesh, PerspectiveCamera, DirectionalLight,6                      SpotLight, PointLight, Scene, Node, OrthographicCamera)7 8 9def test_scenes():10 11    # Basics12    s = Scene()13    assert np.allclose(s.bg_color, np.ones(4))14    assert np.allclose(s.ambient_light, np.zeros(3))15    assert len(s.nodes) == 016    assert s.name is None17    s.name = 'asdf'18    s.bg_color = None19    s.ambient_light = None20    assert np.allclose(s.bg_color, np.ones(4))21    assert np.allclose(s.ambient_light, np.zeros(3))22 23    assert s.nodes == set()24    assert s.cameras == set()25    assert s.lights == set()26    assert s.point_lights == set()27    assert s.spot_lights == set()28    assert s.directional_lights == set()29    assert s.meshes == set()30    assert s.camera_nodes == set()31    assert s.light_nodes == set()32    assert s.point_light_nodes == set()33    assert s.spot_light_nodes == set()34    assert s.directional_light_nodes == set()35    assert s.mesh_nodes == set()36    assert s.main_camera_node is None37    assert np.all(s.bounds == 0)38    assert np.all(s.centroid == 0)39    assert np.all(s.extents == 0)40    assert np.all(s.scale == 0)41 42    # From trimesh scene43    tms = trimesh.load('tests/data/WaterBottle.glb')44    s = Scene.from_trimesh_scene(tms)45    assert len(s.meshes) == 146    assert len(s.mesh_nodes) == 147 48    # Test bg color formatting49    s = Scene(bg_color=[0, 1.0, 0])50    assert np.allclose(s.bg_color, np.array([0.0, 1.0, 0.0, 1.0]))51 52    # Test constructor for nodes53    n1 = Node()54    n2 = Node()55    n3 = Node()56    nodes = [n1, n2, n3]57    s = Scene(nodes=nodes)58    n1.children.append(n2)59    s = Scene(nodes=nodes)60    n3.children.append(n2)61    with pytest.raises(ValueError):62        s = Scene(nodes=nodes)63    n3.children = []64    n2.children.append(n3)65    n3.children.append(n2)66    with pytest.raises(ValueError):67        s = Scene(nodes=nodes)68 69    # Test node accessors70    n1 = Node()71    n2 = Node()72    n3 = Node()73    nodes = [n1, n2]74    s = Scene(nodes=nodes)75    assert s.has_node(n1)76    assert s.has_node(n2)77    assert not s.has_node(n3)78 79    # Test node poses80    for n in nodes:81        assert np.allclose(s.get_pose(n), np.eye(4))82    with pytest.raises(ValueError):83        s.get_pose(n3)84    with pytest.raises(ValueError):85        s.set_pose(n3, np.eye(4))86    tf = np.eye(4)87    tf[:3,3] = np.ones(3)88    s.set_pose(n1, tf)89    assert np.allclose(s.get_pose(n1), tf)90    assert np.allclose(s.get_pose(n2), np.eye(4))91 92    nodes = [n1, n2, n3]93    tf2 = np.eye(4)94    tf2[:3,:3] = np.diag([-1,-1,1])95    n1.children.append(n2)96    n1.matrix = tf97    n2.matrix = tf298    s = Scene(nodes=nodes)99    assert np.allclose(s.get_pose(n1), tf)100    assert np.allclose(s.get_pose(n2), tf.dot(tf2))101    assert np.allclose(s.get_pose(n3), np.eye(4))102 103    n1 = Node()104    n2 = Node()105    n3 = Node()106    n1.children.append(n2)107    s = Scene()108    s.add_node(n1)109    with pytest.raises(ValueError):110        s.add_node(n2)111    s.set_pose(n1, tf)112    assert np.allclose(s.get_pose(n1), tf)113    assert np.allclose(s.get_pose(n2), tf)114    s.set_pose(n2, tf2)115    assert np.allclose(s.get_pose(n2), tf.dot(tf2))116 117    # Test node removal118    n1 = Node()119    n2 = Node()120    n3 = Node()121    n1.children.append(n2)122    n2.children.append(n3)123    s = Scene(nodes=[n1, n2, n3])124    s.remove_node(n2)125    assert len(s.nodes) == 1126    assert n1 in s.nodes127    assert len(n1.children) == 0128    assert len(n2.children) == 1129    s.add_node(n2, parent_node=n1)130    assert len(n1.children) == 1131    n1.matrix = tf132    n3.matrix = tf2133    assert np.allclose(s.get_pose(n3), tf.dot(tf2))134 135    # Now test ADD function136    s = Scene()137    m = Mesh([], name='m')138    cp = PerspectiveCamera(yfov=2.0)139    co = OrthographicCamera(xmag=1.0, ymag=1.0)140    dl = DirectionalLight()141    pl = PointLight()142    sl = SpotLight()143 144    n1 = s.add(m, name='mn')145    assert n1.mesh == m146    assert len(s.nodes) == 1147    assert len(s.mesh_nodes) == 1148    assert n1 in s.mesh_nodes149    assert len(s.meshes) == 1150    assert m in s.meshes151    assert len(s.get_nodes(node=n2)) == 0152    n2 = s.add(m, pose=tf)153    assert len(s.nodes) == len(s.mesh_nodes) == 2154    assert len(s.meshes) == 1155    assert len(s.get_nodes(node=n1)) == 1156    assert len(s.get_nodes(node=n1, name='mn')) == 1157    assert len(s.get_nodes(name='mn')) == 1158    assert len(s.get_nodes(obj=m)) == 2159    assert len(s.get_nodes(obj=m, obj_name='m')) == 2160    assert len(s.get_nodes(obj=co)) == 0161    nsl = s.add(sl, name='sln')162    npl = s.add(pl, parent_name='sln')163    assert nsl.children[0] == npl164    ndl = s.add(dl, parent_node=npl)165    assert npl.children[0] == ndl166    nco = s.add(co)167    ncp = s.add(cp)168 169    assert len(s.light_nodes) == len(s.lights) == 3170    assert len(s.point_light_nodes) == len(s.point_lights) == 1171    assert npl in s.point_light_nodes172    assert len(s.spot_light_nodes) == len(s.spot_lights) == 1173    assert nsl in s.spot_light_nodes174    assert len(s.directional_light_nodes) == len(s.directional_lights) == 1175    assert ndl in s.directional_light_nodes176    assert len(s.cameras) == len(s.camera_nodes) == 2177    assert s.main_camera_node == nco178    s.main_camera_node = ncp179    s.remove_node(ncp)180    assert len(s.cameras) == len(s.camera_nodes) == 1181    assert s.main_camera_node == nco182    s.remove_node(n2)183    assert len(s.meshes) == 1184    s.remove_node(n1)185    assert len(s.meshes) == 0186    s.remove_node(nsl)187    assert len(s.lights) == 0188    s.remove_node(nco)189    assert s.main_camera_node is None190 191    s.add_node(n1)192    s.clear()193    assert len(s.nodes) == 0194 195    # Trigger final errors196    with pytest.raises(ValueError):197        s.main_camera_node = None198    with pytest.raises(ValueError):199        s.main_camera_node = ncp200    with pytest.raises(ValueError):201        s.add(m, parent_node=n1)202    with pytest.raises(ValueError):203        s.add(m, name='asdf')204        s.add(m, name='asdf')205        s.add(m, parent_name='asdf')206    with pytest.raises(ValueError):207        s.add(m, parent_name='asfd')208    with pytest.raises(TypeError):209        s.add(None)210 211    s.clear()212    # Test bounds213    m1 = Mesh.from_trimesh(trimesh.creation.box())214    m2 = Mesh.from_trimesh(trimesh.creation.box())215    m3 = Mesh.from_trimesh(trimesh.creation.box())216    n1 = Node(mesh=m1)217    n2 = Node(mesh=m2, translation=[1.0, 0.0, 0.0])218    n3 = Node(mesh=m3, translation=[0.5, 0.0, 1.0])219    s.add_node(n1)220    s.add_node(n2)221    s.add_node(n3)222    assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [1.5, 0.5, 1.5]])223    s.clear()224    s.add_node(n1)225    s.add_node(n2, parent_node=n1)226    s.add_node(n3, parent_node=n2)227    assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [2.0, 0.5, 1.5]])228    tf = np.eye(4)229    tf[:3,3] = np.ones(3)230    s.set_pose(n3, tf)231    assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [2.5, 1.5, 1.5]])232    s.remove_node(n2)233    assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]])234    s.clear()235    assert np.allclose(s.bounds, 0.0)236