Spaces:
Runtime error
Runtime error
File size: 1,884 Bytes
29a229f |
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 |
import torch
import numpy as np
import pickle
from typing import Optional
import smplx
from smplx.lbs import vertices2joints
from smplx.utils import SMPLOutput
class SMPL(smplx.SMPLLayer):
def __init__(self, *args, joint_regressor_extra: Optional[str] = None, update_hips: bool = False, **kwargs):
"""
Extension of the official SMPL implementation to support more joints.
Args:
Same as SMPLLayer.
joint_regressor_extra (str): Path to extra joint regressor.
"""
super(SMPL, self).__init__(*args, **kwargs)
smpl_to_openpose = [24, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4,
7, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34]
if joint_regressor_extra is not None:
self.register_buffer('joint_regressor_extra', torch.tensor(pickle.load(open(joint_regressor_extra, 'rb'), encoding='latin1'), dtype=torch.float32))
self.register_buffer('joint_map', torch.tensor(smpl_to_openpose, dtype=torch.long))
self.update_hips = update_hips
def forward(self, *args, **kwargs) -> SMPLOutput:
"""
Run forward pass. Same as SMPL and also append an extra set of joints if joint_regressor_extra is specified.
"""
smpl_output = super(SMPL, self).forward(*args, **kwargs)
joints = smpl_output.joints[:, self.joint_map, :]
if self.update_hips:
joints[:,[9,12]] = joints[:,[9,12]] + \
0.25*(joints[:,[9,12]]-joints[:,[12,9]]) + \
0.5*(joints[:,[8]] - 0.5*(joints[:,[9,12]] + joints[:,[12,9]]))
if hasattr(self, 'joint_regressor_extra'):
extra_joints = vertices2joints(self.joint_regressor_extra, smpl_output.vertices)
joints = torch.cat([joints, extra_joints], dim=1)
smpl_output.joints = joints
return smpl_output
|