UPC-ViRVIG / pymotion

A Python library for working with motion data in numpy or PyTorch
MIT License
32 stars 2 forks source link

Having issues with fk operator and converting the results back for saving #6

Open techvd opened 3 weeks ago

techvd commented 3 weeks ago

When I use the fk operator I get the world positions and rotations (pos, rotmats). Is there a way to convert these back to the original local position/rotation values? As a simple test I'm trying to convert these back and trying to save a new bvh file (by using set_data) and that new bvh doesn't come out right. The animation kinda becomes a jumbled mess. I hope that makes sense.

Here is some simplified code I'm using to test this:

bvh = BVH()
bvh.load(bvh_file)
local_rotations, local_positions, parents, offsets, _, _ = bvh.get_data()
global_positions = local_positions[:, 0, :]  # root joint
pos, rotmats = fk(local_rotations, global_positions, offsets, parents)

# now save the first 30 frames to a new bvh file
new_bvh = BVH()
new_bvh.data = bvh.data.copy()
new_positions = np.zeros((30, len(bvh.data['names']), 3))
new_rotations = np.zeros((30, len(bvh.data['names']), 4))
for frame in range(30):
    for j in range(len(bvh.data['names'])):
        position = pos[frame][j]
        rotation = rotmats[frame][j]
        quat = from_matrix(rotation)
        new_positions[frame][j] = position
        new_rotations[frame][j] = quat
new_bvh.set_data(new_rotations, new_positions)
new_bvh.save('test.bvh')
JLPM22 commented 2 weeks ago

Hi! I modified your code; please see below.

There are some changes. First, I have included the from_root_quat() function (which will be soon available in the package, but for now, you can use this version). This function converts root space quaternions into local space ones. As you see, in the code, I first apply forward kinematics to get global space rotations, then apply the inverse of the root to get root space, apply this new function, and done! I hope it helps :)

import numpy as np
from pymotion.io.bvh import BVH
from pymotion.ops.forward_kinematics import fk
import pymotion.rotations.quat as quat

def from_root_quat(q: np.ndarray, parents: np.ndarray) -> np.ndarray:
    """
    Convert root-centered quaternions to the skeleton information.

    Parameters
    ----------
    q: np.ndarray[..., n_joints, 4]
        includes the root joint in global space in the first joint
    parents: np.ndarray[n_joints]

    Returns
    -------
    rotations : np.ndarray[..., n_joints, 4]
    """
    n_joints = q.shape[-2]
    # rotations has shape (frames, n_joints, 4)
    rotations = q.copy()
    # make transformations local to the parents
    # (initially local to the root)
    for j in reversed(range(1, n_joints)):
        parent = parents[j]
        if parent == 0:  # already in root space
            continue
        inv = quat.inverse(rotations[..., parent, :])
        rotations[..., j, :] = quat.mul(inv, rotations[..., j, :])
    return rotations

bvh = BVH()
bvh.load("test.bvh")
local_rotations, local_positions, parents, offsets, _, _ = bvh.get_data()
global_positions = local_positions[:, 0, :]  # root joint
_, rotmats = fk(local_rotations, global_positions, offsets, parents)

# convert to quaternions
quats_global = quat.from_matrix(rotmats)
# convert global space rotations to root space rotations (quats_root will have all joints except the root)
quats_root = quat.mul(quat.inverse(quats_global[:, 0:1]), quats_global[:, 1:])
# add back the root joint in global space
quats_root = np.concatenate([quats_global[:, 0:1], quats_root], axis=1)
# convert root space rotations to local space rotations
quats_local = from_root_quat(quats_root, parents)

# now save to a new bvh file
new_bvh = BVH()
new_bvh.data = bvh.data.copy()
new_bvh.set_data(quats_local, new_bvh.data["positions"])
new_bvh.save("test_out.bvh")
techvd commented 2 weeks ago

That works perfectly! Thank you so much for your time. I also want to say this library is pretty awesome!

techvd commented 2 weeks ago

Sorry I closed this too soon. So the sample code you have above works correctly. However, I'm having trouble with my test. I'm basically using this library to convert some motion files I have into data for training a neural network experimenting with motion generation. I basically 1) load the bvh file, 2) use the fk() function to get pos and rotmats, 3) convert the rotmats to global quaternions using the from_matrix function. So in the end what I have is generated pos and global quaternions. I'm now trying to turn them back into local rotations and positions for the set_data method so it can save the new BVH data. Does all of this make sense or are there better ways to do this back and forth? If this sounds ok, how would I do the final conversion? Thanks again for all your time!