Spaces:
Paused
Paused
| import numpy as np | |
| import random | |
| from NLFPoseExtract.nlf_draw import intrinsic_matrix_from_field_of_view, process_data_to_COCO_format, p3d_to_p2d | |
| import torch | |
| # reshapePool只负责形变，骨骼偏移、丢弃等得从draw层来做 | |
| class reshapePool3d: | |
| def __init__(self, reshape_type, height, width): # 对每个视频只初始化一次 | |
| self.reshape_type = reshape_type | |
| self.height = height | |
| self.width = width | |
| self.shoulder_alpha = 0 | |
| self.upper_arm_alpha = 0 | |
| self.forearm_alpha = 0 | |
| self.body_alpha = 0 | |
| self.thigh_alpha = 0 | |
| self.calf_alpha = 0 | |
| self.face_alpha = random.choices([-0.4, -0.2, 0, 0.2, 0.4], weights=[0.2, 0.15, 0.3, 0.15, 0.2], k=1)[0] | |
| self.body_reshape_methods = [ | |
| self.reshape_body, | |
| self.reshape_arm, | |
| self.reshape_leg, | |
| self.reshape_shoulder, | |
| ] | |
| self.face_reshape_methods = [ | |
| self.reshape_face, | |
| ] | |
| self.body_offset_selected_methods = [] | |
| options = ["normal_human", "dwarf", "slender", "elf", "random_long_arm_long_leg", "king-kong"] | |
| if self.reshape_type == "low": | |
| weights = [0.8, 0, 0, 0.1, 0.1, 0] | |
| self.body_offset_selected_methods = [] | |
| if self.reshape_type == "normal": | |
| weights = [0.4, 0.1, 0.1, 0.1, 0.2, 0.1] | |
| elif self.reshape_type == "high": | |
| weights = [0.3, 0.2, 0.1, 0.1, 0.2, 0.1] | |
| elif self.reshape_type == "dongman": | |
| weights = [0.1, 0.2, 0.2, 0.2, 0.2, 0.1] | |
| self.face_alpha = random.choices([-0.4, -0.2, 0.4], weights=[0.4, 0.2, 0.4], k=1)[0] | |
| choice = random.choices(options, weights=weights, k=1)[0] | |
| self.aug_init(choice) | |
| def pose_reshape_2d_for_face(self, alpha, candidate, face, subset, body_anchor_point, body_affected_points): | |
| anchor_x, anchor_y = candidate[body_anchor_point] | |
| if subset[body_anchor_point] == -1: | |
| return | |
| for face_point_idx in range(len(face)): | |
| face_point_x, face_point_y = face[face_point_idx] | |
| if face_point_x == -1 or face_point_y == -1: | |
| continue | |
| vector_x = face_point_x - anchor_x | |
| vector_y = face_point_y - anchor_y | |
| offset_x = vector_x * alpha | |
| offset_y = vector_y * alpha | |
| face[face_point_idx] = [face_point_x + offset_x, face_point_y + offset_y] | |
| for body_affected_point_idx in body_affected_points: | |
| body_point_x, body_point_y = candidate[body_affected_point_idx] | |
| if subset[body_affected_point_idx] == -1 or body_point_x == -1 or body_point_y == -1: | |
| continue | |
| vector_x = body_point_x - anchor_x | |
| vector_y = body_point_y - anchor_y | |
| offset_x = vector_x * alpha | |
| offset_y = vector_y * alpha | |
| candidate[body_affected_point_idx] = [body_point_x + offset_x, body_point_y + offset_y] | |
| def pose_reshape_3d(self, alpha, smpl_joints, candidate, subset, left_hand, right_hand, face, | |
| anchor_point, end_point, | |
| affected_body_points): | |
| """ | |
| # joints3d: n, 24, 3 | |
| alpha: 变化比例 | |
| anchor_point: 起始点 | |
| anchor_part: 起始点所属的身体部分，0代表body candidate, 1代表face, 2代表hand | |
| end_point: 中止点 | |
| """ | |
| if torch.sum(smpl_joints[anchor_point]) == 0: | |
| right_hand[:] = -1 | |
| left_hand[:] = -1 | |
| face[:] = -1 | |
| subset[:] = -1 | |
| return | |
| anchor_x, anchor_y, anchor_z = smpl_joints[anchor_point] | |
| end_x, end_y, end_z = smpl_joints[end_point] | |
| vector_x = end_x - anchor_x | |
| vector_y = end_y - anchor_y | |
| vector_z = end_z - anchor_z | |
| offset_x = (vector_x * alpha).item() | |
| offset_y = (vector_y * alpha).item() | |
| offset_z = (vector_z * alpha).item() | |
| map_to_2d = {} | |
| map_to_2d[4] = 11 | |
| map_to_2d[7] = 12 | |
| map_to_2d[10] = 13 | |
| map_to_2d[5] = 8 | |
| map_to_2d[8] = 9 | |
| map_to_2d[11] = 10 | |
| map_to_2d[20] = 6 | |
| map_to_2d[22] = 7 | |
| map_to_2d[21] = 3 | |
| map_to_2d[23] = 4 | |
| map_to_2d[18] = 5 | |
| map_to_2d[19] = 2 | |
| for affected_body_point in affected_body_points: | |
| if torch.sum(smpl_joints[affected_body_point]) == 0: | |
| continue | |
| new_smpl_joint = smpl_joints[affected_body_point] + torch.tensor([offset_x, offset_y, offset_z]).to(smpl_joints.device) | |
| new_smpl_joint_2d_offset = p3d_to_p2d(new_smpl_joint.reshape(1,1,3).cpu().numpy(), self.height, self.width)[0][0] - p3d_to_p2d(smpl_joints[affected_body_point].reshape(1,1,3).cpu().numpy(), self.height, self.width)[0][0] | |
| new_smpl_joint_2d_offset = np.array([new_smpl_joint_2d_offset[0] / self.width, new_smpl_joint_2d_offset[1] / self.height]) | |
| smpl_joints[affected_body_point] = new_smpl_joint | |
| if affected_body_point in map_to_2d.keys(): | |
| affected_candidate_point_idx = map_to_2d[affected_body_point] | |
| if subset[affected_candidate_point_idx] != -1 and candidate[affected_candidate_point_idx][0] != -1 and candidate[affected_candidate_point_idx][1] != -1: | |
| candidate[affected_candidate_point_idx] = candidate[affected_candidate_point_idx] + new_smpl_joint_2d_offset # 2d的也移动这么多 | |
| if affected_candidate_point_idx == 4: # dwpose 右手 （反的 | |
| left_hand[:] = left_hand + new_smpl_joint_2d_offset | |
| if affected_candidate_point_idx == 7: # dwpose 左手 （反的 | |
| right_hand[:] = right_hand + new_smpl_joint_2d_offset | |
| def aug_init(self, body_type): | |
| print(f"augmentation: using body_type: {body_type}") | |
| self.shoulder_alpha = 0 | |
| self.upper_arm_alpha = 0 | |
| self.forearm_alpha = 0 | |
| self.body_alpha = 0 | |
| self.thigh_alpha = 0 | |
| self.calf_alpha = 0 | |
| if body_type == "normal_human": | |
| self.body_reshape_selected_methods = [] | |
| elif body_type == "dwarf": # body不动 | |
| self.upper_arm_alpha = random.uniform(-0.3, -0.2) | |
| self.forearm_alpha = self.upper_arm_alpha | |
| self.shoulder_alpha = -0.2 | |
| self.thigh_alpha = random.uniform(-0.3, -0.2) | |
| self.calf_alpha = self.thigh_alpha | |
| self.body_reshape_selected_methods = [self.reshape_shoulder, self.reshape_arm, self.reshape_leg] | |
| self.face_alpha = 0.2 | |
| elif body_type == "slender": | |
| self.upper_arm_alpha = 0.3 | |
| self.forearm_alpha = 0.2 | |
| self.shoulder_alpha = 0.2 | |
| self.thigh_alpha = 0.1 | |
| self.calf_alpha = 0.1 | |
| self.body_reshape_selected_methods = [self.reshape_shoulder, self.reshape_arm, self.reshape_leg] | |
| self.face_alpha = -0.2 | |
| elif body_type == "elf": | |
| self.body_alpha = random.uniform(-0.2, 0.2) | |
| self.shoulder_alpha = 0.1 | |
| self.upper_arm_alpha = 0.1 | |
| self.forearm_alpha = 0.1 | |
| self.thigh_alpha = 0.25 | |
| self.calf_alpha = 0.25 | |
| self.body_reshape_selected_methods = [self.reshape_body, self.reshape_shoulder, self.reshape_arm, self.reshape_leg] | |
| self.face_alpha = 0 | |
| elif body_type == "king-kong": | |
| self.body_alpha = 0.1 | |
| self.thigh_alpha = -0.25 | |
| self.calf_alpha = -0.25 | |
| self.upper_arm_alpha = 0.2 | |
| self.forearm_alpha = 0.2 | |
| self.shoulder_alpha = 0.3 | |
| self.body_reshape_selected_methods = [self.reshape_body, self.reshape_shoulder, self.reshape_arm, self.reshape_leg] | |
| self.face_alpha = 0 | |
| elif body_type == "random_long_arm_long_leg": | |
| self.upper_arm_alpha = random.uniform(-0.2, 0.2) | |
| self.forearm_alpha = random.uniform(-0.2, 0.2) | |
| self.thigh_alpha = random.uniform(-0.2, 0.2) | |
| self.calf_alpha = random.uniform(-0.2, 0.2) | |
| self.body_alpha = random.uniform(-0.1, 0.1) | |
| self.body_reshape_selected_methods = [self.reshape_body, self.reshape_arm, self.reshape_leg] | |
| elif body_type == "test_case_1": | |
| self.upper_arm_alpha = -0.4 | |
| self.forearm_alpha = -0.4 | |
| self.shoulder_alpha = -0.3 | |
| self.thigh_alpha = 0.2 | |
| self.calf_alpha = 0.2 | |
| self.body_alpha = 0.1 | |
| self.face_alpha = 0.4 | |
| self.body_reshape_selected_methods = [self.reshape_body, self.reshape_shoulder, self.reshape_arm, self.reshape_leg] | |
| elif body_type == "test_case_2": | |
| self.upper_arm_alpha = 0.4 | |
| self.forearm_alpha = 0.4 | |
| self.shoulder_alpha = 0.3 | |
| self.thigh_alpha = -0.2 | |
| self.calf_alpha = -0.25 | |
| self.body_alpha = -0.2 | |
| self.face_alpha = -0.2 | |
| self.body_reshape_selected_methods = [self.reshape_body, self.reshape_shoulder, self.reshape_arm, self.reshape_leg] | |
| elif body_type == "test_case_3": | |
| self.upper_arm_alpha = 0.2 | |
| self.forearm_alpha = 0.2 | |
| self.shoulder_alpha = 0.2 | |
| self.thigh_alpha = 0.4 | |
| self.calf_alpha = 0.4 | |
| self.body_alpha = 0.3 | |
| self.face_alpha = 0.15 | |
| self.body_reshape_selected_methods = [self.reshape_body, self.reshape_shoulder, self.reshape_arm, self.reshape_leg] | |
| elif body_type == "normal_human_test": | |
| self.body_reshape_selected_methods = [] | |
| self.face_alpha = 0 | |
| def apply_random_reshapes(self, smpl_joints_list, candidate, left_hand, right_hand, face, subset): | |
| # Apply the two selected reshape methods | |
| for method in self.body_reshape_selected_methods: | |
| method(smpl_joints_list, candidate, subset, left_hand, right_hand, face) | |
| for method in self.body_offset_selected_methods: | |
| method(smpl_joints_list, candidate, left_hand, right_hand, face) | |
| for method in self.face_reshape_methods: | |
| method(candidate, face, subset) | |
| def reshape_body(self, smpl_joints_list, candidate, subset, left_hand, right_hand, face): | |
| self.pose_reshape_3d(self.body_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 12, 1, | |
| [1, 4, 7, 10] | |
| ) | |
| self.pose_reshape_3d(self.body_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 12, 2, | |
| [2, 5, 8, 11] | |
| ) | |
| def reshape_arm(self, smpl_joints_list, candidate, subset, left_hand, right_hand, face): | |
| self.pose_reshape_3d(self.upper_arm_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 16, 18, | |
| [18, 20, 22] | |
| ) | |
| self.pose_reshape_3d(self.upper_arm_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 17, 19, | |
| [19, 21, 23] | |
| ) | |
| self.pose_reshape_3d(self.forearm_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 18, 20, | |
| [20, 22] | |
| ) | |
| self.pose_reshape_3d(self.forearm_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 19, 21, | |
| [21, 23] | |
| ) | |
| def reshape_leg(self, smpl_joints_list, candidate, subset, left_hand, right_hand, face): | |
| self.pose_reshape_3d(self.thigh_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 1, 4, | |
| [4, 7, 10] | |
| ) | |
| self.pose_reshape_3d(self.thigh_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 2, 5, | |
| [5, 8, 11] | |
| ) | |
| self.pose_reshape_3d(self.calf_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 4, 7, | |
| [7, 10] | |
| ) | |
| self.pose_reshape_3d(self.calf_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 5, 8, | |
| [8, 11] | |
| ) | |
| def reshape_shoulder(self, smpl_joints_list, candidate, subset, left_hand, right_hand, face): | |
| self.pose_reshape_3d(self.shoulder_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 12, 16, | |
| [16, 18, 20, 22] | |
| ) | |
| self.pose_reshape_3d(self.shoulder_alpha, smpl_joints_list, candidate, subset, left_hand, right_hand, face, | |
| 12, 17, | |
| [17, 19, 21, 23] | |
| ) | |
| # def offset_3d_all(self, smpl_joints_list, candidate, left_hand, right_hand, face): | |
| # smpl_joints_list = smpl_joints_list + torch.tensor([self.offset_3d_x, self.offset_3d_y, self.offset_3d_z]).to(smpl_joints_list.device) | |
| # offset_2d = p3d_to_p2d(np.array([[[self.offset_3d_x, self.offset_3d_y, self.offset_3d_z]]]), self.height, self.width)[0][0] | |
| # candidate = candidate + np.array([offset_2d[0], offset_2d[1]]) | |
| # left_hand = left_hand + np.array([offset_2d[0], offset_2d[1]]) | |
| # right_hand = right_hand + np.array([offset_2d[0], offset_2d[1]]) | |
| # face = face + np.array([offset_2d[0], offset_2d[1]]) | |
| def reshape_face(self, candidate, face, subset): | |
| self.pose_reshape_2d_for_face(alpha=self.face_alpha, candidate=candidate, face=face, subset=subset, | |
| body_anchor_point=0, body_affected_points=[14, 15, 16, 17]) | |