pymo/writers.py

import numpy as np
import pandas as pd


class BVHWriter:
    def __init__(self):
        pass

    def write(self, X, ofile, framerate=-1, start=0, stop=-1):
        # Writing the skeleton info
        ofile.write("HIERARCHY\n")

        self.motions_ = []
        self._printJoint(X, X.root_name, 0, ofile)

        if stop > 0:
            nframes = stop - start
        else:
            nframes = X.values.shape[0]
            stop = X.values.shape[0]

        # Writing the motion header
        ofile.write("MOTION\n")
        ofile.write("Frames: %d\n" % nframes)

        if framerate > 0:
            ofile.write("Frame Time: %f\n" % float(1.0 / framerate))
        else:
            ofile.write("Frame Time: %f\n" % X.framerate)

        # Writing the data
        self.motions_ = np.asarray(self.motions_).T
        lines = [" ".join(item) for item in self.motions_[start:stop].astype(str)]
        ofile.write("".join("%s\n" % l for l in lines))

    def _printJoint(self, X, joint, tab, ofile):
        if X.skeleton[joint]["parent"] == None:
            ofile.write("ROOT %s\n" % joint)
        elif len(X.skeleton[joint]["children"]) > 0:
            ofile.write("%sJOINT %s\n" % ("\t" * (tab), joint))
        else:
            ofile.write("%sEnd site\n" % ("\t" * (tab)))

        ofile.write("%s{\n" % ("\t" * (tab)))

        ofile.write(
            "%sOFFSET %3.5f %3.5f %3.5f\n"
            % (
                "\t" * (tab + 1),
                X.skeleton[joint]["offsets"][0],
                X.skeleton[joint]["offsets"][1],
                X.skeleton[joint]["offsets"][2],
            )
        )
        rot_order = X.skeleton[joint]["order"]

        # print("rot_order = " + rot_order)
        channels = X.skeleton[joint]["channels"]
        rot = [c for c in channels if ("rotation" in c)]
        pos = [c for c in channels if ("position" in c)]

        n_channels = len(rot) + len(pos)
        ch_str = ""
        if n_channels > 0:
            for ci in range(len(pos)):
                cn = pos[ci]
                self.motions_.append(np.asarray(X.values["%s_%s" % (joint, cn)].values))
                ch_str = ch_str + " " + cn
            for ci in range(len(rot)):
                cn = "%srotation" % (rot_order[ci])
                self.motions_.append(np.asarray(X.values["%s_%s" % (joint, cn)].values))
                ch_str = ch_str + " " + cn
        if len(X.skeleton[joint]["children"]) > 0:
            # ch_str = ''.join(' %s'*n_channels%tuple(channels))
            ofile.write("%sCHANNELS %d%s\n" % ("\t" * (tab + 1), n_channels, ch_str))

            for c in X.skeleton[joint]["children"]:
                self._printJoint(X, c, tab + 1, ofile)

        ofile.write("%s}\n" % ("\t" * (tab)))