-
Notifications
You must be signed in to change notification settings - Fork 0
/
extract_poses_2d.py
39 lines (29 loc) · 1.21 KB
/
extract_poses_2d.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import argparse
from pathlib import Path
import numpy as np
from metadata import SUBJECTS, VIEWS
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--root", type=Path, required=True)
parser.add_argument("--poses_2d", type=Path, required=True)
args = parser.parse_args()
poses_2d = args.poses_2d
root = args.root
with np.load(poses_2d, allow_pickle=True) as file:
data = file["positions_2d"].item()
for subject in SUBJECTS:
for sequence in data[subject]:
for i, values in enumerate(data[subject][sequence]):
sequence = sequence.replace(" ", "_")
view = VIEWS[i]
dst_poses_2d = root / subject / "Sequences" / sequence / "Poses" / "2D" / "CPN" / f"{view}.npz"
if not dst_poses_2d.parent.exists():
dst_poses_2d.parent.mkdir(parents=True)
print(dst_poses_2d)
n, j, d = values.shape
values = values.reshape((n, j, d, 1))
values = values.astype(dtype=np.float32)
keys = np.arange(n)
np.savez_compressed(dst_poses_2d, keys=keys, values=values)
if __name__ == "__main__":
main()