|
|
|
|
| import numpy as np
|
| import copy
|
| from ..retarget_pose import get_retarget_pose
|
|
|
| L_EYE_IDXS = list(range(36, 42))
|
| R_EYE_IDXS = list(range(42, 48))
|
| NOSE_TIP = 30
|
| MOUTH_L = 48
|
| MOUTH_R = 54
|
| JAW_LINE = list(range(0, 17))
|
|
|
|
|
|
|
| def aaposemeta_to_dwpose(meta):
|
| candidate_body = meta['keypoints_body'][:-2][:, :2]
|
| score_body = meta['keypoints_body'][:-2][:, 2]
|
| subset_body = np.arange(len(candidate_body), dtype=float)
|
| subset_body[score_body <= 0] = -1
|
| bodies = {
|
| "candidate": candidate_body,
|
| "subset": np.expand_dims(subset_body, axis=0),
|
| "score": np.expand_dims(score_body, axis=0)
|
| }
|
| hands_coords = np.stack([
|
| meta['keypoints_right_hand'][:, :2],
|
| meta['keypoints_left_hand'][:, :2]
|
| ])
|
| hands_score = np.stack([
|
| meta['keypoints_right_hand'][:, 2],
|
| meta['keypoints_left_hand'][:, 2]
|
| ])
|
| faces_coords = np.expand_dims(meta['keypoints_face'][1:][:, :2], axis=0)
|
| faces_score = np.expand_dims(meta['keypoints_face'][1:][:, 2], axis=0)
|
| dwpose_format = {
|
| "bodies": bodies,
|
| "hands": hands_coords,
|
| "hands_score": hands_score,
|
| "faces": faces_coords,
|
| "faces_score": faces_score
|
| }
|
| return dwpose_format
|
|
|
| def aaposemeta_obj_to_dwpose(pose_meta):
|
| """
|
| Convert an AAPoseMeta object into a dwpose-like data structure
|
| Restore coordinates to relative coordinates (divide by width, height)
|
| Only handle None -> fill with zeros
|
| """
|
| w = pose_meta.width
|
| h = pose_meta.height
|
|
|
|
|
| def safe(arr, like_shape):
|
| if arr is None:
|
| return np.zeros(like_shape, dtype=np.float32)
|
| arr_np = np.array(arr, dtype=np.float32)
|
| arr_np = np.nan_to_num(arr_np, nan=0.0)
|
| return arr_np
|
|
|
| kps_body = safe(pose_meta.kps_body, (pose_meta.kps_body_p.shape[0], 2))
|
| candidate_body = kps_body / np.array([w, h])
|
| score_body = safe(pose_meta.kps_body_p, (candidate_body.shape[0],))
|
| subset_body = np.arange(len(candidate_body), dtype=float)
|
| subset_body[score_body <= 0] = -1
|
| bodies = {
|
| "candidate": candidate_body,
|
| "subset": np.expand_dims(subset_body, axis=0),
|
| "score": np.expand_dims(score_body, axis=0)
|
| }
|
|
|
|
|
| kps_rhand = safe(pose_meta.kps_rhand, (pose_meta.kps_rhand_p.shape[0], 2))
|
| kps_lhand = safe(pose_meta.kps_lhand, (pose_meta.kps_lhand_p.shape[0], 2))
|
| hands_coords = np.stack([
|
| kps_rhand / np.array([w, h]),
|
| kps_lhand / np.array([w, h])
|
| ])
|
| hands_score = np.stack([
|
| safe(pose_meta.kps_rhand_p, (kps_rhand.shape[0],)),
|
| safe(pose_meta.kps_lhand_p, (kps_lhand.shape[0],))
|
| ])
|
|
|
| dwpose_format = {
|
| "bodies": bodies,
|
| "hands": hands_coords,
|
| "hands_score": hands_score,
|
| "faces": None,
|
| "faces_score": None
|
| }
|
| return dwpose_format
|
|
|
|
|
|
|
| def _to_68x2(arr):
|
| if arr.shape == (1, 68, 2):
|
| def to_orig(x):
|
| x = np.asarray(x, dtype=np.float64)
|
| if x.shape != (68, 2):
|
| raise ValueError("to_orig expects (68,2)")
|
| return x[np.newaxis, :, :]
|
| return arr[0].astype(np.float64), to_orig
|
| if arr.shape == (68, 2):
|
| def to_orig(x):
|
| x = np.asarray(x, dtype=np.float64)
|
| if x.shape != (68, 2):
|
| raise ValueError("to_orig expects (68,2)")
|
| return x
|
| return arr.astype(np.float64), to_orig
|
| if arr.shape == (2, 68):
|
| def to_orig(x):
|
| x = np.asarray(x, dtype=np.float64)
|
| if x.shape != (68, 2):
|
| raise ValueError("to_orig expects (68,2)")
|
| return x.T
|
| return arr.T.astype(np.float64), to_orig
|
| raise ValueError(f"faces shape {arr.shape} not supported; expected (1,68,2) or (68,2) or (2,68)")
|
|
|
| def _eye_center(face68, idxs):
|
| return face68[idxs].mean(axis=0)
|
|
|
| def _anchors(face68):
|
| le = _eye_center(face68, L_EYE_IDXS)
|
| re = _eye_center(face68, R_EYE_IDXS)
|
| nose = face68[NOSE_TIP]
|
| lm = face68[MOUTH_L]
|
| rm = face68[MOUTH_R]
|
| if re[0] < le[0]:
|
| le, re = re, le
|
| return np.stack([le, re, nose, lm, rm], axis=0)
|
|
|
| def _face_scale_only(src68, ref68, target_nose_pos, alpha=1.0, anchor_pairs=[[36, 45], [27, 8]]):
|
| """
|
| Rough alignment - adjust the shape of the source face according to the proportions of the reference, and align the nose tip to target_nose_pos.
|
| anchor_pairs:
|
| - [36, 45] for x
|
| - [27, 8] for y
|
| """
|
| src = np.asarray(src68, dtype=np.float64)
|
| ref = np.asarray(ref68, dtype=np.float64)
|
|
|
| center = _anchors(src).mean(axis=0)
|
| src_centered = src - center
|
|
|
| src_w = np.linalg.norm(src[anchor_pairs[0][0]] - src[anchor_pairs[0][1]])
|
| ref_w = np.linalg.norm(ref[anchor_pairs[0][0]] - ref[anchor_pairs[0][1]])
|
|
|
| src_h = np.linalg.norm(src[anchor_pairs[1][0]] - src[anchor_pairs[1][1]])
|
| ref_h = np.linalg.norm(ref[anchor_pairs[1][0]] - ref[anchor_pairs[1][1]])
|
|
|
| scale_x = ref_w / src_w if src_w > 1e-6 else 1.0
|
| scale_y = ref_h / src_h if src_h > 1e-6 else 1.0
|
|
|
| scaled_local = src_centered.copy()
|
| scaled_local[:, 0] *= (1 - alpha) + scale_x * alpha
|
| scaled_local[:, 1] *= (1 - alpha) + scale_y * alpha
|
| scaled_global = scaled_local + center
|
|
|
| nose_idx = NOSE_TIP
|
| current_nose = scaled_global[nose_idx]
|
| offset = target_nose_pos - current_nose
|
| scaled_global += offset
|
|
|
| return scaled_global
|
|
|
|
|
|
|
|
|
| def scale_and_translate_pose(tgt_pose, ref_pose, conf_th=0.9, return_ratio=False):
|
| aligned_pose = copy.deepcopy(tgt_pose)
|
| th = 1e-6
|
| ref_kpt = ref_pose['bodies']['candidate'].astype(np.float32)
|
| tgt_kpt = aligned_pose['bodies']['candidate'].astype(np.float32)
|
|
|
| ref_sc = ref_pose['bodies'].get('score', np.ones(ref_kpt.shape[0])).astype(np.float32).reshape(-1)
|
| tgt_sc = tgt_pose['bodies'].get('score', np.ones(tgt_kpt.shape[0])).astype(np.float32).reshape(-1)
|
|
|
| ref_shoulder_valid = (ref_sc[2] >= conf_th) and (ref_sc[5] >= conf_th)
|
| tgt_shoulder_valid = (tgt_sc[2] >= conf_th) and (tgt_sc[5] >= conf_th)
|
| shoulder_ok = ref_shoulder_valid and tgt_shoulder_valid
|
|
|
| ref_hip_valid = (ref_sc[8] >= conf_th) and (ref_sc[11] >= conf_th)
|
| tgt_hip_valid = (tgt_sc[8] >= conf_th) and (tgt_sc[11] >= conf_th)
|
| hip_ok = ref_hip_valid and tgt_hip_valid
|
|
|
| if shoulder_ok and hip_ok:
|
| ref_shoulder_w = abs(ref_kpt[5, 0] - ref_kpt[2, 0])
|
| tgt_shoulder_w = abs(tgt_kpt[5, 0] - tgt_kpt[2, 0])
|
| x_ratio = ref_shoulder_w / tgt_shoulder_w if tgt_shoulder_w > th else 1.0
|
|
|
| ref_torso_h = abs(np.mean(ref_kpt[[8, 11], 1]) - np.mean(ref_kpt[[2, 5], 1]))
|
| tgt_torso_h = abs(np.mean(tgt_kpt[[8, 11], 1]) - np.mean(tgt_kpt[[2, 5], 1]))
|
| y_ratio = ref_torso_h / tgt_torso_h if tgt_torso_h > th else 1.0
|
| scale_ratio = (x_ratio + y_ratio) / 2
|
|
|
| elif shoulder_ok:
|
| ref_sh_dist = np.linalg.norm(ref_kpt[2] - ref_kpt[5])
|
| tgt_sh_dist = np.linalg.norm(tgt_kpt[2] - tgt_kpt[5])
|
| scale_ratio = ref_sh_dist / tgt_sh_dist if tgt_sh_dist > th else 1.0
|
|
|
| else:
|
| ref_ear_dist = np.linalg.norm(ref_kpt[16] - ref_kpt[17])
|
| tgt_ear_dist = np.linalg.norm(tgt_kpt[16] - tgt_kpt[17])
|
| scale_ratio = ref_ear_dist / tgt_ear_dist if tgt_ear_dist > th else 1.0
|
|
|
| if return_ratio:
|
| return scale_ratio
|
|
|
|
|
| anchor_idx = 1
|
| anchor_pt_before_scale = tgt_kpt[anchor_idx].copy()
|
| def scale(arr):
|
| if arr is not None and arr.size > 0:
|
| arr[..., 0] = anchor_pt_before_scale[0] + (arr[..., 0] - anchor_pt_before_scale[0]) * scale_ratio
|
| arr[..., 1] = anchor_pt_before_scale[1] + (arr[..., 1] - anchor_pt_before_scale[1]) * scale_ratio
|
| scale(tgt_kpt)
|
| scale(aligned_pose.get('faces'))
|
| scale(aligned_pose.get('hands'))
|
|
|
|
|
| offset = ref_kpt[anchor_idx] - tgt_kpt[anchor_idx]
|
| def translate(arr):
|
| if arr is not None and arr.size > 0:
|
| arr += offset
|
| translate(tgt_kpt)
|
| translate(aligned_pose.get('faces'))
|
| translate(aligned_pose.get('hands'))
|
| aligned_pose['bodies']['candidate'] = tgt_kpt
|
|
|
| return aligned_pose, shoulder_ok, hip_ok
|
|
|
|
|
|
|
| def align_to_reference(ref_pose_meta, tpl_pose_metas, tpl_dwposes, anchor_idx=None):
|
|
|
|
|
| ref_pose_dw = aaposemeta_to_dwpose(ref_pose_meta)
|
| best_idx = anchor_idx
|
| tpl_pose_meta_best = tpl_pose_metas[best_idx]
|
|
|
| tpl_retarget_pose_metas = get_retarget_pose(
|
| tpl_pose_meta_best,
|
| ref_pose_meta,
|
| tpl_pose_metas,
|
| None, None
|
| )
|
|
|
| retarget_dwposes = [aaposemeta_obj_to_dwpose(pm) for pm in tpl_retarget_pose_metas]
|
|
|
| if ref_pose_dw['faces'] is not None:
|
| ref68, _ = _to_68x2(ref_pose_dw['faces'])
|
| for frame_idx, (tpl_dw, rt_dw) in enumerate(zip(tpl_dwposes, retarget_dwposes)):
|
| if tpl_dw['faces'] is None:
|
| continue
|
| src68, to_orig = _to_68x2(tpl_dw['faces'])
|
| target_nose_pos = rt_dw['bodies']['candidate'][0]
|
| scaled68 = _face_scale_only(src68, ref68, target_nose_pos, alpha=1.0)
|
| rt_dw['faces'] = to_orig(scaled68)
|
| rt_dw['faces_score'] = tpl_dw['faces_score']
|
|
|
| return retarget_dwposes
|
|
|
|
|
|
|
|
|
| def compute_ratios_stepwise(ref_scores, source_scores, ref_pts, src_pts, conf_th=0.9, th=1e-6):
|
|
|
| def keypoint_valid(idx):
|
| return ref_scores[0, idx] >= conf_th and source_scores[0, idx] >= conf_th
|
|
|
| def safe_ratio(p1, p2):
|
| len_ref = np.linalg.norm(ref_pts[p1] - ref_pts[p2])
|
| len_src = np.linalg.norm(src_pts[p1] - src_pts[p2])
|
| if len_src > th:
|
| return len_ref / len_src
|
| else:
|
| return 1.0
|
|
|
| ratio_pairs = [
|
| (0,1),(1,2),(1,5),(2,3),(3,4),(5,6),(6,7),
|
| (0,14),(0,15),(14,16),(15,17),
|
| (8,9),(9,10),(11,12),(12,13),
|
| (1,8),(1,11)
|
| ]
|
| ratios = {p: 1.0 for p in ratio_pairs}
|
|
|
| parent_map = {
|
| (3, 4): (2, 3),
|
| (6, 7): (5, 6),
|
| (9, 10): (8, 9),
|
| (12, 13): (11, 12)
|
| }
|
|
|
|
|
| if all(keypoint_valid(i) for i in [0,1,14,15,16,17]):
|
| ratios[(0,1)] = safe_ratio(0,1)
|
| ratios[(0,14)] = safe_ratio(0,14)
|
| ratios[(0,15)] = safe_ratio(0,15)
|
| ratios[(14,16)]= safe_ratio(14,16)
|
| ratios[(15,17)]= safe_ratio(15,17)
|
|
|
|
|
| if all(keypoint_valid(i) for i in [0,1,2,5,14,15,16,17]):
|
| ratios[(1,2)] = safe_ratio(1,2)
|
| ratios[(1,5)] = safe_ratio(1,5)
|
|
|
|
|
| if all(keypoint_valid(i) for i in [0,1,2,5,14,15,16,17,3,6]):
|
| ratios[(2,3)] = safe_ratio(2,3)
|
| ratios[(5,6)] = safe_ratio(5,6)
|
| ratios[(3,4)] = ratios[parent_map[(3,4)]]
|
| ratios[(6,7)] = ratios[parent_map[(6,7)]]
|
|
|
|
|
| if all(keypoint_valid(i) for i in [0,1,2,5,14,15,16,17,3,6,8,11]):
|
| ratios[(1,8)] = safe_ratio(1,8)
|
| ratios[(1,11)] = safe_ratio(1,11)
|
|
|
|
|
| if all(keypoint_valid(i) for i in [0,1,2,5,14,15,16,17,3,6,8,11,4,7]):
|
| ratios[(3,4)] = safe_ratio(3,4)
|
| ratios[(6,7)] = safe_ratio(6,7)
|
|
|
|
|
| if all(keypoint_valid(i) for i in [0,1,2,5,14,15,16,17,3,6,8,11,4,7,9,12]):
|
| ratios[(8,9)] = safe_ratio(8,9)
|
| ratios[(11,12)] = safe_ratio(11,12)
|
| ratios[(9,10)] = ratios[parent_map[(9,10)]]
|
| ratios[(12,13)]= ratios[parent_map[(12,13)]]
|
|
|
|
|
| if all(keypoint_valid(i) for i in range(18)):
|
| for p in ratio_pairs:
|
| ratios[p] = safe_ratio(*p)
|
|
|
| symmetric_pairs = [
|
| ((1, 2), (1, 5)),
|
| ((2, 3), (5, 6)),
|
| ((3, 4), (6, 7)),
|
| ((8, 9), (11, 12)),
|
| ((9, 10), (12, 13))
|
| ]
|
| for left_key, right_key in symmetric_pairs:
|
| left_val = ratios.get(left_key)
|
| right_val = ratios.get(right_key)
|
| if left_val is not None and right_val is not None:
|
| avg_val = (left_val + right_val) / 2.0
|
| ratios[left_key] = avg_val
|
| ratios[right_key] = avg_val
|
|
|
| eye_pairs = [
|
| ((13, 15), (14, 16))
|
| ]
|
| for left_key, right_key in eye_pairs:
|
| left_val = ratios.get(left_key)
|
| right_val = ratios.get(right_key)
|
| if left_val is not None and right_val is not None:
|
| avg_val = (left_val + right_val) / 2.0
|
| ratios[left_key] = avg_val
|
| ratios[right_key] = avg_val
|
|
|
| return ratios
|
|
|
| def align_to_pose(ref_dwpose, tpl_dwposes,anchor_idx=None,conf_th=0.9,):
|
| detected_poses = copy.deepcopy(tpl_dwposes)
|
|
|
| best_pose = tpl_dwposes[anchor_idx]
|
| ref_pose_scaled, _, _ = scale_and_translate_pose(ref_dwpose, best_pose, conf_th=conf_th)
|
|
|
| ref_candidate = ref_pose_scaled['bodies']['candidate'].astype(np.float32)
|
| ref_scores = ref_pose_scaled['bodies']['score'].astype(np.float32)
|
|
|
| source_candidate = best_pose['bodies']['candidate'].astype(np.float32)
|
| source_scores = best_pose['bodies']['score'].astype(np.float32)
|
|
|
| has_ref_face = 'faces' in ref_pose_scaled and ref_pose_scaled['faces'] is not None and ref_pose_scaled['faces'].size > 0
|
| if has_ref_face:
|
| try:
|
| ref68, _ = _to_68x2(ref_pose_scaled['faces'])
|
| except Exception as e:
|
| print("Reference face conversion failed:", e)
|
| has_ref_face = False
|
|
|
| ratios = compute_ratios_stepwise(ref_scores, source_scores, ref_candidate, source_candidate, conf_th=conf_th, th=1e-6)
|
|
|
| for pose in detected_poses:
|
| candidate = pose['bodies']['candidate']
|
| hands = pose['hands']
|
|
|
|
|
| ratio = ratios[(0, 1)]
|
| x_offset = (candidate[1][0] - candidate[0][0]) * (1. - ratio)
|
| y_offset = (candidate[1][1] - candidate[0][1]) * (1. - ratio)
|
| candidate[[0, 14, 15, 16, 17], 0] += x_offset
|
| candidate[[0, 14, 15, 16, 17], 1] += y_offset
|
|
|
|
|
| ratio = ratios[(1, 2)]
|
| x_offset = (candidate[1][0] - candidate[2][0]) * (1. - ratio)
|
| y_offset = (candidate[1][1] - candidate[2][1]) * (1. - ratio)
|
| candidate[[2, 3, 4], 0] += x_offset
|
| candidate[[2, 3, 4], 1] += y_offset
|
| hands[1, :, 0] += x_offset
|
| hands[1, :, 1] += y_offset
|
|
|
|
|
| ratio = ratios[(1, 5)]
|
| x_offset = (candidate[1][0] - candidate[5][0]) * (1. - ratio)
|
| y_offset = (candidate[1][1] - candidate[5][1]) * (1. - ratio)
|
| candidate[[5, 6, 7], 0] += x_offset
|
| candidate[[5, 6, 7], 1] += y_offset
|
| hands[0, :, 0] += x_offset
|
| hands[0, :, 1] += y_offset
|
|
|
|
|
| ratio = ratios[(2, 3)]
|
| x_offset = (candidate[2][0] - candidate[3][0]) * (1. - ratio)
|
| y_offset = (candidate[2][1] - candidate[3][1]) * (1. - ratio)
|
| candidate[[3, 4], 0] += x_offset
|
| candidate[[3, 4], 1] += y_offset
|
| hands[1, :, 0] += x_offset
|
| hands[1, :, 1] += y_offset
|
|
|
|
|
| ratio = ratios[(3, 4)]
|
| x_offset = (candidate[3][0] - candidate[4][0]) * (1. - ratio)
|
| y_offset = (candidate[3][1] - candidate[4][1]) * (1. - ratio)
|
| candidate[4, 0] += x_offset
|
| candidate[4, 1] += y_offset
|
| hands[1, :, 0] += x_offset
|
| hands[1, :, 1] += y_offset
|
|
|
|
|
| ratio = ratios[(5, 6)]
|
| x_offset = (candidate[5][0] - candidate[6][0]) * (1. - ratio)
|
| y_offset = (candidate[5][1] - candidate[6][1]) * (1. - ratio)
|
| candidate[[6, 7], 0] += x_offset
|
| candidate[[6, 7], 1] += y_offset
|
| hands[0, :, 0] += x_offset
|
| hands[0, :, 1] += y_offset
|
|
|
|
|
| ratio = ratios[(6, 7)]
|
| x_offset = (candidate[6][0] - candidate[7][0]) * (1. - ratio)
|
| y_offset = (candidate[6][1] - candidate[7][1]) * (1. - ratio)
|
| candidate[7, 0] += x_offset
|
| candidate[7, 1] += y_offset
|
| hands[0, :, 0] += x_offset
|
| hands[0, :, 1] += y_offset
|
|
|
|
|
| for (p1, p2) in [(0,14),(0,15),(14,16),(15,17)]:
|
| ratio = ratios[(p1,p2)]
|
| x_offset = (candidate[p1][0] - candidate[p2][0]) * (1. - ratio)
|
| y_offset = (candidate[p1][1] - candidate[p2][1]) * (1. - ratio)
|
| candidate[p2, 0] += x_offset
|
| candidate[p2, 1] += y_offset
|
|
|
|
|
| ratio = ratios[(1, 8)]
|
| x_offset = (candidate[1][0] - candidate[8][0]) * (1. - ratio)
|
| y_offset = (candidate[1][1] - candidate[8][1]) * (1. - ratio)
|
| candidate[8, 0] += x_offset
|
| candidate[8, 1] += y_offset
|
|
|
| ratio = ratios[(1, 11)]
|
| x_offset = (candidate[1][0] - candidate[11][0]) * (1. - ratio)
|
| y_offset = (candidate[1][1] - candidate[11][1]) * (1. - ratio)
|
| candidate[11, 0] += x_offset
|
| candidate[11, 1] += y_offset
|
|
|
|
|
| ratio = ratios[(8, 9)]
|
| x_offset = (candidate[9][0] - candidate[8][0]) * (ratio - 1.)
|
| y_offset = (candidate[9][1] - candidate[8][1]) * (ratio - 1.)
|
| candidate[[9, 10], 0] += x_offset
|
| candidate[[9, 10], 1] += y_offset
|
|
|
| ratio = ratios[(9, 10)]
|
| x_offset = (candidate[10][0] - candidate[9][0]) * (ratio - 1.)
|
| y_offset = (candidate[10][1] - candidate[9][1]) * (ratio - 1.)
|
| candidate[10, 0] += x_offset
|
| candidate[10, 1] += y_offset
|
|
|
| ratio = ratios[(11, 12)]
|
| x_offset = (candidate[12][0] - candidate[11][0]) * (ratio - 1.)
|
| y_offset = (candidate[12][1] - candidate[11][1]) * (ratio - 1.)
|
| candidate[[12, 13], 0] += x_offset
|
| candidate[[12, 13], 1] += y_offset
|
|
|
| ratio = ratios[(12, 13)]
|
| x_offset = (candidate[13][0] - candidate[12][0]) * (ratio - 1.)
|
| y_offset = (candidate[13][1] - candidate[12][1]) * (ratio - 1.)
|
| candidate[13, 0] += x_offset
|
| candidate[13, 1] += y_offset
|
|
|
|
|
| if has_ref_face and 'faces' in pose and pose['faces'] is not None and pose['faces'].size > 0:
|
| try:
|
| src68, to_orig = _to_68x2(pose['faces'])
|
| scaled68 = _face_scale_only(src68, ref68, candidate[0], alpha=1.0)
|
| pose['faces'] = to_orig(scaled68)
|
| except Exception as e:
|
| print("Reference face conversion failed:", e)
|
| continue
|
|
|
| return detected_poses
|
|
|
