# ----------------------------------------------------- # Copyright (c) Shanghai Jiao Tong University. All rights reserved. # Written by Jiefeng Li (jeff.lee.sjtu@gmail.com) # ----------------------------------------------------- import numpy as np import torch from scipy.optimize import linear_sum_assignment sigmas = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89]) def candidate_reselect(bboxes, bboxes_scores, pose_preds): ''' Grouping ''' # Group same keypointns together kp_groups = grouping(bboxes, bboxes_scores, pose_preds) ''' Re-select ''' # Generate Matrix human_num = len(pose_preds.keys()) costMatrix = [] for k in range(17): kp_group = kp_groups[k] joint_num = len(kp_group.keys()) costMatrix.append(np.zeros((human_num, joint_num))) group_size = {k: {} for k in range(17)} for n, person in pose_preds.items(): h_id = n assert 0 <= h_id < human_num for k in range(17): g_id = person['group_id'][k] if g_id is not None: if g_id not in group_size[k].keys(): group_size[k][g_id] = 0 group_size[k][g_id] += 1 g_id = int(g_id) - 1 _, _, score = person[k][0] h_score = person['human_score'] if score < 0.05: costMatrix[k][h_id][g_id] = 0 else: costMatrix[k][h_id][g_id] = -(h_score * score) pose_preds = matching(pose_preds, costMatrix, kp_groups) # To JSON final_result = [] for n, person in pose_preds.items(): final_pose = torch.zeros(17, 2) final_score = torch.zeros(17, 1) max_score = 0 mean_score = 0 xmax, xmin = 0, 1e5 ymax, ymin = 0, 1e5 for k in range(17): assert len(person[k]) > 0 x, y, s = person[k][0] xmax = max(xmax, x) xmin = min(xmin, x) ymax = max(ymax, y) ymin = min(ymin, y) final_pose[k][0] = x.item() - 0.3 final_pose[k][1] = y.item() - 0.3 final_score[k] = s.item() mean_score += (s.item() / 17) max_score = max(max_score, s.item()) if torch.max(final_score).item() < 0.1: continue if (1.5 ** 2 * (xmax - xmin) * (ymax - ymin) < 40 * 40): continue final_result.append({ 'keypoints': final_pose, 'kp_score': final_score, 'proposal_score': mean_score + max_score + person['bbox_score'] }) return final_result def grouping(bboxes, bboxes_scores, pose_preds): kp_groups = {} for k in range(17): kp_groups[k] = {} ids = np.zeros(17) for n, person in pose_preds.items(): pose_preds[n]['bbox'] = bboxes[n] pose_preds[n]['bbox_score'] = bboxes_scores[n] pose_preds[n]['group_id'] = {} s = 0 for k in range(17): pose_preds[n]['group_id'][k] = None pose_preds[n][k] = np.array(pose_preds[n][k]) assert len(pose_preds[n][k]) > 0 s += pose_preds[n][k][0][-1] s = s / 17 pose_preds[n]['human_score'] = s for k in range(17): latest_id = ids[k] kp_group = kp_groups[k] assert len(person[k]) > 0 x0, y0, s0 = person[k][0] if s0 < 0.05: continue for g_id, g in kp_group.items(): x_c, y_c = kp_group[g_id]['group_center'] ''' Get Average Box Size ''' group_area = kp_group[g_id]['group_area'] group_area = group_area[0] * group_area[1] / (group_area[2] ** 2) ''' Groupingn Criterion ''' # Joint Group dist = np.sqrt( ((x_c - x0) ** 2 + (y_c - y0) ** 2) / group_area) if dist <= 0.1 * sigmas[k]: # Small Distance if s0 >= 0.3: kp_group[g_id]['kp_list'][0] += x0 * s0 kp_group[g_id]['kp_list'][1] += y0 * s0 kp_group[g_id]['kp_list'][2] += s0 kp_group[g_id]['group_area'][0] += (person['bbox'][2] - person['bbox'][0]) * person['human_score'] kp_group[g_id]['group_area'][1] += (person['bbox'][3] - person['bbox'][1]) * person['human_score'] kp_group[g_id]['group_area'][2] += person['human_score'] x_c = kp_group[g_id]['kp_list'][0] / kp_group[g_id]['kp_list'][2] y_c = kp_group[g_id]['kp_list'][1] / kp_group[g_id]['kp_list'][2] kp_group[g_id]['group_center'] = (x_c, y_c) pose_preds[n]['group_id'][k] = g_id break else: # A new keypoint group latest_id += 1 kp_group[latest_id] = { 'kp_list': None, 'group_center': person[k][0].copy()[:2], 'group_area': None } x, y, s = person[k][0] kp_group[latest_id]['kp_list'] = np.array((x * s, y * s, s)) # Ref Area ref_width = person['bbox'][2] - person['bbox'][0] ref_height = person['bbox'][3] - person['bbox'][1] ref_score = person['human_score'] kp_group[latest_id]['group_area'] = np.array(( ref_width * ref_score, ref_height * ref_score, ref_score)) pose_preds[n]['group_id'][k] = latest_id ids[k] = latest_id return kp_groups def matching(pose_preds, matrix, kp_groups): index = [] for k in range(17): human_ind, joint_ind = linear_sum_assignment(matrix[k]) # human_ind, joint_ind = greedy_matching(matrix[k]) index.append(list(zip(human_ind, joint_ind))) for n, person in pose_preds.items(): for k in range(17): g_id = person['group_id'][k] if g_id is not None: g_id = int(g_id) - 1 h_id = n x, y, s = pose_preds[n][k][0] if ((h_id, g_id) not in index[k]) and len(pose_preds[n][k]) > 1: pose_preds[n][k] = np.delete(pose_preds[n][k], 0, 0) elif ((h_id, g_id) not in index[k]) and len(person[k]) == 1: x, y, _ = pose_preds[n][k][0] pose_preds[n][k][0] = (x, y, 1e-5) pass elif ((h_id, g_id) in index[k]): x, y = kp_groups[k][g_id + 1]['group_center'] s = pose_preds[n][k][0][2] pose_preds[n][k][0] = (x, y, s) return pose_preds def greedy_matching(matrix): num_human, num_joint = matrix.shape if num_joint <= num_human or True: human_ind = np.argmin(matrix, axis=0) joint_ind = np.arange(num_joint) else: pass return human_ind.tolist(), joint_ind.tolist()