forked from nkolot/SPIN
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pw3d.py
73 lines (66 loc) · 3.05 KB
/
pw3d.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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import os
import cv2
import numpy as np
import pickle
def pw3d_extract(dataset_path, out_path):
# scale factor
scaleFactor = 1.2
# structs we use
imgnames_, scales_, centers_, parts_ = [], [], [], []
poses_, shapes_, genders_ = [], [], []
# get a list of .pkl files in the directory
dataset_path = os.path.join(dataset_path, 'sequenceFiles', 'test')
files = [os.path.join(dataset_path, f)
for f in os.listdir(dataset_path) if f.endswith('.pkl')]
# go through all the .pkl files
for filename in files:
with open(filename, 'rb') as f:
data = pickle.load(f, encoding='latin1')
smpl_pose = data['poses']
smpl_betas = data['betas']
poses2d = data['poses2d']
global_poses = data['cam_poses']
genders = data['genders']
valid = np.array(data['campose_valid']).astype(np.bool)
num_people = len(smpl_pose)
num_frames = len(smpl_pose[0])
seq_name = str(data['sequence'])
img_names = np.array(['imageFiles/' + seq_name + '/image_%s.jpg' % str(i).zfill(5) for i in range(num_frames)])
# get through all the people in the sequence
for i in range(num_people):
valid_pose = smpl_pose[i][valid[i]]
valid_betas = np.tile(smpl_betas[i][:10].reshape(1,-1), (num_frames, 1))
valid_betas = valid_betas[valid[i]]
valid_keypoints_2d = poses2d[i][valid[i]]
valid_img_names = img_names[valid[i]]
valid_global_poses = global_poses[valid[i]]
gender = genders[i]
# consider only valid frames
for valid_i in range(valid_pose.shape[0]):
part = valid_keypoints_2d[valid_i,:,:].T
part = part[part[:,2]>0,:]
bbox = [min(part[:,0]), min(part[:,1]),
max(part[:,0]), max(part[:,1])]
center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2]
scale = scaleFactor*max(bbox[2]-bbox[0], bbox[3]-bbox[1])/200
# transform global pose
pose = valid_pose[valid_i]
extrinsics = valid_global_poses[valid_i][:3,:3]
pose[:3] = cv2.Rodrigues(np.dot(extrinsics, cv2.Rodrigues(pose[:3])[0]))[0].T[0]
imgnames_.append(valid_img_names[valid_i])
centers_.append(center)
scales_.append(scale)
poses_.append(pose)
shapes_.append(valid_betas[valid_i])
genders_.append(gender)
# store data
if not os.path.isdir(out_path):
os.makedirs(out_path)
out_file = os.path.join(out_path,
'3dpw_test.npz')
np.savez(out_file, imgname=imgnames_,
center=centers_,
scale=scales_,
pose=poses_,
shape=shapes_,
gender=genders_)