You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
233 lines
7.1 KiB
233 lines
7.1 KiB
# -----------------------------------------------------
|
|
# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
|
|
# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
|
|
# -----------------------------------------------------
|
|
|
|
from copy import deepcopy
|
|
|
|
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()
|