# 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()