|
|
|
|
| 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))
|
|
|
