| 2 | import scipy.ndimage.filters as filters |
| 3 | |
| 4 | class Skeleton(object): |
| 5 | def __init__(self, offset, kinematic_tree, device): |
| 6 | self.device = device |
| 7 | self._raw_offset_np = offset.numpy() |
| 8 | self._raw_offset = offset.clone().detach().to(device).float() |
| 9 | self._kinematic_tree = kinematic_tree |
| 10 | self._offset = None |
| 11 | self._parents = [0] * len(self._raw_offset) |
| 12 | self._parents[0] = -1 |
| 13 | for chain in self._kinematic_tree: |
| 14 | for j in range(1, len(chain)): |
| 15 | self._parents[chain[j]] = chain[j-1] |
| 16 | |
| 17 | def njoints(self): |
| 18 | return len(self._raw_offset) |
| 19 | |
| 20 | def offset(self): |
| 21 | return self._offset |
| 22 | |
| 23 | def set_offset(self, offsets): |
| 24 | self._offset = offsets.clone().detach().to(self.device).float() |
| 25 | |
| 26 | def kinematic_tree(self): |
| 27 | return self._kinematic_tree |
| 28 | |
| 29 | def parents(self): |
| 30 | return self._parents |
| 31 | |
| 32 | # joints (batch_size, joints_num, 3) |
| 33 | def get_offsets_joints_batch(self, joints): |
| 34 | assert len(joints.shape) == 3 |
| 35 | _offsets = self._raw_offset.expand(joints.shape[0], -1, -1).clone() |
| 36 | for i in range(1, self._raw_offset.shape[0]): |
| 37 | _offsets[:, i] = torch.norm(joints[:, i] - joints[:, self._parents[i]], p=2, dim=1)[:, None] * _offsets[:, i] |
| 38 | |
| 39 | self._offset = _offsets.detach() |
| 40 | return _offsets |
| 41 | |
| 42 | # joints (joints_num, 3) |
| 43 | def get_offsets_joints(self, joints): |
| 44 | assert len(joints.shape) == 2 |
| 45 | _offsets = self._raw_offset.clone() |
| 46 | for i in range(1, self._raw_offset.shape[0]): |
| 47 | # print(joints.shape) |
| 48 | _offsets[i] = torch.norm(joints[i] - joints[self._parents[i]], p=2, dim=0) * _offsets[i] |
| 49 | |
| 50 | self._offset = _offsets.detach() |
| 51 | return _offsets |
| 52 | |
| 53 | # face_joint_idx should follow the order of right hip, left hip, right shoulder, left shoulder |
| 54 | # joints (batch_size, joints_num, 3) |
| 55 | def inverse_kinematics_np(self, joints, face_joint_idx, smooth_forward=False): |
| 56 | assert len(face_joint_idx) == 4 |
| 57 | '''Get Forward Direction''' |
| 58 | l_hip, r_hip, sdr_r, sdr_l = face_joint_idx |
| 59 | across1 = joints[:, r_hip] - joints[:, l_hip] |
| 60 | across2 = joints[:, sdr_r] - joints[:, sdr_l] |
| 61 | across = across1 + across2 |
no outgoing calls
no test coverage detected