SMPL/preprocess.py at master · whitec10ud/SMPL · GitHub

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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
#!/usr/bin/python2

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys
import os
import numpy as np
import pickle as pkl
import json


def main(args):
    """Main entrance.
    Arguments
    ----------
    - args: list of strings
        Command line arguments.
    Returns
    ----------
    """
    gender = args[1]
    raw_model_path = args[2]
    save_dir = args[3]

    if gender == 'female':
        NP_SAVE_FILE = 'smpl_female.npz'
        JSON_SVAE_FILE = 'smpl_female.json'
    elif gender == 'male':
        NP_SAVE_FILE = 'smpl_male.npz'
        JSON_SVAE_FILE = 'smpl_male.json'
    else:
        raise SystemError('Please specify gender of the model!\n'
                          'USAGE: \'*f*.pkl\' - female, '
                          '\'*m*.pkl\' - male')

    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    np_save_path = os.path.join(save_dir, NP_SAVE_FILE)
    json_save_path = os.path.join(save_dir, JSON_SVAE_FILE)

    # * Model Data Description * #
    # vertices_template: global vertex locations of template - (6890, 3)
    # face_indices: vertex indices of each face (triangles) - (13776, 3)
    # joint_regressor: joint regressor - (24, 6890)
    # kinematic_tree_table: table of kinematic tree - (2, 24)
    # weights: weights for linear blend skinning - (6890, 24)
    # shape_blend_shapes: shape blend shapes - (6890, 3, 10)
    # pose_blend_shapes: pose blend shapes - (6890, 3, 207)

    # * Extra Data Description *
    # Besides the data above, the official model provide the following things.
    # The pickle file downloaded from SMPL website seems to be redundant or
    # some of the contents are used for training the model. None of them will
    # be used to generate a new skinning.
    #
    # bs_stype: blend skinning style - (default)linear blend skinning
    # bs_type: blend skinning type - (default) linear rotation minimization
    # J: global joint locations of the template mesh - (24, 3)
    # J_regressor_prior: prior joint regressor - (24, 6890)
    # pose_training_info: pose training information - string list with 6
    #                     elements.
    # vert_sym_idxs: symmetrical corresponding vertex indices - (6890, )
    # weights_prior: prior weights for linear blend skinning
    with open(raw_model_path, 'rb') as f:
        raw_model_data = pkl.load(f)
    vertices_template = np.array(raw_model_data['v_template'])
    face_indices = np.array(raw_model_data['f'] + 1)  # starts from 1
    weights = np.array(raw_model_data['weights'])
    shape_blend_shapes = np.array(raw_model_data['shapedirs'])
    pose_blend_shapes = np.array(raw_model_data['posedirs'])
    joint_regressor = np.array(raw_model_data['J_regressor'].toarray())
    kinematic_tree = np.array(raw_model_data['kintree_table'])

    model_data_np = {
        'vertices_template': vertices_template,
        'face_indices': face_indices,
        'weights': weights,
        'shape_blend_shapes': shape_blend_shapes,
        'pose_blend_shapes': pose_blend_shapes,
        'joint_regressor': joint_regressor,
        'kinematic_tree': kinematic_tree
    }

    # Data must be converted to list before storing as json.
    model_data_json = {
        'vertices_template': vertices_template.tolist(),
        'face_indices': face_indices.tolist(),
        'weights': weights.tolist(),
        'shape_blend_shapes': shape_blend_shapes.tolist(),
        'pose_blend_shapes': pose_blend_shapes.tolist(),
        'joint_regressor': joint_regressor.tolist(),
        'kinematic_tree': kinematic_tree.tolist()
    }

    np.savez(np_save_path, **model_data_np)
    with open(json_save_path, 'wb+') as f:
        json.dump(model_data_json, f, indent=4, sort_keys=True)

    print('Save SMPL Model to: ', os.path.abspath(save_dir))


if __name__ == '__main__':
    if sys.version_info[0] != 2:
        raise EnvironmentError('Run this file with Python2!')
    if len(sys.argv) < 4:
        raise SystemError('Too few arguments!\n'
                          'USAGE: python2 preprocess.py '
                          '<gender> <path-to-the-pkl> '
                          '<dir-to-the-model>')
    elif len(sys.argv) > 4:
        raise SystemError('Too many arguments, only one model at a time!\n'
                          'USAGE: python2 preprocess.py '
                          '<path-to-the-pkl> <dir-to-the-model>')

    main(sys.argv)