Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Human3.6M Flat dataset #76

Merged
merged 2 commits into from
Jan 31, 2024
Merged

Human3.6M Flat dataset #76

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3061d53
Flatten dataset for easier index retrieval
kaseris Jan 31, 2024
53f2225
Modify skeleton to align with QuaterNet implementation (c/p)
kaseris Jan 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
236 changes: 147 additions & 89 deletions src/skelcast/data/human36m/human36m.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,28 +16,29 @@ def __init__(self, path, skeleton, fps):
self._fps = fps
self._use_gpu = False
self._skeleton = skeleton

def cuda(self):
self._use_gpu = True
self._skeleton.cuda()
return self

def _load(self, path):
result = {}
data = np.load(path, 'r')
for i, (trajectory, rotations, subject, action) in enumerate(zip(data['trajectories'],
data['rotations'],
data['subjects'],
data['actions'])):
data = np.load(path, allow_pickle=True)
for i, (trajectory, rotations, subject, action) in enumerate(
zip(
data["positions"],
data["rotations"],
data["subjects"],
data["actions"],
)
):
if subject not in result:
result[subject] = {}

result[subject][action] = {
'rotations': rotations,
'trajectory': trajectory
}

result[subject][action] = {"rotations": rotations, "trajectory": trajectory}
return result

def downsample(self, factor, keep_strides=True):
"""
Downsample this dataset by an integer factor, keeping all strides of the data
Expand All @@ -47,95 +48,107 @@ def downsample(self, factor, keep_strides=True):
will have '_d0', ... '_dn' appended to their names.
"""
assert self._fps % factor == 0

for subject in self._data.keys():
new_actions = {}
for action in list(self._data[subject].keys()):
for idx in range(factor):
tup = {}
for k in self._data[subject][action].keys():
tup[k] = self._data[subject][action][k][idx::factor]
new_actions[action + '_d' + str(idx)] = tup
new_actions[action + "_d" + str(idx)] = tup
if not keep_strides:
break
self._data[subject] = new_actions

self._fps //= factor

def _mirror_sequence(self, sequence):
mirrored_rotations = sequence['rotations'].copy()
mirrored_trajectory = sequence['trajectory'].copy()
mirrored_rotations = sequence["rotations"].copy()
mirrored_trajectory = sequence["trajectory"].copy()

joints_left = self._skeleton.joints_left()
joints_right = self._skeleton.joints_right()

# Flip left/right joints
mirrored_rotations[:, joints_left] = sequence['rotations'][:, joints_right]
mirrored_rotations[:, joints_right] = sequence['rotations'][:, joints_left]
mirrored_rotations[:, joints_left] = sequence["rotations"][:, joints_right]
mirrored_rotations[:, joints_right] = sequence["rotations"][:, joints_left]

mirrored_rotations[:, :, [2, 3]] *= -1
mirrored_trajectory[:, 0] *= -1

return {
'rotations': qfix(mirrored_rotations),
'trajectory': mirrored_trajectory
"rotations": qfix(mirrored_rotations),
"trajectory": mirrored_trajectory,
}

def mirror(self):
"""
Perform data augmentation by mirroring every sequence in the dataset.
The mirrored sequences will have '_m' appended to the action name.
"""
for subject in self._data.keys():
for action in list(self._data[subject].keys()):
if '_m' in action:
if "_m" in action:
continue
self._data[subject][action + '_m'] = self._mirror_sequence(self._data[subject][action])

self._data[subject][action + "_m"] = self._mirror_sequence(
self._data[subject][action]
)

def compute_euler_angles(self, order):
for subject in self._data.values():
for action in subject.values():
action['rotations_euler'] = qeuler_np(action['rotations'], order, use_gpu=self._use_gpu)

action["rotations_euler"] = qeuler_np(
action["rotations"], order, use_gpu=self._use_gpu
)

def compute_positions(self):
for subject in self._data.values():
for action in subject.values():
rotations = torch.from_numpy(action['rotations'].astype('float32')).unsqueeze(0)
trajectory = torch.from_numpy(action['trajectory'].astype('float32')).unsqueeze(0)
rotations = torch.from_numpy(
action["rotations"].astype("float32")
).unsqueeze(0)
trajectory = torch.from_numpy(
action["trajectory"].astype("float32")
).unsqueeze(0)
if self._use_gpu:
rotations = rotations.cuda()
trajectory = trajectory.cuda()
action['positions_world'] = self._skeleton.forward_kinematics(rotations, trajectory).squeeze(0).cpu().numpy()

action["positions_world"] = (
self._skeleton.forward_kinematics(rotations, trajectory)
.squeeze(0)
.cpu()
.numpy()
)

# Absolute translations across the XY plane are removed here
trajectory[:, :, [0, 2]] = 0
action['positions_local'] = self._skeleton.forward_kinematics(rotations, trajectory).squeeze(0).cpu().numpy()


def __getitem__(self, key):
return self._data[key]


