--- 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" ---
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
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:
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)
Reference implementation from SPL:
"""
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
Converting this to PyTorch:
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
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())
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)
Creating new Forward Kinematics model that incorporates this:
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"))
fk = SMPL_ForwardKinematics_Sparse()
plot_pose(fk.from_rotmat(sparse)[0], parents=fk.parents)
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:
train.txt
...%%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
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
for pkl_path in iter_pkl_in(dips_path):
print(pkl_path)
break
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")
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)
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()
Trying to open this new processed dataset using the AMASS Dataset class:
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
)
for example in dataset:
print(example)
for k, v in example.items():
print(k, v.size())
break
Is it possible to load all of DIP into memory?
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()
!du -hs /nobackup/gngdb/dips/train/tensor_dataset.pt