mujoco_worldgen/util/obj_util.py

import numpy as np
from collections import OrderedDict
# TODO: Write more comments.

from mujoco_worldgen.util.types import accepts, returns, maybe


def get_camera_xyaxes(camera_pos, target_pos):
    '''
    Calculate the "xyaxis" frame orientation for a camera,
        given the camera position and target position.
    Returns a 6-vector of the x axis and y axis of the camera frame.
        See http://www.mujoco.org/book/modeling.html#COrientation "xyaxes"
    '''
    camera_pos = np.array(camera_pos)
    target_pos = np.array(target_pos)
    assert camera_pos.shape == (
        3,), "Bad camera position {}".format(camera_pos)
    assert target_pos.shape == (
        3,), "Bad target position {}".format(target_pos)
    vector = target_pos - camera_pos
    cross = np.cross(vector, np.array([0, 0, 1]))
    cross2 = np.cross(cross, vector)
    return np.concatenate((cross, cross2))


@accepts(OrderedDict, str, maybe(np.ndarray))
def add_annotation_bound(xml_dict, annotation_name, bound):
    '''
    Add an annotation bounding box to and XML dictionary.
    Annotation name will be "annotation:" + annotation_name
    Bound is given as a 2 x 3 np.ndarray, and represents:
        [[min_x, min_y, min_z], [max_x, max_y, max_z]]
    '''
    if bound is None:
        return xml_dict  # Nothing to do here
    assert bound.shape == (2, 3), "Bound must be 2 x 3 (see docstring)."
    assert 'worldbody' in xml_dict, "XML must have worldbody"
    worldbody = xml_dict['worldbody']
    assert 'body' in worldbody, "XML worldbody must have bodies"
    name = 'annotation:' + annotation_name
    # Remove old annotations with the same name before inserting new annotation
    bodies = [body for body in worldbody['body'] if body.get('@name') != name]
    rgba = np.random.uniform(size=4)
    rgba[3] = 0.1  # annotation is almost transparent.
    geom = OrderedDict([('@conaffinity', 0),
                        ('@contype', 0),
                        ('@mass', 0.0),
                        ('@pos', np.zeros(3)),
                        ('@rgba', rgba),
                        ('@size', (bound[1] - bound[0]) / 2),  # halfsize
                        ('@type', 'box')])
    annotation = OrderedDict([('@name', name),
                              ('@pos', bound.mean(axis=0)),  # center pos
                              ('geom', [geom])])
    bodies.append(annotation)
    worldbody['body'] = bodies
    print('adding annotation bound (size)', name, bound[1] - bound[0])


@accepts(OrderedDict)
@returns(OrderedDict)
def get_xml_meshes(xml_dict):
    ''' Get dictionary of all the mesh names -> filenames in a parsed XML. '''
    meshes = OrderedDict()
    for mesh in xml_dict.get('asset', {}).get('mesh', []):
        assert '@name' in mesh, "Mesh missing name: {}".format(mesh)
        assert '@file' in mesh, "Mesh missing file: {}".format(mesh)
        scale = np.ones(3)
        if "@scale" in mesh:
            scale = mesh["@scale"]
        meshes[mesh['@name']] = (mesh['@file'], scale)
    return meshes


@accepts(OrderedDict, str)
def recursive_rename(xml_dict, prefix):
    attrs = ["@name", "@joint", "@jointinparent", "@class", "@source",
             "@target", "@childclass", "@body1", "@body2", "@mesh",
             "@joint1", "@joint2", "@geom", "@geom1", "@geom2", "@site",
             "@material", "@texture", "@tendon", "@sidesite", "@actuator"]
    names = ["geom", "joint", "jointinparent", "body", "motor", "freejoint", "general",
             "position", "default", "weld", "exclude", "mesh",
             "site", "pair", "jointpos", "touch", "texture", "material",
             "fixed", "spatial", "motor", "actuatorfrc"]
    if not isinstance(xml_dict, OrderedDict):
        return
    for key in list(xml_dict.keys()):
        value_dict = xml_dict[key]
        if isinstance(value_dict, OrderedDict):
            value_dict = [value_dict]
        if key in names:
            assert isinstance(
                value_dict, list), "Invalid type for value {}".format(value_dict)
            for value in value_dict:
                for attr in list(value.keys()):
                    if attr in attrs:
                        if not value[attr].startswith(prefix + ':'):
                            value[attr] = prefix + ':' + value[attr]
        if isinstance(value_dict, list):
            for value in value_dict:
                recursive_rename(value, prefix)


def establish_size(min_size, max_size):
    if isinstance(min_size, (float, int)):
        min_size = np.ones(3) * float(min_size)
    if isinstance(max_size, (float, int)):
        max_size = np.ones(3) * float(max_size)
    if max_size is None and min_size is not None:
        max_size = min_size
    if max_size is None and min_size is None:
        min_size = np.ones(3) * 0.1
        max_size = np.ones(3) * 0.1
    if isinstance(min_size, (list, tuple)):
        min_size = np.array(min_size, dtype=np.float64)
    if isinstance(max_size, (list, tuple)):
        max_size = np.array(max_size, dtype=np.float64)
    assert(isinstance(min_size[0], float))
    assert(isinstance(max_size[0], float))
    for i in range(3):
        assert(max_size[i] >= min_size[i])
    return min_size, max_size


def get_name_index(name_indexes, name):
    '''
    Update the name index and return new name
    name - name to look up index for, e.g. "geom"
    name_indexes - dictionary to keep track of names, e.g.
        {'geom': 4} means there are 4 geom objects, and the next
        geom object should be called "geom4" and the dictionary updated
        to be {'geom': 5}
    Returns name with index attached, e.g. "geom4"
    '''
    if name not in name_indexes:
        name_indexes[name] = 0
    result = '%s%d' % (name, name_indexes[name])
    name_indexes[name] += 1
    return result


def get_body_xml_node(name, use_joints=False):
    '''
    Build a body XML dict for use in object models.
        name - name for the body (should be unique in the model, e.g. "geom4")
        joints - if True, add 6 degrees of freedom joints (slide, hinge)
    Returns named XML body node.
    '''
    body = OrderedDict()
    body['@name'] = name
    body['@pos'] = np.zeros(3)

    if use_joints:
        joints = []
        for axis_type in ('slide', 'hinge'):
            for i, axis in enumerate(np.eye(3)):
                joint = OrderedDict()
                joint['@name'] = "%s:%s%d" % (name, axis_type, i)
                joint['@axis'] = axis
                joint['@type'] = axis_type
                joint['@damping'] = 0.01
                joint['@pos'] = np.zeros(3)
                joints.append(joint)
        body['joint'] = joints
    return body


@accepts(np.ndarray)
@returns(int)
def get_axis_index(axis):
    '''
    Returns axis index from a string:
    # return 0 for axis = 1 0 0
    # return 1 for axis = 0 1 0
    # return 2 for axis = 0 0 1
    '''
    assert axis.shape[0] == 3
    for i in range(3):
        if axis[i] != 0:
            return i
    assert False, "axis should be of a form (1 0 0), or (0 1 0), or (0 0 1)." \
                  "Current axis = %s, it's not. Failing." % str(axis)