action["positions_local"] = (
self._skeleton.forward_kinematics(rotations, trajectory)
.squeeze(0)
.cpu()
.numpy()
)

# def __getitem__(self, key):
# return self._data[key]

def subjects(self):
return self._data.keys()



def subject_actions(self, subject):
return self._data[subject].keys()



def all_actions(self):
result = []
for subject, actions in self._data.items():
for action in actions.keys():
result.append((subject, action))
return result



def fps(self):
return self._fps



def skeleton(self):
return self._skeleton

Expand All @@ -145,44 +158,89 @@ class Human36MDataset(MocapDataset):
"""
TODO: Possibly add a flatten_data method for easy accessing with a single index.
"""

def __init__(self, path, seq_len=27):
skeleton = Skeleton(offsets=[
[ 0. , 0. , 0. ],
[-132.948591, 0. , 0. ],
[ 0. , -442.894612, 0. ],
[ 0. , -454.206447, 0. ],
[ 0. , 0. , 162.767078],
[ 0. , 0. , 74.999437],
[ 132.948826, 0. , 0. ],
[ 0. , -442.894413, 0. ],
[ 0. , -454.20659 , 0. ],
[ 0. , 0. , 162.767426],
[ 0. , 0. , 74.999948],
[ 0. , 0.1 , 0. ],
[ 0. , 233.383263, 0. ],
[ 0. , 257.077681, 0. ],
[ 0. , 121.134938, 0. ],
[ 0. , 115.002227, 0. ],
[ 0. , 257.077681, 0. ],
[ 0. , 151.034226, 0. ],
[ 0. , 278.882773, 0. ],
[ 0. , 251.733451, 0. ],
[ 0. , 0. , 0. ],
[ 0. , 0. , 99.999627],
[ 0. , 100.000188, 0. ],
[ 0. , 0. , 0. ],
[ 0. , 257.077681, 0. ],
[ 0. , 151.031437, 0. ],
[ 0. , 278.892924, 0. ],
[ 0. , 251.72868 , 0. ],
[ 0. , 0. , 0. ],
[ 0. , 0. , 99.999888],
[ 0. , 137.499922, 0. ],
[ 0. , 0. , 0. ]
],
parents=[-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 12,
16, 17, 18, 19, 20, 19, 22, 12, 24, 25, 26, 27, 28, 27, 30],
joints_left=[1, 2, 3, 4, 5, 24, 25, 26, 27, 28, 29, 30, 31],
joints_right=[6, 7, 8, 9, 10, 16, 17, 18, 19, 20, 21, 22, 23])
skeleton = Skeleton(
offsets=[
[0.0, 0.0, 0.0],
[-132.948591, 0.0, 0.0],
[0.0, -442.894612, 0.0],
[0.0, -454.206447, 0.0],
[0.0, 0.0, 162.767078],
[0.0, 0.0, 74.999437],
[132.948826, 0.0, 0.0],
[0.0, -442.894413, 0.0],
[0.0, -454.20659, 0.0],
[0.0, 0.0, 162.767426],
[0.0, 0.0, 74.999948],
[0.0, 0.1, 0.0],
[0.0, 233.383263, 0.0],
[0.0, 257.077681, 0.0],
[0.0, 121.134938, 0.0],
[0.0, 115.002227, 0.0],
[0.0, 257.077681, 0.0],
[0.0, 151.034226, 0.0],
[0.0, 278.882773, 0.0],
[0.0, 251.733451, 0.0],
[0.0, 0.0, 0.0],
[0.0, 0.0, 99.999627],
[0.0, 100.000188, 0.0],
[0.0, 0.0, 0.0],
[0.0, 257.077681, 0.0],
[0.0, 151.031437, 0.0],
[0.0, 278.892924, 0.0],
[0.0, 251.72868, 0.0],
[0.0, 0.0, 0.0],
[0.0, 0.0, 99.999888],
[0.0, 137.499922, 0.0],
[0.0, 0.0, 0.0],
],
parents=[
-1,
0,
1,
2,
3,
4,
0,
6,
7,
8,
9,
0,
11,
12,
13,
14,
12,
16,
17,
18,
19,
20,
19,
22,
12,
24,
25,
26,
27,
28,
27,
30,
],
joints_left=[1, 2, 3, 4, 5, 24, 25, 26, 27, 28, 29, 30, 31],
joints_right=[6, 7, 8, 9, 10, 16, 17, 18, 19, 20, 21, 22, 23],
)
super().__init__(path, skeleton, fps=50)
self.compute_positions()
self._dataset_flat = []
for subject in ['S1', 'S5', 'S6', 'S7', 'S8', 'S9', 'S11']:
for action in list(self._data[subject].keys()):
self._dataset_flat.append(self._data[subject][action]['rotations'])

def __getitem__(self, index):
return self._dataset_flat[index]

def __len__(self):
return len(self._dataset_flat)
Loading
Loading