Source code for rofunc.utils.robolab.formatter.mjcf_parser.skin

# Copyright 2020 The dm_control Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""Utilities for parsing and writing MuJoCo skin files.

The file format is described at http://mujoco.org/book/XMLreference.html#skin.
"""

import collections
import io
import struct

import numpy as np

MAX_BODY_NAME_LENGTH = 40

Skin = collections.namedtuple(
    'Skin', ('vertices', 'texcoords', 'faces', 'bones'))

Bone = collections.namedtuple(
    'Bone', ('body', 'bindpos', 'bindquat', 'vertex_ids', 'vertex_weights'))


def parse(contents, body_getter):
    """Parses the contents of a MuJoCo skin file.

    Args:
      contents: a bytes-like object containing the contents of a skin file.
      body_getter: a callable that takes a string and returns the `mjcf.Element`
        instance of a MuJoCo body of the specified name.

    Returns:
      A `Skin` named tuple.
    """
    f = io.BytesIO(contents)
    nvertex, ntexcoord, nface, nbone = struct.unpack('<iiii', f.read(4 * 4))
    vertices = np.frombuffer(f.read(4 * (3 * nvertex)), dtype='<f4').reshape(-1, 3)
    texcoords = np.frombuffer(f.read(4 * (2 * ntexcoord)), dtype='<f4').reshape(-1, 2)
    faces = np.frombuffer(f.read(4 * (3 * nface)), dtype='<i4').reshape(-1, 3)
    bones = []
    for _ in range(nbone):
        body_name = f.read(MAX_BODY_NAME_LENGTH).decode().split('\0')[0]
        body = lambda body_name=body_name: body_getter(body_name)
        bindpos = np.asarray(struct.unpack('<fff', f.read(4 * 3)), dtype=float)
        bindquat = np.asarray(struct.unpack('<ffff', f.read(4 * 4)), dtype=float)
        vertex_count = struct.unpack('<i', f.read(4))[0]
        vertex_ids = np.frombuffer(f.read(4 * vertex_count), dtype='<i4')
        vertex_weights = np.frombuffer(f.read(4 * vertex_count), dtype='<f4')
        bones.append(Bone(body=body, bindpos=bindpos, bindquat=bindquat,
                          vertex_ids=vertex_ids, vertex_weights=vertex_weights))

    return Skin(vertices=vertices, texcoords=texcoords, faces=faces, bones=bones)


def serialize(skin):
    """Serializes a `Skin` named tuple into the contents of a MuJoCo skin file.

    Args:
      skin: a `Skin` named tuple.

    Returns:
      A `bytes` object representing the content of a MuJoCo skin file.
    """
    out = io.BytesIO()
    nvertex = len(skin.vertices)
    ntexcoord = len(skin.texcoords)
    nface = len(skin.faces)
    nbone = len(skin.bones)
    out.write(struct.pack('<iiii', nvertex, ntexcoord, nface, nbone))
    out.write(skin.vertices.astype('<f4').tobytes())
    out.write(skin.texcoords.astype('<f4').tobytes())
    out.write(skin.faces.astype('<i4').tobytes())
    for bone in skin.bones:
        body_bytes = bone.body().full_identifier.encode('utf-8')
        if len(body_bytes) > MAX_BODY_NAME_LENGTH:
            raise ValueError(
                'body name is longer than  permitted by the skin file format '
                '(maximum 40): {:r}'.format(body_bytes))
        out.write(body_bytes)
        out.write(b'\0' * (MAX_BODY_NAME_LENGTH - len(body_bytes)))
        out.write(bone.bindpos.astype('<f4').tobytes())
        out.write(bone.bindquat.astype('<f4').tobytes())
        assert len(bone.vertex_ids) == len(bone.vertex_weights)
        out.write(struct.pack('<i', len(bone.vertex_ids)))
        out.write(bone.vertex_ids.astype('<i4').tobytes())
        out.write(bone.vertex_weights.astype('<f4').tobytes())
    return out.getvalue()