Spaces:
Starting
on
L40S
Starting
on
L40S
File size: 3,531 Bytes
d7e58f0 |
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 |
import numpy as np
import torch
import torch.nn as nn
from mmcv.runner.base_module import BaseModule
from detrsmpl.utils.geometry import rot6d_to_rotmat
class HMRHead(BaseModule):
def __init__(self,
feat_dim,
smpl_mean_params=None,
npose=144,
nbeta=10,
ncam=3,
hdim=1024,
init_cfg=None):
super(HMRHead, self).__init__(init_cfg=init_cfg)
self.fc1 = nn.Linear(feat_dim + npose + nbeta + ncam, hdim)
self.drop1 = nn.Dropout()
self.fc2 = nn.Linear(hdim, hdim)
self.drop2 = nn.Dropout()
self.decpose = nn.Linear(hdim, npose)
self.decshape = nn.Linear(hdim, nbeta)
self.deccam = nn.Linear(hdim, ncam)
nn.init.xavier_uniform_(self.decpose.weight, gain=0.01)
nn.init.xavier_uniform_(self.decshape.weight, gain=0.01)
nn.init.xavier_uniform_(self.deccam.weight, gain=0.01)
if smpl_mean_params is None:
init_pose = torch.zeros([1, npose])
init_shape = torch.zeros([1, nbeta])
init_cam = torch.FloatTensor([[1, 0, 0]])
else:
mean_params = np.load(smpl_mean_params)
init_pose = torch.from_numpy(mean_params['pose'][:]).unsqueeze(0)
init_shape = torch.from_numpy(
mean_params['shape'][:].astype('float32')).unsqueeze(0)
init_cam = torch.from_numpy(mean_params['cam']).unsqueeze(0)
self.register_buffer('init_pose', init_pose)
self.register_buffer('init_shape', init_shape)
self.register_buffer('init_cam', init_cam)
def forward(self,
x,
init_pose=None,
init_shape=None,
init_cam=None,
n_iter=3):
# hmr head only support one layer feature
if isinstance(x, list) or isinstance(x, tuple):
x = x[-1]
output_seq = False
if len(x.shape) == 4:
# use feature from the last layer of the backbone
# apply global average pooling on the feature map
x = x.mean(dim=-1).mean(dim=-1)
elif len(x.shape) == 3:
# temporal feature
output_seq = True
B, T, L = x.shape
x = x.view(-1, L)
batch_size = x.shape[0]
if init_pose is None:
init_pose = self.init_pose.expand(batch_size, -1)
if init_shape is None:
init_shape = self.init_shape.expand(batch_size, -1)
if init_cam is None:
init_cam = self.init_cam.expand(batch_size, -1)
pred_pose = init_pose
pred_shape = init_shape
pred_cam = init_cam
for i in range(n_iter):
xc = torch.cat([x, pred_pose, pred_shape, pred_cam], 1)
xc = self.fc1(xc)
xc = self.drop1(xc)
xc = self.fc2(xc)
xc = self.drop2(xc)
pred_pose = self.decpose(xc) + pred_pose
pred_shape = self.decshape(xc) + pred_shape
pred_cam = self.deccam(xc) + pred_cam
pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3)
if output_seq:
pred_rotmat = pred_rotmat.view(B, T, 24, 3, 3)
pred_shape = pred_shape.view(B, T, 10)
pred_cam = pred_cam.view(B, T, 3)
output = {
'pred_pose': pred_rotmat,
'pred_shape': pred_shape,
'pred_cam': pred_cam
}
return output
|