--- title: DIP keywords: fastai sidebar: home_sidebar summary: "Load the prototype version of AMASS." description: "Load the prototype version of AMASS." nb_path: "06_dip.ipynb" ---
{% raw %}
{% endraw %} {% raw %}
import os
import pickle
from pathlib import Path
import numpy as np
import torch
import llamass.core
import llamass.transforms
from einops import repeat, rearrange
from scipy.spatial.transform import Rotation as R
{% endraw %}

A prototypical version of AMASS can be found in the downloads of the DIP project. This is used in some papers so being able to train on it is useful.

Sidenote: The official documentation for the DIPS dataset says the angles stored in data["poses"] in it's archives are stored as axis-angle vectors, but when I looked at the code to open this dataset in the SPL repository here I was surprised to see they converted it from a rotation matrix to other formalisms. So, I opened it up myself and tested if the data stored was a rotation matrix or not and it was! However, that was even more confusing because there were only 15 joint angles and I knew that SMPL is 24 joint angles but it turns out they didn't include all of them, which you can see if you look at the list of SMPL_MAJOR_JOINTS in the SPL forward kinematics model.

To do:

  • Visualize an example after applying the SMPL forward kinematics module to verify it's loaded correctly
{% raw %}
dips_path = Path("/nobackup/gngdb/Synthetic_60FPS/")

# open a random pickle file in the DIPS dataset
example_file = None
for dirpath, dirnames, filenames in os.walk(dips_path):
    dirpath = Path(dirpath)
    for filename in filenames:
        filename = Path(filename)
        if not filename.is_dir() and filename.suffix == ".pkl":
            example_file = dirpath / filename
    if example_file is not None:
        break

def pkl_load(pkl_file):
    try:
        with open(pkl_file, "rb") as f:
            return pickle.load(f)
    except UnicodeDecodeError:
        with open(pkl_file, "rb") as f:
            u = pickle._Unpickler(f)
            u.encoding = "latin1"
            return u.load()

cdata = pkl_load(example_file)

n = len(cdata['poses'])
d = len(cdata['poses'][0])
print(f"Contains {n} in a {d} dim vector containing {d//9} joint angles")

r = np.array(cdata['poses'][0]).reshape(-1,3,3)[0]
assert np.allclose(r.dot(r.T), np.eye(3))
print("Example rotation matrix equals identity when multiplied by it's transpose")
r.dot(r.T)
Contains 975 in a 135 dim vector containing 15 joint angles
Example rotation matrix equals identity when multiplied by it's transpose
array([[1.00000000e+00, 3.27052040e-17, 4.41648625e-17],
       [3.27052040e-17, 1.00000000e+00, 1.40562284e-17],
       [4.41648625e-17, 1.40562284e-17, 1.00000000e+00]])
{% endraw %}

Sparse to Full

The data is provided with only 15 major joints, the remaining joints must be padded.

Reference implementation from SPL:

