| import torch | |
| from utils.transforms import quat2mat, repr6d2mat, euler2mat | |
| class ForwardKinematics: | |
| def __init__(self, parents, offsets=None): | |
| self.parents = parents | |
| if offsets is not None and len(offsets.shape) == 2: | |
| offsets = offsets.unsqueeze(0) | |
| self.offsets = offsets | |
| def forward(self, rots, offsets=None, global_pos=None): | |
| """ | |
| Forward Kinematics: returns a per-bone transformation | |
| @param rots: local joint rotations (batch_size, bone_num, 3, 3) | |
| @param offsets: (batch_size, bone_num, 3) or None | |
| @param global_pos: global_position: (batch_size, 3) or keep it as in offsets (default) | |
| @return: (batch_szie, bone_num, 3, 4) | |
| """ | |
| rots = rots.clone() | |
| if offsets is None: | |
| offsets = self.offsets.to(rots.device) | |
| if global_pos is None: | |
| global_pos = offsets[:, 0] | |
| pos = torch.zeros((rots.shape[0], rots.shape[1], 3), device=rots.device) | |
| rest_pos = torch.zeros_like(pos) | |
| res = torch.zeros((rots.shape[0], rots.shape[1], 3, 4), device=rots.device) | |
| pos[:, 0] = global_pos | |
| rest_pos[:, 0] = offsets[:, 0] | |
| for i, p in enumerate(self.parents): | |
| if i != 0: | |
| rots[:, i] = torch.matmul(rots[:, p], rots[:, i]) | |
| pos[:, i] = torch.matmul(rots[:, p], offsets[:, i].unsqueeze(-1)).squeeze(-1) + pos[:, p] | |
| rest_pos[:, i] = rest_pos[:, p] + offsets[:, i] | |
| res[:, i, :3, :3] = rots[:, i] | |
| res[:, i, :, 3] = torch.matmul(rots[:, i], -rest_pos[:, i].unsqueeze(-1)).squeeze(-1) + pos[:, i] | |
| return res | |
| def accumulate(self, local_rots): | |
| """ | |
| Get global joint rotation from local rotations | |
| @param local_rots: (batch_size, n_bone, 3, 3) | |
| @return: global_rotations | |
| """ | |
| res = torch.empty_like(local_rots) | |
| for i, p in enumerate(self.parents): | |
| if i == 0: | |
| res[:, i] = local_rots[:, i] | |
| else: | |
| res[:, i] = torch.matmul(res[:, p], local_rots[:, i]) | |
| return res | |
| def unaccumulate(self, global_rots): | |
| """ | |
| Get local joint rotation from global rotations | |
| @param global_rots: (batch_size, n_bone, 3, 3) | |
| @return: local_rotations | |
| """ | |
| res = torch.empty_like(global_rots) | |
| inv = torch.empty_like(global_rots) | |
| for i, p in enumerate(self.parents): | |
| if i == 0: | |
| inv[:, i] = global_rots[:, i].transpose(-2, -1) | |
| res[:, i] = global_rots[:, i] | |
| continue | |
| res[:, i] = torch.matmul(inv[:, p], global_rots[:, i]) | |
| inv[:, i] = torch.matmul(res[:, i].transpose(-2, -1), inv[:, p]) | |
| return res | |
| class ForwardKinematicsJoint: | |
| def __init__(self, parents, offset): | |
| self.parents = parents | |
| self.offset = offset | |
| ''' | |
| rotation should have shape batch_size * Joint_num * (3/4) * Time | |
| position should have shape batch_size * 3 * Time | |
| offset should have shape batch_size * Joint_num * 3 | |
| output have shape batch_size * Time * Joint_num * 3 | |
| ''' | |
| def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset=None, | |
| world=True): | |
| ''' | |
| if not quater and rotation.shape[-2] != 3: raise Exception('Unexpected shape of rotation') | |
| if quater and rotation.shape[-2] != 4: raise Exception('Unexpected shape of rotation') | |
| rotation = rotation.permute(0, 3, 1, 2) | |
| position = position.permute(0, 2, 1) | |
| ''' | |
| if rotation.shape[-1] == 6: | |
| transform = repr6d2mat(rotation) | |
| elif rotation.shape[-1] == 4: | |
| norm = torch.norm(rotation, dim=-1, keepdim=True) | |
| rotation = rotation / norm | |
| transform = quat2mat(rotation) | |
| elif rotation.shape[-1] == 3: | |
| transform = euler2mat(rotation) | |
| else: | |
| raise Exception('Only accept quaternion rotation input') | |
| result = torch.empty(transform.shape[:-2] + (3,), device=position.device) | |
| if offset is None: | |
| offset = self.offset | |
| offset = offset.reshape((-1, 1, offset.shape[-2], offset.shape[-1], 1)) | |
| result[..., 0, :] = position | |
| for i, pi in enumerate(self.parents): | |
| if pi == -1: | |
| assert i == 0 | |
| continue | |
| result[..., i, :] = torch.matmul(transform[..., pi, :, :], offset[..., i, :, :]).squeeze() | |
| transform[..., i, :, :] = torch.matmul(transform[..., pi, :, :].clone(), transform[..., i, :, :].clone()) | |
| if world: result[..., i, :] += result[..., pi, :] | |
| return result | |
| class InverseKinematicsJoint: | |
| def __init__(self, rotations: torch.Tensor, positions: torch.Tensor, offset, parents, constrains): | |
| self.rotations = rotations.detach().clone() | |
| self.rotations.requires_grad_(True) | |
| self.position = positions.detach().clone() | |
| self.position.requires_grad_(True) | |
| self.parents = parents | |
| self.offset = offset | |
| self.constrains = constrains | |
| self.optimizer = torch.optim.Adam([self.position, self.rotations], lr=1e-3, betas=(0.9, 0.999)) | |
| self.criteria = torch.nn.MSELoss() | |
| self.fk = ForwardKinematicsJoint(parents, offset) | |
| self.glb = None | |
| def step(self): | |
| self.optimizer.zero_grad() | |
| glb = self.fk.forward(self.rotations, self.position) | |
| loss = self.criteria(glb, self.constrains) | |
| loss.backward() | |
| self.optimizer.step() | |
| self.glb = glb | |
| return loss.item() | |
| class InverseKinematicsJoint2: | |
| def __init__(self, rotations: torch.Tensor, positions: torch.Tensor, offset, parents, constrains, cid, | |
| lambda_rec_rot=1., lambda_rec_pos=1., use_velo=False): | |
| self.use_velo = use_velo | |
| self.rotations_ori = rotations.detach().clone() | |
| self.rotations = rotations.detach().clone() | |
| self.rotations.requires_grad_(True) | |
| self.position_ori = positions.detach().clone() | |
| self.position = positions.detach().clone() | |
| if self.use_velo: | |
| self.position[1:] = self.position[1:] - self.position[:-1] | |
| self.position.requires_grad_(True) | |
| self.parents = parents | |
| self.offset = offset | |
| self.constrains = constrains.detach().clone() | |
| self.cid = cid | |
| self.lambda_rec_rot = lambda_rec_rot | |
| self.lambda_rec_pos = lambda_rec_pos | |
| self.optimizer = torch.optim.Adam([self.position, self.rotations], lr=1e-3, betas=(0.9, 0.999)) | |
| self.criteria = torch.nn.MSELoss() | |
| self.fk = ForwardKinematicsJoint(parents, offset) | |
| self.glb = None | |
| def step(self): | |
| self.optimizer.zero_grad() | |
| if self.use_velo: | |
| position = torch.cumsum(self.position, dim=0) | |
| else: | |
| position = self.position | |
| glb = self.fk.forward(self.rotations, position) | |
| self.constrain_loss = self.criteria(glb[:, self.cid], self.constrains) | |
| self.rec_loss_rot = self.criteria(self.rotations, self.rotations_ori) | |
| self.rec_loss_pos = self.criteria(self.position, self.position_ori) | |
| loss = self.constrain_loss + self.rec_loss_rot * self.lambda_rec_rot + self.rec_loss_pos * self.lambda_rec_pos | |
| loss.backward() | |
| self.optimizer.step() | |
| self.glb = glb | |
| return loss.item() | |
| def get_position(self): | |
| if self.use_velo: | |
| position = torch.cumsum(self.position.detach(), dim=0) | |
| else: | |
| position = self.position.detach() | |
| return position | |