You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
83 lines
2.5 KiB
83 lines
2.5 KiB
|
|
import numpy as np
|
|
|
|
|
|
class Skeleton:
|
|
def __init__(self, parents, joints_left, joints_right):
|
|
assert len(joints_left) == len(joints_right)
|
|
|
|
self._parents = np.array(parents)
|
|
self._joints_left = joints_left
|
|
self._joints_right = joints_right
|
|
self._compute_metadata()
|
|
|
|
def num_joints(self):
|
|
return len(self._parents)
|
|
|
|
def parents(self):
|
|
return self._parents
|
|
|
|
def has_children(self):
|
|
return self._has_children
|
|
|
|
def children(self):
|
|
return self._children
|
|
|
|
def remove_joints(self, joints_to_remove):
|
|
|
|
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]]
|
|
|
|
index_offsets = np.zeros(len(self._parents), dtype=int)
|
|
new_parents = []
|
|
for i, parent in enumerate(self._parents):
|
|
if i not in joints_to_remove:
|
|
new_parents.append(parent - index_offsets[parent])
|
|
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._compute_metadata()
|
|
|
|
return valid_joints
|
|
|
|
def joints_left(self):
|
|
return self._joints_left
|
|
|
|
def joints_right(self):
|
|
return self._joints_right
|
|
|
|
def _compute_metadata(self):
|
|
self._has_children = np.zeros(len(self._parents)).astype(bool)
|
|
for i, parent in enumerate(self._parents):
|
|
if parent != -1:
|
|
self._has_children[parent] = True
|
|
|
|
self._children = []
|
|
for i, parent in enumerate(self._parents):
|
|
self._children.append([])
|
|
for i, parent in enumerate(self._parents):
|
|
if parent != -1:
|
|
self._children[parent].append(i)
|
|
|
|
|
|
|