def forward(self,
pose=None,
pose2=None,
shape_id=None,
hiRes=False,
residual=False,
requires_skeleton=False):
if hiRes:
raise Exception(
"High resolution support is abandoned in MIT license's SMPL implementation"
)
if pose is None:
pose = torch.zeros((1, self.bone_num * 3), device=self.device)
if shape_id is None:
shape_id = [
random.randint(0,
len(self) - 1) for _ in range(pose.shape[0])
]
pose = pose.reshape(pose.shape[0], -1, 3)
if pose2 is not None:
pose2 = pose2.reshape(pose.shape[0], -1, 3)
t_poses = self.t_pose_list[shape_id]
offsets = self.offset_list[shape_id]
if not requires_skeleton:
root_loc = offsets[:, 0, :]
else:
root_loc = offsets
if not hiRes:
t_poses = t_poses[:, :self.smpl.num_verts]
weight = self.weight
smpl = self.smpl
else:
weight = self.weight_hires
smpl = self.smpl_hires
mat = self.fk.forward(aa2mat(pose), offsets)
if pose2 is not None:
mat2 = self.fk.forward(aa2mat(pose2), offsets)
if residual:
offset = smpl.pose_blendshapes(pose.reshape(pose.shape[0], -1))
offset2 = smpl.pose_blendshapes(pose2.reshape(
pose2.shape[0], -1)) if pose2 is not None else None
else:
offset = 0
offset2 = 0
res = deform_with_offset(t_poses, weight, mat, offset=offset)
if pose2 is not None:
res2 = deform_with_offset(t_poses, weight, mat2, offset=offset2)
return res, res2, t_poses, root_loc
else:
return res, t_poses, root_loc
def forward_multipose(self,
pose,
n_shape,
residual=False,
requires_skeleton=False):
shape_id = [random.randint(0, len(self) - 1) for _ in range(n_shape)]
pose = pose.reshape(pose.shape[0], -1, 3)
t_poses = self.t_pose_list[shape_id]
t_poses = t_poses[:, :self.smpl.num_verts]
offsets = self.offset_list[shape_id]
if not requires_skeleton:
root_loc = offsets[:, 0, :]
else:
root_loc = offsets
if residual:
v_offset = self.smpl.pose_blendshapes(
pose.reshape(pose.shape[0], -1))
else:
v_offset = 0
deformed = []
posemat = aa2mat(pose)
for i in range(n_shape):
mat = self.fk.forward(posemat, offsets[[i]])
deformed.append(
deform_with_offset(t_poses[[i]],
self.weight,
mat,
offset=v_offset)[None, :])
deformed = torch.cat(deformed, dim=0)
return deformed, t_poses, root_loc
def forward_with_shape(self,
shape_id,
shapes,
poses=None,
hiRes=False,
residual=False):
# todo: reduce the number of calling smpl.forward()
if poses is None:
poses = torch.zeros((1, self.bone_num * 3), device=self.device)
if hiRes:
smpl = self.smpl_hires
weight = self.weight_hires
else:
smpl = self.smpl
weight = self.weight
t_poses, _ = smpl.forward(torch.zeros_like(poses), shapes)
offsets = smpl.get_offset(shapes)
cloths = self.cloth_all[shape_id, :smpl.num_verts]
clothed_t_poses = t_poses + cloths
mat = self.fk.forward(aa2mat(poses.reshape(poses.shape[0], -1, 3)),
offsets)
if residual:
offset = smpl.pose_blendshapes(poses)
else:
offset = 0
res = deform_with_offset(clothed_t_poses, weight, mat, offset=offset)
return res, clothed_t_poses
def get_coff(self, pose):
"""
@return: (batch_size, n_vert, n_basis_per_bone)
"""
batch_size = pose.shape[0]
device = pose.device
if len(pose.shape) == 2:
pose_repr = aa2mat(pose.reshape(pose.shape[0], -1, 3))
elif len(pose.shape) == 4:
pose_repr = pose.reshape(batch_size, -1, 3, 3)
else:
raise Exception('Wrong input format')
pose_repr = pose_repr[:, 1:]
pose_repr = pose_repr.reshape(-1, 9)
identical = torch.eye(3, device=device).reshape(-1)
pose_repr = pose_repr - identical
pose_repr = pose_repr.reshape(pose.shape[0], self.n_joint, -1)
coff = []
for i in range(pose_repr.shape[1]):
coff.append(self.coff_branch[i](pose_repr[:, i]).unsqueeze(1))
coff = torch.cat(coff, dim=1)
return coff
def forward(self,
shape_id,
poses=None,
hiRes=False,
v_offset=0,
residual=False):
if poses is None:
poses = torch.zeros((1, self.bone_num * 3), device=self.device)
poses = poses.reshape(poses.shape[0], -1, 3)
t_poses = self.t_pose_list[shape_id]
offsets = self.offset_list[shape_id]
# offsets[:, 0] = 0
if not hiRes:
t_poses = t_poses[:, :self.smpl.num_verts]
weight = self.weight
smpl = self.smpl
else:
weight = self.weight_hires
smpl = self.smpl_hires
mat = self.fk.forward(aa2mat(poses), offsets)
if residual:
offset = smpl.pose_blendshapes(poses)
else:
offset = 0
res = deform_with_offset(t_poses,
weight,
mat,
offset=v_offset + offset)
return res, t_poses
def run_single_mesh(verts, topo_id, pose, env_model, res_model):
skinning_weight, skeleton = eval_envelop(verts, topo_id, env_model)
offset = eval_residual(verts, topo_id, pose, res_model)
local_mat = aa2mat(pose.reshape(pose.shape[0], -1, 3))
global_mat = env_model.fk.forward(local_mat, skeleton.unsqueeze(0))
mask = env_model.rec_model.topo_loader.v_masks[topo_id]
verts = verts[mask]
vs = deform_with_offset(verts, skinning_weight, global_mat, offset)
vs_lbs = deform_with_offset(verts, skinning_weight, global_mat)
return skinning_weight, skeleton, vs, vs_lbs
def forward(self, t_pose, pose, topo_id, skeletons):
self.loss = 0
self.set_topology(topo_id)
self.t_pose_mapped = self.rec_model.apply_topo(t_pose)
# Normalize
if self.args.normalize:
self.bb = BoundingBox(self.t_pose_mapped)
else:
self.bb = BoundingBox()
self.t_pose_mapped = self.bb.normalize(self.t_pose_mapped)
# Get skinning weight
with torch.no_grad():
self.att, self.att_vert = self.forward_att(self.t_pose_mapped)
self.att = self.att.detach()
self.att_vert = self.att_vert.detach()
# Generate blend shapes
self.deep_v = self.forward_geometry(self.t_pose_mapped)
offset_input = torch.cat((self.deep_v, self.att), dim=-1)
basis = self.forward_gen(offset_input)
basis = basis.reshape(basis.shape[:-1] +
(-1, 3)) # (batch_size, n_vert, n_basis, 3)
self.basis = basis
self.res_verts = basis
self.t_pose_mapped = self.bb.denormalize(self.t_pose_mapped)
self.basis = self.bb.denormalize_basis(self.basis)
if self.args.basis_only:
return self.basis
# Generate deformed shape
self.res_verts = []
self.offsets = []
local_mat = aa2mat(pose.reshape(pose.shape[0], -1, 3))
for i in range(t_pose.shape[0]):
offset = self.forward_residual(local_mat, self.basis[i],
self.att_vert[i])
self.offsets.append(offset)
if skeletons is None:
continue
trans_b = self.fk.forward(local_mat, skeletons[[i]])
res_vert = deform_with_offset(
self.t_pose_mapped[[i]].expand(pose.shape[0], -1, -1),
self.att_vert[[i]], trans_b, offset)
self.res_verts.append(res_vert[None, ...])
if skeletons is not None:
self.res_verts = torch.cat(self.res_verts, dim=0)
return self.res_verts
def forward_lbs(self, poses, shapes=None, v_offsets=0, dqs=False):
if shapes is None:
shapes = torch.zeros((poses.shape[0], 10), device=poses.device)
t_pose = self.forward(torch.zeros_like(poses), shapes)[0]
offsets = self.get_offset(shapes)
local_mat = aa2mat(poses.reshape(poses.shape[0], -1, 3))
global_mat = self.fk.forward(local_mat, offsets)
return deform_with_offset(t_pose,
self.th_weights,
global_mat,
offset=v_offsets,
dqs=dqs)
def forward(self, t_pose, pose, topo_id, weight=None, pose_ee=None):
self.set_topology(topo_id)
self.t_pose_mapped = self.rec_model.apply_topo(t_pose)
# Normalize
if self.args.normalize:
self.bb = BoundingBox(self.t_pose_mapped)
else:
self.bb = BoundingBox()
self.t_pose_mapped = self.bb.normalize(self.t_pose_mapped)
# Generate skinning weight (attention)
self.att, self.att_vert = self.forward_att(self.t_pose_mapped)
self.deep_v = self.forward_geometry(self.t_pose_mapped)
# Deal with skeleton generation
deep_s = attention_pooling(self.deep_v, self.att)
deep_s = deep_s.reshape(deep_s.shape[0], -1)
skeleton = self.models['gen'](deep_s)
skeleton = skeleton.reshape(deep_s.shape[0], -1, 3)
self.skeleton = skeleton = self.bb.denormalize_offset(skeleton)
self.t_pose_mapped = self.bb.denormalize(self.t_pose_mapped)
if pose is None:
return
# Generate deformed shape
if pose_ee is not None:
pose = torch.cat((pose, pose_ee), dim=0)
pose = pose.reshape(pose.shape[0], -1, 3)
self.local_mat = aa2mat(pose)
if pose_ee is not None:
self.local_mat_ee = self.local_mat[-pose_ee.shape[0]:]
self.local_mat = self.local_mat[:-pose_ee.shape[0]]
trans_b = self.fk.forward(self.local_mat, skeleton)
self.res_verts = deform_with_offset(self.t_pose_mapped, self.att_vert,
trans_b)
if pose_ee is not None:
trans_b_ee = self.fk.forward(self.local_mat_ee, skeleton)
self.res_verts_ee = deform_with_offset(self.t_pose_mapped,
self.att_vert, trans_b_ee)
return self.res_verts
def run_single_mesh(verts,
topo_id,
pose,
env_model,
res_model,
requires_lbs=False):
skinning_weight, skeleton = eval_envelop(verts, topo_id, env_model)
if res_model is not None:
offset, basis, coff = eval_residual(verts, topo_id, pose, res_model)
else:
offset = 0
basis = None
coff = None
local_mat = aa2mat(pose.reshape(pose.shape[0], -1, 3))
global_mat = env_model.fk.forward(local_mat, skeleton.unsqueeze(0))
mask = env_model.rec_model.topo_loader.v_masks[topo_id]
verts = verts[mask]
vs = deform_with_offset(verts, skinning_weight, global_mat, offset)
if requires_lbs:
vs_lbs = deform_with_offset(verts, skinning_weight, global_mat)
return skinning_weight, skeleton, vs, vs_lbs, basis, coff
return skinning_weight, skeleton, vs, basis, coff