def get_all_vertices()

in robogym/envs/rearrange/common/utils.py [0:0]

• 28 lines of code
• 6 McCabe index (conditional complexity)

def get_all_vertices(sim, object_name, subdivide_threshold=None) -> np.ndarray:
    """
    Return all vertices for given object.
    :param sim: The MjSim instance.
    :param object_name: The object name.
    :param subdivide_threshold: If provided, subdivide mesh into smaller faces.
      See subdivide_mesh for detail of this parameter.
    :return: Array of all vertices for this object.
    """
    all_verts: List[np.ndarray] = []
    all_faces: List[Optional[np.ndarray]] = []

    object_rot_mat = quat2mat(
        quat_conjugate(mat2quat(sim.data.get_body_xmat(object_name)))
    )
    geom_ids = geom_ids_of_body(sim, object_name)

    for geom_id in geom_ids:
        pos = sim.model.geom_pos[geom_id]
        quat = quat_conjugate(sim.model.geom_quat[geom_id])
        mat = quat2mat(quat)

        # Get all vertices associated with the current geom.
        verts = get_geom_vertices(sim, geom_id)
        faces = get_geom_faces(sim, geom_id)

        # Translate from geom's to body's coordinate frame.
        geom_ref_verts = verts @ mat
        geom_ref_verts = pos + geom_ref_verts

        all_verts.append(geom_ref_verts)
        all_faces.append(faces)

    if subdivide_threshold is not None and all(f is not None for f in all_faces):
        # We can only subdivide mesh with faces.
        mesh = trimesh.util.concatenate(
            [
                trimesh.Trimesh(vertices=verts, faces=faces)
                for verts, faces in zip(all_verts, all_faces)
            ]
        )

        verts = subdivide_mesh(mesh.vertices, mesh.faces, subdivide_threshold)
    else:
        verts = np.concatenate(all_verts, axis=0)

    return verts @ object_rot_mat