| |
|
|
| import cv2 |
| import numpy as np |
| import math |
| import copy |
|
|
| eps = 0.01 |
|
|
| DROP_FACE_POINTS = {0, 14, 15, 16, 17} |
| DROP_UPPER_POINTS = {0, 14, 15, 16, 17, 2, 1, 5, 3, 6} |
| DROP_LOWER_POINTS = {8, 9, 10, 11, 12, 13} |
|
|
| 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 warp_ref_to_pose(tgt_img, |
| ref_pose: dict, |
| tgt_pose: dict, |
| bg_val=(0, 0, 0), |
| conf_th=0.9, |
| align_center=False): |
|
|
| H, W = tgt_img.shape[:2] |
| img_tgt_pose = draw_pose_aligned(tgt_pose, H, W, without_face=True) |
|
|
| tgt_kpt = tgt_pose['bodies']['candidate'].astype(np.float32) |
| ref_kpt = ref_pose['bodies']['candidate'].astype(np.float32) |
|
|
| scale_ratio = scale_and_translate_pose(tgt_pose, ref_pose, conf_th=conf_th, return_ratio=True) |
|
|
| anchor_idx = 1 |
| x0 = tgt_kpt[anchor_idx][0] * W |
| y0 = tgt_kpt[anchor_idx][1] * H |
|
|
| ref_x = ref_kpt[anchor_idx][0] * W if not align_center else W/2 |
| ref_y = ref_kpt[anchor_idx][1] * H |
|
|
| dx = ref_x - x0 |
| dy = ref_y - y0 |
|
|
| |
| M = np.array([[scale_ratio, 0, (1-scale_ratio)*x0 + dx], |
| [0, scale_ratio, (1-scale_ratio)*y0 + dy]], |
| dtype=np.float32) |
| img_warp = cv2.warpAffine(tgt_img, M, (W, H), |
| flags=cv2.INTER_LINEAR, |
| borderValue=bg_val) |
| img_tgt_pose_warp = cv2.warpAffine(img_tgt_pose, M, (W, H), |
| flags=cv2.INTER_LINEAR, |
| borderValue=bg_val) |
| zeros = np.zeros((H, W), dtype=np.uint8) |
| mask_warp = cv2.warpAffine(zeros, M, (W, H), |
| flags=cv2.INTER_NEAREST, |
| borderValue=255) |
| return img_warp, img_tgt_pose_warp, mask_warp |
|
|
| def hsv_to_rgb(hsv): |
| hsv = np.asarray(hsv, dtype=np.float32) |
| in_shape = hsv.shape |
| hsv = hsv.reshape(-1, 3) |
|
|
| h, s, v = hsv[:, 0], hsv[:, 1], hsv[:, 2] |
|
|
| i = (h * 6.0).astype(int) |
| f = (h * 6.0) - i |
| i = i % 6 |
|
|
| p = v * (1.0 - s) |
| q = v * (1.0 - s * f) |
| t = v * (1.0 - s * (1.0 - f)) |
|
|
| rgb = np.zeros_like(hsv) |
| rgb[i == 0] = np.stack([v[i == 0], t[i == 0], p[i == 0]], axis=1) |
| rgb[i == 1] = np.stack([q[i == 1], v[i == 1], p[i == 1]], axis=1) |
| rgb[i == 2] = np.stack([p[i == 2], v[i == 2], t[i == 2]], axis=1) |
| rgb[i == 3] = np.stack([p[i == 3], q[i == 3], v[i == 3]], axis=1) |
| rgb[i == 4] = np.stack([t[i == 4], p[i == 4], v[i == 4]], axis=1) |
| rgb[i == 5] = np.stack([v[i == 5], p[i == 5], q[i == 5]], axis=1) |
|
|
| gray_mask = s == 0 |
| rgb[gray_mask] = np.stack([v[gray_mask]] * 3, axis=1) |
|
|
| return (rgb.reshape(in_shape) * 255) |
|
|
| def get_stickwidth(W, H, stickwidth=4): |
| if max(W, H) < 512: |
| ratio = 1.0 |
| elif max(W, H) < 1080: |
| ratio = 1.5 |
| elif max(W, H) < 2160: |
| ratio = 2.0 |
| elif max(W, H) < 3240: |
| ratio = 2.5 |
| elif max(W, H) < 4320: |
| ratio = 3.5 |
| elif max(W, H) < 5400: |
| ratio = 4.5 |
| else: |
| ratio = 4.0 |
| return int(stickwidth * ratio) |
|
|
|
|
| def alpha_blend_color(color, alpha): |
| return [int(c * alpha) for c in color] |
|
|
|
|
| def draw_bodypose_aligned(canvas, candidate, subset, score, plan=None): |
| H, W, C = canvas.shape |
| candidate = np.array(candidate) |
| subset = np.array(subset) |
| stickwidth = get_stickwidth(W, H, stickwidth=3) |
|
|
| limbSeq = [ |
| [2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], |
| [2, 9], [9, 10], [10, 11], [2, 12], [12, 13], [13, 14], |
| [2, 1], [1, 15], [15, 17], [1, 16], [16, 18], [3, 17], [6, 18]] |
| colors = [ |
| [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], |
| [85, 255, 0], [0, 255, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255], |
| [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], |
| [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]] |
|
|
| HIDE_JOINTS = set() |
| stretch_limb_idx = None |
| stretch_scale = None |
| if plan: |
| if plan["mode"] == "drop_point": |
| HIDE_JOINTS.add(plan["point_idx"]) |
| elif plan["mode"] == "drop_region": |
| HIDE_JOINTS |= set(plan["points"]) |
| elif plan["mode"] == "stretch_limb": |
| stretch_limb_idx = plan["limb_idx"] |
| stretch_scale = plan["stretch_scale"] |
|
|
| hide_joint = np.zeros_like(subset, dtype=bool) |
|
|
| for i in range(17): |
| for n in range(len(subset)): |
| idx_pair = limbSeq[i] |
|
|
| if any(j in HIDE_JOINTS for j in idx_pair): |
| continue |
|
|
| index = subset[n][np.array(idx_pair) - 1] |
| conf = score[n][np.array(idx_pair) - 1] |
| if -1 in index: |
| continue |
| |
| alpha = max(conf[0] * conf[1], 0) if conf[0]>0 and conf[1]>0 else 0.35 |
| if conf[0] == 0 or conf[1] == 0: |
| alpha = 0 |
|
|
| Y = candidate[index.astype(int), 0] * float(W) |
| X = candidate[index.astype(int), 1] * float(H) |
|
|
| if stretch_limb_idx == i: |
| vec_x = X[1] - X[0] |
| vec_y = Y[1] - Y[0] |
| X[1] = X[0] + vec_x * stretch_scale |
| Y[1] = Y[0] + vec_y * stretch_scale |
| hide_joint[n, idx_pair[1]-1] = True |
|
|
| mX = np.mean(X) |
| mY = np.mean(Y) |
| length = ((X[0]-X[1])**2 + (Y[0]-Y[1])**2) ** 0.5 |
| angle = math.degrees(math.atan2(X[0]-X[1], Y[0]-Y[1])) |
| polygon = cv2.ellipse2Poly((int(mY), int(mX)), |
| (int(length/2), stickwidth), int(angle), 0, 360, 1) |
| cv2.fillConvexPoly(canvas, polygon, alpha_blend_color(colors[i], alpha)) |
|
|
| canvas = (canvas * 0.6).astype(np.uint8) |
|
|
| for i in range(18): |
| if i in HIDE_JOINTS: |
| continue |
| for n in range(len(subset)): |
| if hide_joint[n, i]: |
| continue |
| index = int(subset[n][i]) |
| if index == -1: |
| continue |
| x, y = candidate[index][0:2] |
| conf = score[n][i] |
|
|
| alpha = 0 if conf==-2 else max(conf, 0) |
| x = int(x * W) |
| y = int(y * H) |
| cv2.circle(canvas, (x, y), stickwidth, alpha_blend_color(colors[i], alpha), thickness=-1) |
|
|
| return canvas |
|
|
|
|
| def draw_handpose_aligned(canvas, all_hand_peaks, all_hand_scores, draw_th=0.3): |
| H, W, C = canvas.shape |
| stickwidth = get_stickwidth(W, H, stickwidth=2) |
| line_thickness = get_stickwidth(W, H, stickwidth=2) |
|
|
| edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \ |
| [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] |
|
|
| for peaks, scores in zip(all_hand_peaks, all_hand_scores): |
| for ie, e in enumerate(edges): |
| if scores[e[0]] < draw_th or scores[e[1]] < draw_th: |
| continue |
| x1, y1 = peaks[e[0]] |
| x2, y2 = peaks[e[1]] |
| x1 = int(x1 * W) |
| y1 = int(y1 * H) |
| x2 = int(x2 * W) |
| y2 = int(y2 * H) |
|
|
| score = int(scores[e[0]] * scores[e[1]] * 255) |
| if x1 > eps and y1 > eps and x2 > eps and y2 > eps: |
| color = hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]).flatten() |
| color = tuple(int(c * score / 255) for c in color) |
| cv2.line(canvas, (x1, y1), (x2, y2), color, thickness=line_thickness) |
|
|
| for i, keyponit in enumerate(peaks): |
| if scores[i] < draw_th: |
| continue |
|
|
| x, y = keyponit |
| x = int(x * W) |
| y = int(y * H) |
| score = int(scores[i] * 255) |
| if x > eps and y > eps: |
| cv2.circle(canvas, (x, y), stickwidth, (0, 0, score), thickness=-1) |
| return canvas |
|
|
|
|
| def draw_facepose_aligned(canvas, all_lmks, all_scores, draw_th=0.3,face_change=False): |
| H, W, C = canvas.shape |
| stickwidth = get_stickwidth(W, H, stickwidth=2) |
| SKIP_IDX = set(range(0, 17)) |
| SKIP_IDX |= set(range(27, 36)) |
|
|
| for lmks, scores in zip(all_lmks, all_scores): |
| for idx, (lmk, score) in enumerate(zip(lmks, scores)): |
| |
| if idx in SKIP_IDX: |
| continue |
| if score < draw_th: |
| continue |
| x, y = lmk |
| x = int(x * W) |
| y = int(y * H) |
| conf = int(score * 255) |
| |
| if face_change: |
| conf = int(conf * 0.35) |
|
|
| if x > eps and y > eps: |
| cv2.circle(canvas, (x, y), stickwidth, (conf, conf, conf), thickness=-1) |
| return canvas |
|
|
|
|
| def draw_pose_aligned(pose, H, W, ref_w=2160, without_face=False, pose_plan=None, head_strength="full", face_change=False): |
| bodies = pose['bodies'] |
| faces = pose['faces'] |
| hands = pose['hands'] |
| candidate = bodies['candidate'] |
| subset = bodies['subset'] |
| body_score = bodies['score'].copy() |
| |
| if head_strength == "weak": |
| target_joints = [0, 14, 15, 16, 17] |
| body_score[:, target_joints] = -2 |
| elif head_strength == "none": |
| target_joints = [0, 14, 15, 16, 17] |
| body_score[:, target_joints] = 0 |
|
|
| sz = min(H, W) |
| sr = (ref_w / sz) if sz != ref_w else 1 |
| canvas = np.zeros(shape=(int(H*sr), int(W*sr), 3), dtype=np.uint8) |
|
|
| canvas = draw_bodypose_aligned(canvas, candidate, subset, |
| score=body_score, |
| plan=pose_plan,) |
|
|
| canvas = draw_handpose_aligned(canvas, hands, pose['hands_score']) |
|
|
| if not without_face: |
| canvas = draw_facepose_aligned(canvas, faces, pose['faces_score'],face_change=face_change) |
|
|
| return cv2.resize(canvas, (W, H)) |
|
|
