import os import matplotlib.pyplot as plt import numpy as np import torch import torch.nn.functional as torch_f SNAP_PARENT = [ 0, # 0's parent 0, # 1's parent 1, 2, 3, 0, # 5's parent 5, 6, 7, 0, # 9's parent 9, 10, 11, 0, # 13's parent 13, 14, 15, 0, # 17's parent 17, 18, 19, ] JOINT_ROOT_IDX = 9 REF_BONE_LINK = (0, 9) # mid mcp # bone indexes in 20 bones setting ID_ROOT_bone = np.array([0, 4, 8, 12, 16]) # ROOT_bone from wrist to MCP ID_PIP_bone = np.array([1, 5, 9, 13, 17]) # PIP_bone from MCP to PIP ID_DIP_bone = np.array([2, 6, 10, 14, 18]) # DIP_bone from PIP to DIP ID_TIP_bone = np.array([3, 7, 11, 15, 19]) # TIP_bone from DIP to TIP def plot_hull(theta, hull): del_rdp_hull = hull.detach().cpu().numpy() theta = theta.detach().cpu().numpy() fig = plt.figure() figManager = plt.get_current_fig_manager() figManager.window.showMaximized() ax = fig.add_subplot(111) ax.scatter(theta[:, 0], theta[:, 1], s=10, c='r') ax.set_xlabel("flexion") ax.set_ylabel("abduction") plt.plot(del_rdp_hull[:, 0], del_rdp_hull[:, 1], '-yo', linewidth=2) plt.xticks(np.arange(-3, 4, 1)) plt.yticks(np.arange(-2, 2, 0.5)) plt.show() def two_norm(a): ''' Args: a: B*M*2 or B*M*3 Returns: ''' return torch.norm(a, dim=-1) def one_norm(a): ''' Args: a: B*M*2 or B*M*3 Returns: ''' return torch.norm(a, dim=-1, p=1) def calculate_joint_angle_loss(thetas, hulls): ''' Args: Theta: B*15*2 hulls: list Returns: ''' loss = torch.Tensor([0]).cuda() for i in range(15): # print("i=",i) hull = hulls[i] # (M*2) theta = thetas[:, i] # (B*2) hull = torch.cat((hull, hull[0].unsqueeze(0)), dim=0) v = (hull[1:] - hull[:-1]).unsqueeze(0) # (M-1)*2 w = - hull[:-1].unsqueeze(0) + theta.unsqueeze(1).repeat(1, hull[:-1].shape[0], 1) # B*(M-1)*2 cross_product_2d = w[:, :, 0] * v[:, :, 1] - w[:, :, 1] * v[:, :, 0] tmp = torch.sum(cross_product_2d < 1e-6, dim=-1) is_outside = (tmp != (hull.shape[0] - 1)) if not torch.sum(is_outside): sub_loss = torch.Tensor([0]).cuda() else: outside_theta = theta[is_outside] outside_theta = outside_theta.unsqueeze(1).repeat(1, hull[:-1].shape[0], 1) w_outside = - hull[:-1].unsqueeze(0) + outside_theta # B*(M-1)*2 t = torch.clamp(inner_product(w_outside, v) / (two_norm(v) ** 2), min=0, max=1).unsqueeze(2) p = hull[:-1] + t * v D = one_norm(torch.cos(outside_theta) - torch.cos(p)) + one_norm(torch.sin(outside_theta) - torch.sin(p)) sub_loss = torch.sum(torch.min(D, dim=-1)[0]) vis = 0 if vis: print(theta) plot_hull(theta, hull) loss += sub_loss loss /= (15 * thetas.shape[0]) return loss def angle_between(v1, v2): epsilon = 1e-7 cos = torch_f.cosine_similarity(v1, v2, dim=-1).clamp(-1 + epsilon, 1 - epsilon) # (B) theta = torch.acos(cos) # (B) return theta def normalize(vec): return torch_f.normalize(vec, p=2, dim=-1) def inner_product(x1, x2): return torch.sum(x1 * x2, dim=-1) def cross_product(x1, x2): return torch.cross(x1, x2, dim=-1) def axangle2mat_torch(axis, angle, is_normalized=False): """ Rotation matrix for rotation angle `angle` around `axis` Parameters ---------- axis : [B,M, 3] element sequence vector specifying axis for rotation. angle :[B,M, ] scalar angle of rotation in radians. is_normalized : bool, optional True if `axis` is already normalized (has norm of 1). Default False. Returns ------- mat : array shape (B, M,3,3) rotation matrix for specified rotation Notes ----- From: http://en.wikipedia.org/wiki/Rotation_matrix#Axis_and_angle """ B = axis.shape[0] M = axis.shape[1] if not is_normalized: norm_axis = axis.norm(p=2, dim=-1, keepdim=True) normed_axis = axis / norm_axis else: normed_axis = axis x, y, z = normed_axis[:, :, 0], normed_axis[:, :, 1], normed_axis[:, :, 2] c = torch.cos(angle) s = torch.sin(angle) C = 1 - c xs = x * s; ys = y * s; zs = z * s # noqa xC = x * C; yC = y * C; zC = z * C # noqa xyC = x * yC; yzC = y * zC; zxC = z * xC # noqa TMP = torch.stack([x * xC + c, xyC - zs, zxC + ys, xyC + zs, y * yC + c, yzC - xs, zxC - ys, yzC + xs, z * zC + c], dim=-1) return TMP.reshape(B, M, 3, 3) def interval_loss(value, min, max): ''' calculate interval loss Args: value: B*M max: M min: M Returns: ''' batch_3d_size = value.shape[0] min = min.repeat(value.shape[0], 1) max = max.repeat(value.shape[0], 1) loss1 = torch.max(min - value, torch.Tensor([0]).cuda()) loss2 = torch.max(value - max, torch.Tensor([0]).cuda()) loss = (loss1 + loss2).sum() loss /= (batch_3d_size * value.shape[1]) return loss class BMCLoss: def __init__( self, lambda_bl=0., lambda_rb=0., lambda_a=0., ): self.lambda_bl = lambda_bl self.lambda_rb = lambda_rb self.lambda_a = lambda_a self.lp = os.path.join(os.path.dirname(__file__), "../../_DATA/BMC") # resolved_path = os.path.abspath(self.lp) # print(f"Resolved path: {resolved_path}") # print(f"Current working directory: {os.getcwd()}") # exit() self.bone_len_max = np.load(os.path.join(self.lp, "bone_len_max.npy")) self.bone_len_min = np.load(os.path.join(self.lp, "bone_len_min.npy")) self.rb_curvatures_max = np.load(os.path.join(self.lp, "curvatures_max.npy")) self.rb_curvatures_min = np.load(os.path.join(self.lp, "curvatures_min.npy")) self.rb_PHI_max = np.load(os.path.join(self.lp, "PHI_max.npy")) self.rb_PHI_min = np.load(os.path.join(self.lp, "PHI_min.npy")) self.joint_angle_limit = np.load(os.path.join(self.lp, "CONVEX_HULLS.npy"), allow_pickle=True) LEN_joint_angle_limit = len(self.joint_angle_limit) self.bl_max = torch.from_numpy(self.bone_len_max).float().cuda() self.bl_min = torch.from_numpy(self.bone_len_min).float().cuda() self.rb_curvatures_max = torch.from_numpy(self.rb_curvatures_max).float().cuda() self.rb_curvatures_min = torch.from_numpy(self.rb_curvatures_min).float().cuda() self.rb_PHI_max = torch.from_numpy(self.rb_PHI_max).float().cuda() self.rb_PHI_min = torch.from_numpy(self.rb_PHI_min).float().cuda() self.joint_angle_limit = [torch.from_numpy(self.joint_angle_limit[i]).float().cuda() for i in range(LEN_joint_angle_limit)] def compute_loss(self, ori_joints, valid_mask): ''' Args: joints: B*21*3 Returns: ''' joints = ori_joints[valid_mask] print(joints.shape, ori_joints.shape) batch_size = joints.shape[0] final_loss = torch.Tensor([0]).cuda() BMC_losses = {"bmc_bl": torch.Tensor([0]).cuda(), "bmc_rb": torch.Tensor([0]).cuda(), "bmc_a": torch.Tensor([0]).cuda()} if (self.lambda_bl < 1e-6) and (self.lambda_rb < 1e-6) and (self.lambda_a < 1e-6): return final_loss, BMC_losses ALL_bones = [ ( joints[:, i, :] - joints[:, SNAP_PARENT[i], :] ) for i in range(21) ] ALL_bones = torch.stack(ALL_bones[1:], dim=1) # (B,20,3) ROOT_bones = ALL_bones[:, ID_ROOT_bone] # (B,5,3) PIP_bones = ALL_bones[:, ID_PIP_bone] DIP_bones = ALL_bones[:, ID_DIP_bone] TIP_bones = ALL_bones[:, ID_TIP_bone] ALL_Z_axis = normalize(ALL_bones) PIP_Z_axis = ALL_Z_axis[:, ID_ROOT_bone] DIP_Z_axis = ALL_Z_axis[:, ID_PIP_bone] TIP_Z_axis = ALL_Z_axis[:, ID_DIP_bone] normals = normalize(cross_product(ROOT_bones[:, 1:], ROOT_bones[:, :-1])) # compute loss of bone length bl_loss = torch.Tensor([0]).cuda() if self.lambda_bl: bls = two_norm(ALL_bones) # (B,20,1) bl_loss = interval_loss(value=bls, min=self.bl_min, max=self.bl_max) final_loss += self.lambda_bl * bl_loss BMC_losses["bmc_bl"] = bl_loss # compute loss of Root bones rb_loss = torch.Tensor([0]).cuda() if self.lambda_rb: edge_normals = torch.zeros_like(ROOT_bones).cuda() # (B,5,3) edge_normals[:, [0, 4]] = normals[:, [0, 3]] edge_normals[:, 1:4] = normalize(normals[:, 1:4] + normals[:, :3]) curvatures = inner_product(edge_normals[:, 1:] - edge_normals[:, :4], ROOT_bones[:, 1:] - ROOT_bones[:, :4]) / \ (two_norm(ROOT_bones[:, 1:] - ROOT_bones[:, :4]) ** 2) PHI = angle_between(ROOT_bones[:, :4], ROOT_bones[:, 1:]) # (B) rb_loss = interval_loss(value=curvatures, min=self.rb_curvatures_min, max=self.rb_curvatures_max) + \ interval_loss(value=PHI, min=self.rb_PHI_min, max=self.rb_PHI_max) final_loss += self.lambda_rb * rb_loss BMC_losses["bmc_rb"] = rb_loss # compute loss of Joint angles a_loss = torch.Tensor([0]).cuda() if self.lambda_a: # PIP bones PIP_X_axis = torch.zeros([batch_size, 5, 3]).cuda() # (B,5,3) PIP_X_axis[:, [0, 1, 4], :] = -normals[:, [0, 1, 3]] PIP_X_axis[:, 2:4] = -normalize(normals[:, 2:4] + normals[:, 1:3]) # (B,2,3) PIP_Y_axis = normalize(cross_product(PIP_Z_axis, PIP_X_axis)) # (B,5,3) PIP_bones_xz = PIP_bones - inner_product(PIP_bones, PIP_Y_axis).unsqueeze(2) * PIP_Y_axis PIP_theta_flexion = angle_between(PIP_bones_xz, PIP_Z_axis) # in global coordinate (B) PIP_theta_abduction = angle_between(PIP_bones_xz, PIP_bones) # in global coordinate (B) # x-component of the bone vector tmp = inner_product(PIP_bones, PIP_X_axis) PIP_theta_flexion = torch.where(tmp < 1e-6, -PIP_theta_flexion, PIP_theta_flexion) # y-component of the bone vector tmp = torch.sum((PIP_bones * PIP_Y_axis), dim=-1) PIP_theta_abduction = torch.where(tmp < 1e-6, -PIP_theta_abduction, PIP_theta_abduction) temp_axis = normalize(cross_product(PIP_Z_axis, PIP_bones)) temp_alpha = angle_between(PIP_Z_axis, PIP_bones) # alpha belongs to [pi/2, pi] temp_R = axangle2mat_torch(axis=temp_axis, angle=temp_alpha, is_normalized=True) # DIP bones DIP_X_axis = torch.matmul(temp_R, PIP_X_axis.unsqueeze(3)).squeeze() DIP_Y_axis = torch.matmul(temp_R, PIP_Y_axis.unsqueeze(3)).squeeze() DIP_bones_xz = DIP_bones - inner_product(DIP_bones, DIP_Y_axis).unsqueeze(2) * DIP_Y_axis DIP_theta_flexion = angle_between(DIP_bones_xz, DIP_Z_axis) # in global coordinate DIP_theta_abduction = angle_between(DIP_bones_xz, DIP_bones) # in global coordinate # x-component of the bone vector tmp = inner_product(DIP_bones, DIP_X_axis) DIP_theta_flexion = torch.where(tmp < 0, -DIP_theta_flexion, DIP_theta_flexion) # y-component of the bone vector tmp = inner_product(DIP_bones, DIP_Y_axis) DIP_theta_abduction = torch.where(tmp < 0, -DIP_theta_abduction, DIP_theta_abduction) temp_axis = normalize(cross_product(DIP_Z_axis, DIP_bones)) temp_alpha = angle_between(DIP_Z_axis, DIP_bones) # alpha belongs to [pi/2, pi] temp_R = axangle2mat_torch(axis=temp_axis, angle=temp_alpha, is_normalized=True) # TIP bones TIP_X_axis = torch.matmul(temp_R, DIP_X_axis.unsqueeze(3)).squeeze() TIP_Y_axis = torch.matmul(temp_R, DIP_Y_axis.unsqueeze(3)).squeeze() TIP_bones_xz = TIP_bones - inner_product(TIP_bones, TIP_Y_axis).unsqueeze(2) * TIP_Y_axis TIP_theta_flexion = angle_between(TIP_bones_xz, TIP_Z_axis) # in global coordinate TIP_theta_abduction = angle_between(TIP_bones_xz, TIP_bones) # in global coordinate # x-component of the bone vector tmp = inner_product(TIP_bones, TIP_X_axis) TIP_theta_flexion = torch.where(tmp < 1e-6, -TIP_theta_flexion, TIP_theta_flexion) # y-component of the bone vector tmp = inner_product(TIP_bones, TIP_Y_axis) TIP_theta_abduction = torch.where(tmp < 1e-6, -TIP_theta_abduction, TIP_theta_abduction) # ALL ALL_theta_flexion = torch.cat((PIP_theta_flexion, DIP_theta_flexion, TIP_theta_flexion), dim=-1) ALL_theta_abduction = torch.cat((PIP_theta_abduction, DIP_theta_abduction, TIP_theta_abduction), dim=-1) ALL_theta = torch.stack((ALL_theta_flexion, ALL_theta_abduction), dim=-1) a_loss = calculate_joint_angle_loss(ALL_theta, self.joint_angle_limit) final_loss += self.lambda_a * a_loss BMC_losses["bmc_a"] = a_loss return final_loss, BMC_losses