{% raw %}
"""
SPL: training and evaluation of neural networks with a structured prediction layer.
Copyright (C) 2019 ETH Zurich, Emre Aksan, Manuel Kaufmann
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""
def _sparse_to_full(joint_angles_sparse, sparse_joints_idxs, tot_nr_joints, rep="rotmat"):
    """
    Pad the given sparse joint angles with identity elements to retrieve a full skeleton with `tot_nr_joints`
    many joints.
    Args:
        joint_angles_sparse: An np array of shape (N, len(sparse_joints_idxs) * dof)
          or (N, len(sparse_joints_idxs), dof)
        sparse_joints_idxs: A list of joint indices pointing into the full skeleton given by range(0, tot_nr_joints)
        tot_nr_jonts: Total number of joints in the full skeleton.
        rep: Which representation is used, rotmat or quat
    Returns:
        The padded joint angles as an array of shape (N, tot_nr_joints*dof)
    """
    joint_idxs = sparse_joints_idxs
    assert rep in ["rotmat", "quat", "aa"]
    dof = 9 if rep == "rotmat" else 4 if rep == "quat" else 3
    n_sparse_joints = len(sparse_joints_idxs)
    angles_sparse = np.reshape(joint_angles_sparse, [-1, n_sparse_joints, dof])

    # fill in the missing indices with the identity element
    smpl_full = np.zeros(shape=[angles_sparse.shape[0], tot_nr_joints, dof])  # (N, tot_nr_joints, dof)
    if rep == "quat":
        smpl_full[..., 0] = 1.0
    elif rep == "rotmat":
        smpl_full[..., 0] = 1.0
        smpl_full[..., 4] = 1.0
        smpl_full[..., 8] = 1.0
    else:
        pass  # nothing to do for angle-axis

    smpl_full[:, joint_idxs] = angles_sparse
    smpl_full = np.reshape(smpl_full, [-1, tot_nr_joints * dof])
    return smpl_full
{% endraw %}

Converting this to PyTorch:

{% raw %}

sparse_to_full[source]

sparse_to_full(joint_angles_sparse, sparse_joints_idxs, tot_nr_joints, rep='rotmat')

Pad the given sparse joint angles with identity elements to retrieve a full skeleton with tot_nr_joints many joints. Args: joint_angles_sparse: Tensor of shape (N, len(sparse_joints_idxs) dof) or (N, len(sparse_joints_idxs), dof) sparse_joints_idxs: A list of joint indices pointing into the full skeleton given by range(0, tot_nr_joints) tot_nr_jonts: Total number of joints in the full skeleton. rep: Which representation is used, rotmat or quat Returns: The padded joint angles as an array of shape (N, tot_nr_jointsdof)

{% endraw %} {% raw %}
def sparse_to_full(joint_angles_sparse, sparse_joints_idxs, tot_nr_joints, rep="rotmat"):
    """
    Pad the given sparse joint angles with identity elements to retrieve a full skeleton with `tot_nr_joints`
    many joints.
    Args:
        joint_angles_sparse: Tensor of shape (N, len(sparse_joints_idxs) * dof)
          or (N, len(sparse_joints_idxs), dof)
        sparse_joints_idxs: A list of joint indices pointing into the full skeleton given by range(0, tot_nr_joints)
        tot_nr_jonts: Total number of joints in the full skeleton.
        rep: Which representation is used, rotmat or quat
    Returns:
        The padded joint angles as an array of shape (N, tot_nr_joints*dof)
    """
    device = joint_angles_sparse.device
    joint_idxs = sparse_joints_idxs
    joint_idx_mapping = {j:i for i,j in enumerate(joint_idxs)}
    assert rep in ["rotmat", "quat", "aa"]
    dof = 9 if rep == "rotmat" else 4 if rep == "quat" else 3
    n_sparse_joints = len(sparse_joints_idxs)
    angles_sparse = joint_angles_sparse.view(-1, n_sparse_joints, dof)

    
    
    # fill in the missing indices with the identity element
    N = angles_sparse.size(0)
    #smpl_full = torch.zeros((N, tot_nr_joints, dof)).to(device)
    if rep == "quat":
        smpl_full = torch.tensor([1.0, 0., 0., 0.]).to(device)
        #smpl_full[..., 0] = 1.0
    elif rep == "rotmat":
        smpl_full = torch.eye(3).view(-1).to(device)
        #smpl_full[..., 0] = 1.0
        #smpl_full[..., 4] = 1.0
        #smpl_full[..., 8] = 1.0
    else:
        smpl_full = torch.zeros(3).to(device)
    
    # repeat these tensors along the N axis
    smpl_full = repeat(smpl_full, 'd -> N () d', N=N)
    
    # make a list of tensors for each joint
    joint_tensors = []
    for j in range(tot_nr_joints):
        if j in joint_idxs:
            k = joint_idx_mapping[j]
            joint_tensors.append(angles_sparse[:, [k]])
        else:
            joint_tensors.append(smpl_full)

    smpl_full =  torch.cat(joint_tensors, 1)
    smpl_full = smpl_full.view(-1, tot_nr_joints*dof)
    return smpl_full
{% endraw %} {% raw %}
fk = llamass.transforms.SMPL_ForwardKinematics()
sparse = torch.tensor(np.array(cdata['poses'][0])).float()
_full = _sparse_to_full(sparse.numpy(), fk.major_joints, fk.n_joints, rep="rotmat")
full = sparse_to_full(sparse, fk.major_joints, fk.n_joints, rep="rotmat")
assert np.allclose(_full, full.numpy())
{% endraw %} {% raw %}
import matplotlib.pyplot as plt

positions = fk.from_rotmat(full)

def plot_pose(positions, skeleton=None, parents=None, save_to=None):
    if parents:
        skeleton = [(i,j) for i, j in enumerate(parents)][1:]
    assert skeleton is not None
    fig, axes = plt.subplots(1, 3, figsize=(10,6))

    for d, ax in enumerate(axes):
        dims_to_plot = [i for i in range(3) if i != d]
        joints = positions
        j = joints[:, dims_to_plot]
        ax.scatter(*j.T, color="b", s=0.5)
        for bone in skeleton:
            if set(bone) <= set(range(positions.shape[0])):
                a = j[bone[0]]
                b = j[bone[1]]
                x, y = list(zip(a, b))
                ax.plot(x, y, color="r", alpha=0.5)
        ax.axes.xaxis.set_ticklabels([])
        ax.axes.yaxis.set_ticklabels([])
        ax.set_aspect('equal', adjustable='box')
    if save_to is not None:
        plt.tight_layout()
        plt.savefig(save_to)
        plt.close()
    else:
        plt.show()

plot_pose(positions[0], parents=fk.parents)
{% endraw %}

Creating new Forward Kinematics model that incorporates this:

{% raw %}

class SMPL_ForwardKinematics_Sparse[source]

SMPL_ForwardKinematics_Sparse() :: SMPL_ForwardKinematics

Forward Kinematics for the skeleton defined by SMPL.

{% endraw %} {% raw %}
class SMPL_ForwardKinematics_Sparse(llamass.transforms.SMPL_ForwardKinematics):
    def from_rotmat(self, joint_angles):
        mj, nj = self.major_joints, self.n_joints
        return super().from_rotmat(sparse_to_full(joint_angles, mj, nj, rep="rotmat"))
        
    def from_aa(self, joint_angles):
        mj, nj = self.major_joints, self.n_joints
        return super().from_aa(sparse_to_full(joint_angles, mj, nj, rep="aa"))
{% endraw %} {% raw %}
fk = SMPL_ForwardKinematics_Sparse()
plot_pose(fk.from_rotmat(sparse)[0], parents=fk.parents)
{% endraw %}

Converting to npz

AMASS proper is stored as npz files, after converting this dataset to that representation, it will take up less space and be easy to load with the AMASS Dataset class.

I've downloaded the txt files containing the lists of files to go into each of training, test and validation splits from here.

To do:

  • Console utility to process the data into npz files like the AMASS dataset
    • Get list of all the files
    • Split the list of files according to the train.txt...
    • Process each archive in parallel with joblib, saving to new location
{% raw %}
%%capture
%%bash
wget -O training_fnames.txt https://raw.githubusercontent.com/eth-ait/spl/master/preprocessing/training_fnames.txt
wget -O test_fnames.txt https://raw.githubusercontent.com/eth-ait/spl/master/preprocessing/test_fnames.txt
wget -O validation_fnames.txt https://raw.githubusercontent.com/eth-ait/spl/master/preprocessing/validation_fnames.txt
{% endraw %} {% raw %}

iter_pkl_in[source]

iter_pkl_in(dip_dir)

{% endraw %} {% raw %}
def iter_pkl_in(dip_dir):
    # open a random pickle file in the DIPS dataset
    example_file = None
    for dirpath, dirnames, filenames in os.walk(dips_path):
        dirpath = Path(dirpath)
        for filename in filenames:
            filename = Path(filename)
            if not filename.is_dir() and filename.suffix == ".pkl":
                yield dirpath/filename
{% endraw %} {% raw %}
for pkl_path in iter_pkl_in(dips_path):
    print(pkl_path)
    break
/nobackup/gngdb/Synthetic_60FPS/AMASS_ACCAD/Female1General_c3dA10_SB__SB2__SB_lie_SB_to_SB_crouch_dynamics.pkl
{% endraw %} {% raw %}
def load_txt(fpath):
    with open(fpath) as f:
        return [Path(p.rstrip()) for p in f.readlines()]
split_paths = {}
split_paths["train"] = load_txt("training_fnames.txt")
split_paths["validation"] = load_txt("validation_fnames.txt")
split_paths["test"] = load_txt("test_fnames.txt")
{% endraw %} {% raw %}
pkey = lambda p: (p.parts[-2].split("_")[-1], p.name)
split_fnames = {pkey(p):n for n,v in split_paths.items() for p in v}
separated_paths = {k:[] for k in split_paths}
allocated, unallocated = [], []
for pkl_path in iter_pkl_in(dips_path):
    try:
        split = split_fnames[pkey(pkl_path)]
        assert pkey(pkl_path) not in allocated
        allocated.append(pkey(pkl_path))
        separated_paths[split].append(pkl_path)
    except KeyError:
        unallocated.append(pkl_path)
n = len(set(p for k,v in separated_paths.items() for p in v))
prescribed = set(k for k in split_fnames)
assert n == len(prescribed)
{% endraw %} {% raw %}
from tqdm.auto import tqdm
def process_to_npz():
    with tqdm(total=len(prescribed)) as pbar:
        for split, pkl_paths in separated_paths.items():
            for pkl_path in pkl_paths:
                dip_path = pkl_path.parents[2]
                dest_dir =  dip_path / split / pkl_path.parts[-2].split("_")[-1]
                dest_path = dest_dir / (pkl_path.parts[-1].split(".")[0] + ".npz")
                dest_dir.mkdir(exist_ok=True, parents=True)
                with open(pkl_path, "rb") as f:
                    cdata = pickle.load(f, encoding='latin1')
                np.savez(dest_path, poses=np.array(cdata["poses"]))
                pbar.update(1)
# process_to_npz()
{% endraw %}

Trying to open this new processed dataset using the AMASS Dataset class:

{% raw %}
dip_path = Path("/nobackup/gngdb/dips/train/")
dataset = llamass.core.AMASS(
    dip_path,
    clip_length=144,
    overlapping=False,
    transform=torch.tensor,
    data_keys=("poses",),
    strict=False
)
{% endraw %} {% raw %}
for example in dataset:
    print(example)
    for k, v in example.items():
        print(k, v.size())
    break
{'poses': tensor([[ 0.9903, -0.1298,  0.0501,  ..., -0.8230, -0.5596,  0.0976],
        [ 0.9904, -0.1307,  0.0457,  ..., -0.8153, -0.5775,  0.0428],
        [ 0.9904, -0.1320,  0.0400,  ..., -0.8051, -0.5930, -0.0100],
        ...,
        [ 0.9402, -0.2332,  0.2483,  ..., -0.7959, -0.5205, -0.3091],
        [ 0.9378, -0.2353,  0.2554,  ..., -0.7994, -0.5183, -0.3038],
        [ 0.9377, -0.2353,  0.2557,  ..., -0.7965, -0.5208, -0.3071]])}
poses torch.Size([144, 135])
{% endraw %}

Is it possible to load all of DIP into memory?

{% raw %}
def build_tensor_dataset():
    dips_examples = []
    for example in tqdm(dataset):
        dips_examples.append(example)
    dip = torch.stack([e["poses"] for e in dips_examples])
    print(dip.size(), f" number of poses={dip.size(0)*dip.size(1)}")
    torch.save(dip, dip_path/"tensor_dataset.pt")
# build_tensor_dataset()
{% endraw %} {% raw %}
!du -hs /nobackup/gngdb/dips/train/tensor_dataset.pt
3.0G	/nobackup/gngdb/dips/train/tensor_dataset.pt
{% endraw %}

Adding Additional Training Data

AMASS contains extra data, can this be added to assist during training?