Modify skeleton to align with QuaterNet implementation (c/p)

src/skelcast/data/human36m/skeleton.py

@@ -6,18 +6,29 @@
 #
 
 import numpy as np
+import torch
+
+from skelcast.data.human36m.quaternion import qmul_np, qmul, qrot
 
 class Skeleton:
- def __init__(self, parents, joints_left, joints_right):
- assert len(joints_left) == len(joints_right)
+ def __init__(self, offsets, parents, joints_left=None, joints_right=None):
+ assert len(offsets) == len(parents)
 
+ self._offsets = torch.FloatTensor(offsets)
  self._parents = np.array(parents)
  self._joints_left = joints_left
  self._joints_right = joints_right
  self._compute_metadata()
 
+ def cuda(self):
+ self._offsets = self._offsets.cuda()
+ return self
+
  def num_joints(self):
- return len(self._parents)
+ return self._offsets.shape[0]
+
+ def offsets(self):
+ return self._offsets
 
  def parents(self):
  return self._parents
@@ -28,19 +39,27 @@ def has_children(self):
  def children(self):
  return self._children
 
- def remove_joints(self, joints_to_remove):
+ def remove_joints(self, joints_to_remove, dataset):
  """
- Remove the joints specified in 'joints_to_remove'.
+ Remove the joints specified in 'joints_to_remove', both from the
+ skeleton definition and from the dataset (which is modified in place).
+ The rotations of removed joints are propagated along the kinematic chain.
  """
  valid_joints = []
  for joint in range(len(self._parents)):
  if joint not in joints_to_remove:
  valid_joints.append(joint)
-
- for i in range(len(self._parents)):
- while self._parents[i] in joints_to_remove:
- self._parents[i] = self._parents[self._parents[i]]
 
+ # Update all transformations in the dataset
+ for subject in dataset.subjects():
+ for action in dataset[subject].keys():
+ rotations = dataset[subject][action]['rotations']
+ for joint in joints_to_remove:
+ for child in self._children[joint]:
+ rotations[:, child] = qmul_np(rotations[:, joint], rotations[:, child])
+ rotations[:, joint] = [1, 0, 0, 0] # Identity
+ dataset[subject][action]['rotations'] = rotations[:, valid_joints]
+
  index_offsets = np.zeros(len(self._parents), dtype=int)
  new_parents = []
  for i, parent in enumerate(self._parents):
@@ -49,24 +68,41 @@ def remove_joints(self, joints_to_remove):
  else:
  index_offsets[i:] += 1
  self._parents = np.array(new_parents)
-
-
- if self._joints_left is not None:
- new_joints_left = []
- for joint in self._joints_left:
- if joint in valid_joints:
- new_joints_left.append(joint - index_offsets[joint])
- self._joints_left = new_joints_left
- if self._joints_right is not None:
- new_joints_right = []
- for joint in self._joints_right:
- if joint in valid_joints:
- new_joints_right.append(joint - index_offsets[joint])
- self._joints_right = new_joints_right
 
+ self._offsets = self._offsets[valid_joints]
  self._compute_metadata()
 
- return valid_joints
+ def forward_kinematics(self, rotations, root_positions):
+ """
+ Perform forward kinematics using the given trajectory and local rotations.
+ Arguments (where N = batch size, L = sequence length, J = number of joints):
+ -- rotations: (N, L, J, 4) tensor of unit quaternions describing the local rotations of each joint.
+ -- root_positions: (N, L, 3) tensor describing the root joint positions.
+ """
+ assert len(rotations.shape) == 4
+ assert rotations.shape[-1] == 4
+
+ positions_world = []
+ rotations_world = []
+
+ expanded_offsets = self._offsets.expand(rotations.shape[0], rotations.shape[1],
+ self._offsets.shape[0], self._offsets.shape[1])
+
+ # Parallelize along the batch and time dimensions
+ for i in range(self._offsets.shape[0]):
+ if self._parents[i] == -1:
+ positions_world.append(root_positions)
+ rotations_world.append(rotations[:, :, 0])
+ else:
+ positions_world.append(qrot(rotations_world[self._parents[i]], expanded_offsets[:, :, i]) \
+ + positions_world[self._parents[i]])
+ if self._has_children[i]:
+ rotations_world.append(qmul(rotations_world[self._parents[i]], rotations[:, :, i]))
+ else:
+ # This joint is a terminal node -> it would be useless to compute the transformation
+ rotations_world.append(None)
+
+ return torch.stack(positions_world, dim=3).permute(0, 1, 3, 2)
 
  def joints_left(self):
  return self._joints_left