File size: 2,292 Bytes
d7a991a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import tempfile

import numpy as np
import torch

from mmpose.models.utils import SMPL
from tests.utils.mesh_utils import generate_smpl_weight_file


def test_smpl():
    """Test smpl model."""

    # build smpl model
    smpl = None
    with tempfile.TemporaryDirectory() as tmpdir:
        # generate weight file for SMPL model.
        generate_smpl_weight_file(tmpdir)

        smpl_cfg = dict(
            smpl_path=tmpdir,
            joints_regressor=osp.join(tmpdir, 'test_joint_regressor.npy'))
        smpl = SMPL(**smpl_cfg)

    assert smpl is not None, 'Fail to build SMPL model'

    # test get face function
    faces = smpl.get_faces()
    assert isinstance(faces, np.ndarray)

    betas = torch.zeros(3, 10)
    body_pose = torch.zeros(3, 23 * 3)
    global_orient = torch.zeros(3, 3)
    transl = torch.zeros(3, 3)
    gender = torch.LongTensor([-1, 0, 1])

    # test forward with body_pose and global_orient in axis-angle format
    smpl_out = smpl(
        betas=betas, body_pose=body_pose, global_orient=global_orient)
    assert isinstance(smpl_out, dict)
    assert smpl_out['vertices'].shape == torch.Size([3, 6890, 3])
    assert smpl_out['joints'].shape == torch.Size([3, 24, 3])

    # test forward with body_pose and global_orient in rotation matrix format
    body_pose = torch.eye(3).repeat([3, 23, 1, 1])
    global_orient = torch.eye(3).repeat([3, 1, 1, 1])
    _ = smpl(betas=betas, body_pose=body_pose, global_orient=global_orient)

    # test forward with translation
    _ = smpl(
        betas=betas,
        body_pose=body_pose,
        global_orient=global_orient,
        transl=transl)

    # test forward with gender
    _ = smpl(
        betas=betas,
        body_pose=body_pose,
        global_orient=global_orient,
        transl=transl,
        gender=gender)

    # test forward when all samples in the same gender
    gender = torch.LongTensor([0, 0, 0])
    _ = smpl(
        betas=betas,
        body_pose=body_pose,
        global_orient=global_orient,
        transl=transl,
        gender=gender)

    # test forward when batch size = 0
    _ = smpl(
        betas=torch.zeros(0, 10),
        body_pose=torch.zeros(0, 23 * 3),
        global_orient=torch.zeros(0, 3))