def getHandJoints(self, xs, b_size, seq_size):
# Initialize Camera Parameters
scale = xs[:, 0] # Scale 1
trans = xs[:, 1:3] # Translate 2
rotation = xs[:, 3:6] # Rotation 3
theta = xs[:, 6:16] # Theta 10
beta = xs[:, 16:] # Beta 10
scale = scale * torch.tensor([self.camera_s['std']]).cuda() + torch.tensor([self.camera_s['mean']]).cuda()
trans = trans * torch.tensor([self.camera_u['std'], self.camera_v['std']]).unsqueeze(0).cuda() + torch.tensor([self.camera_u['mean'], self.camera_v['mean']]).unsqueeze(0).cuda()
batch_size = theta.size(0)
theta_ = torch.cat((torch.tensor([math.pi, 0., 0.]).unsqueeze(0).repeat(batch_size, 1).float().cuda(), theta), 1) # [pi, 0, 0, theta]
verts_3d_pred, joint_3d_pred = self.mano_layerR(theta_, beta)
R = batch_rodrigues(rotation).view(batch_size, 3, 3)
joint_3d_pred = torch.transpose(torch.matmul(R, torch.transpose(joint_3d_pred, 1, 2)), 1, 2) / 1000 # [mm] >> [m]
verts_3d_pred = torch.transpose(torch.matmul(R, torch.transpose(verts_3d_pred, 1, 2)), 1, 2) / 1000 # [mm] >> [m]
joint_2d_pred = joint_3d_pred[:, :, :2] * scale.unsqueeze(1).unsqueeze(2) + trans.unsqueeze(1) # [pixel]
verts_2d_pred = verts_3d_pred[:, :, :2] * scale.unsqueeze(1).unsqueeze(2) + trans.unsqueeze(1) # [pixel]
x3d = torch.cat((joint_3d_pred, verts_3d_pred), dim=1)
x2d = torch.cat((joint_2d_pred, verts_2d_pred), dim=1).view(batch_size, -1)
x2d = x2d.view(b_size, seq_size, -1)
x3d = x3d.view(b_size, seq_size, -1)
xs = xs.view(b_size, seq_size, -1)
theta = theta.view(b_size, seq_size, -1)
beta = beta.view(b_size, seq_size, -1)
return x2d, x3d, xs, theta, beta
def model_forward():
param_dict = {}
for key, var in var_dict.items():
if part_key == key:
param_dict[key] = batch_rodrigues(
var.reshape(-1, 3)).reshape(len(var), -1, 3, 3)
param_dict[key][:, jidx] = batch_rodrigues(
part.reshape(-1, 3)).reshape(-1, 3, 3)
else:
# Decode the axis-angles
if 'pose' in key or 'orient' in key:
param_dict[key] = batch_rodrigues(
var.reshape(-1, 3)).reshape(len(var), -1, 3, 3)
else:
# Simply pass the variable
param_dict[key] = var
return body_model(
return_full_pose=True, get_skin=True, **param_dict)
def summary_closure(gt_vertices, var_dict, body_model, mask_ids=None):
param_dict = {}
for key, var in var_dict.items():
# Decode the axis-angles
if 'pose' in key or 'orient' in key:
param_dict[key] = batch_rodrigues(
var.reshape(-1, 3)).reshape(len(var), -1, 3, 3)
else:
# Simply pass the variable
param_dict[key] = var
body_model_output = body_model(
return_full_pose=True, get_skin=True, **param_dict)
est_vertices = body_model_output['vertices']
if mask_ids is not None:
est_vertices = est_vertices[:, mask_ids]
gt_vertices = gt_vertices[:, mask_ids]
v2v = (est_vertices - gt_vertices).pow(2).sum(dim=-1).sqrt().mean()
return {
'Vertex-to-Vertex': v2v * 1000}
def run_fitting(
exp_cfg,
batch: Dict[str, Tensor],
body_model: nn.Module,
def_matrix: Tensor,
mask_ids: Optional = None
) -> Dict[str, Tensor]:
''' Runs fitting
'''
vertices = batch['vertices']
faces = batch['faces']
batch_size = len(vertices)
dtype, device = vertices.dtype, vertices.device
summary_steps = exp_cfg.get('summary_steps')
interactive = exp_cfg.get('interactive')
# Get the parameters from the model
var_dict = get_variables(batch_size, body_model)
# Build the optimizer object for the current batch
optim_cfg = exp_cfg.get('optim', {})
def_vertices = apply_deformation_transfer(def_matrix, vertices, faces)
if mask_ids is None:
f_sel = np.ones_like(body_model.faces[:, 0], dtype=np.bool_)
else:
f_per_v = [[] for _ in range(body_model.get_num_verts())]
[f_per_v[vv].append(iff) for iff, ff in enumerate(body_model.faces)
for vv in ff]
f_sel = list(set(tuple(sum([f_per_v[vv] for vv in mask_ids], []))))
vpe = get_vertices_per_edge(
body_model.v_template.detach().cpu().numpy(), body_model.faces[f_sel])
def log_closure():
return summary_closure(def_vertices, var_dict, body_model,
mask_ids=mask_ids)
edge_fitting_cfg = exp_cfg.get('edge_fitting', {})
edge_loss = build_loss(type='vertex-edge', gt_edges=vpe, est_edges=vpe,
**edge_fitting_cfg)
edge_loss = edge_loss.to(device=device)
vertex_fitting_cfg = exp_cfg.get('vertex_fitting', {})
vertex_loss = build_loss(**vertex_fitting_cfg)
vertex_loss = vertex_loss.to(device=device)
per_part = edge_fitting_cfg.get('per_part', True)
logger.info(f'Per-part: {per_part}')
# Optimize edge-based loss to initialize pose
if per_part:
for key, var in tqdm(var_dict.items(), desc='Parts'):
if 'pose' not in key:
continue
for jidx in tqdm(range(var.shape[1]), desc='Joints'):
part = torch.zeros(
[batch_size, 3], dtype=dtype, device=device,
requires_grad=True)
# Build the optimizer for the current part
optimizer_dict = build_optimizer([part], optim_cfg)
closure = build_edge_closure(
body_model, var_dict, edge_loss, optimizer_dict,
def_vertices, per_part=per_part, part_key=key, jidx=jidx,
part=part)
minimize(optimizer_dict['optimizer'], closure,
params=[part],
summary_closure=log_closure,
summary_steps=summary_steps,
interactive=interactive,
**optim_cfg)
with torch.no_grad():
var[:, jidx] = part
else:
optimizer_dict = build_optimizer(list(var_dict.values()), optim_cfg)
closure = build_edge_closure(
body_model, var_dict, edge_loss, optimizer_dict,
def_vertices, per_part=per_part)
minimize(optimizer_dict['optimizer'], closure,
params=var_dict.values(),
summary_closure=log_closure,
summary_steps=summary_steps,
interactive=interactive,
**optim_cfg)
if 'translation' in var_dict:
optimizer_dict = build_optimizer([var_dict['translation']], optim_cfg)
closure = build_vertex_closure(
body_model, var_dict,
optimizer_dict,
def_vertices,
vertex_loss=vertex_loss,
mask_ids=mask_ids,
per_part=False,
params_to_opt=[var_dict['translation']],
)
# Optimize translation
minimize(optimizer_dict['optimizer'],
closure,
params=[var_dict['translation']],
summary_closure=log_closure,
summary_steps=summary_steps,
interactive=interactive,
**optim_cfg)
# Optimize all model parameters with vertex-based loss
optimizer_dict = build_optimizer(list(var_dict.values()), optim_cfg)
closure = build_vertex_closure(
body_model, var_dict,
optimizer_dict,
def_vertices,
vertex_loss=vertex_loss,
per_part=False,
mask_ids=mask_ids)
minimize(optimizer_dict['optimizer'], closure,
params=list(var_dict.values()),
summary_closure=log_closure,
summary_steps=summary_steps,
interactive=interactive,
**optim_cfg)
param_dict = {}
for key, var in var_dict.items():
# Decode the axis-angles
if 'pose' in key or 'orient' in key:
param_dict[key] = batch_rodrigues(
var.reshape(-1, 3)).reshape(len(var), -1, 3, 3)
else:
# Simply pass the variable
param_dict[key] = var
body_model_output = body_model(
return_full_pose=True, get_skin=True, **param_dict)
var_dict.update(body_model_output)
var_dict['faces'] = body_model.faces
return var_dict
def forward(self, beta, theta, get_skin=False):
"""
Obtain SMPL with shape (beta) & pose (theta) inputs.
Theta includes the global rotation.
Args:
beta: N x 10
theta: N x 72 (with 3-D axis-angle rep)
Updates:
self.J_transformed: N x 24 x 3 joint location after shaping
& posing with beta and theta
Returns:
- joints: N x 19 or 14 x 3 joint locations depending on joint_type
If get_skin is True, also returns
- Verts: N x 6980 x 3
"""
num_batch = beta.shape[0]
# 1. Add shape blend shapes
# (N x 10) x (10 x 6890*3) = N x 6890 x 3
beta = torch.unsqueeze(beta, 0)
v_shaped = torch.bmm(beta, self.shapedirs).view(
-1, self.size[0], self.size[1]) + self.v_template
# 2. Infer shape-dependent joint locations.
Jx = torch.mm(v_shaped[:, :, 0], self.J_regressor)
Jy = torch.mm(v_shaped[:, :, 1], self.J_regressor)
Jz = torch.mm(v_shaped[:, :, 2], self.J_regressor)
J = torch.stack([Jx, Jy, Jz], 2)
# 3. Add pose blend shapes
# N x 24 x 3 x 3
Rs = batch_rodrigues(theta.contiguous().view(-1, 3)).view(-1, 24, 3, 3)
pose_feature = (Rs[:, 1:, :, :] - torch.eye(3).cuda()).view(-1, 207)
# (N x 207) x (207, 20670) -> N x 6890 x 3
v_posed = torch.mm(pose_feature, self.posedirs).view(
-1, self.size[0], self.size[1]) + v_shaped
# 4. Get the global joint location
self.J_transformed, A = batch_global_rigid_transformation(
Rs, J, self.parents, False)
# 5. Do skinning:
# W is N x 6890 x 24
W = self.weights.repeat(num_batch, 1).view(num_batch, -1, 24)
# (N x 6890 x 24) x (N x 24 x 16)
T = torch.bmm(W, A.view(num_batch, 24, 16)).view(num_batch, -1, 4, 4)
v_posed_homo = torch.cat(
[v_posed,
torch.ones((num_batch, v_posed.shape[1], 1)).cuda()], 2)
v_homo = torch.bmm(T.view(-1, 4, 4),
torch.unsqueeze(v_posed_homo, -1).view(-1, 4, 1))
v_homo = v_homo.view(num_batch, 6890, 4, 1)
verts = v_homo[:, :, :3, 0]
# Get cocoplus or lsp joints:
joint_x = torch.mm(verts[:, :, 0], self.joint_regressor)
joint_y = torch.mm(verts[:, :, 1], self.joint_regressor)
joint_z = torch.mm(verts[:, :, 2], self.joint_regressor)
joints = torch.stack([joint_x, joint_y, joint_z], dim=2)
if get_skin:
return verts, joints, Rs
else:
return joints