in visualization/POF/utils/PAF.py [0:0]
def detect_keypoints2d_PAF(scoremaps, PAF, objtype=0, weight_conn=1.0, mean_shift=False, prev_frame=None):
print('PAF_type {}'.format(PAF_type))
if len(scoremaps.shape) == 4:
scoremaps = np.squeeze(scoremaps)
s = scoremaps.shape
assert len(s) == 3, "This function was only designed for 3D Scoremaps."
assert (s[2] < s[1]) and (s[2] < s[0]), "Probably the input is not correct, because [H, W, C] is expected."
num_candidate = 5
local_maxs = []
for i in range(s[2]):
candidates = skimage.feature.peak_local_max(scoremaps[:, :, i], num_peaks=num_candidate)
if candidates.shape[0] < num_candidate:
# if less than that, replicate the first element
if candidates.shape[0] > 0:
candidates = np.concatenate([candidates[0][np.newaxis, :]] * (num_candidate - candidates.shape[0]) + [candidates], axis=0)
else:
candidates = np.zeros((5, 2), dtype=int)
local_maxs.append(candidates)
if objtype == 0:
root_idx = 1 # starting constructing the tree from root_idx
else:
assert objtype == 1
root_idx = 0
joint_idx_list = [root_idx]
candidate_idx_list = [[c] for c in range(num_candidate)]
sum_score_list = [scoremaps[local_maxs[root_idx][c, 0], local_maxs[root_idx][c, 1], root_idx] for c in range(num_candidate)]
if prev_frame is not None:
for c in range(num_candidate):
sum_score_list[c] -= 20 * nl.norm(local_maxs[root_idx][candidate_idx_list[c][0]][::-1] - prev_frame[c]) / (s[0] + s[1])
# dynamic programming
for iconn, conn in enumerate(PAFConnection[objtype]):
if objtype == 0:
if PAF_type == 0:
if iconn in (9, 13) or iconn >= 21: # unused PAF connection
continue
elif PAF_type == 1:
if iconn in (9, 13) or iconn >= 20:
continue
joint_idx_list.append(conn[1])
candidates = local_maxs[conn[1]]
new_candidate_idx_list = []
new_sum_score_list = []
for ican, candidate in enumerate(candidates):
best_sum_score = -np.inf
best_candidate_idx = None
B = candidate[::-1]
for candidate_idx, sum_score in zip(candidate_idx_list, sum_score_list):
parent_idx = conn[0]
parent_candidate_idx = candidate_idx[joint_idx_list.index(parent_idx)]
A = local_maxs[parent_idx][parent_candidate_idx][::-1]
connection_score = connection_score_2d(A, B, PAF[:, :, 3 * iconn:3 * iconn + 3])
new_sum_score = sum_score + scoremaps[candidate[0], candidate[1], conn[1]] + weight_conn * connection_score # TODO
if prev_frame is not None:
new_sum_score -= 20 * nl.norm(prev_frame[conn[1]] - B) / (s[0] + s[1])
if new_sum_score > best_sum_score:
best_sum_score = new_sum_score
best_candidate_idx = candidate_idx
assert best_candidate_idx is not None
new_sum_score_list.append(best_sum_score)
new_candidate_idx_list.append(best_candidate_idx + [ican])
sum_score_list = new_sum_score_list
candidate_idx_list = new_candidate_idx_list
best_candidate_idx = candidate_idx_list[np.argmax(sum_score_list)]
best_candidate_idx_joint_order = np.zeros_like(best_candidate_idx)
best_candidate_idx_joint_order[np.array(joint_idx_list, dtype=int)] = best_candidate_idx
best_candidate = np.array([local_maxs[i][j] for i, j in enumerate(best_candidate_idx_joint_order)])
coord2d = best_candidate[:, ::-1]
if objtype == 0:
assert coord2d.shape[0] == 19 or coord2d.shape[0] == 20
if objtype == 1:
assert coord2d.shape[0] == 21
scores = []
for i in range(coord2d.shape[0]):
scores.append(scoremaps[coord2d[i, 1], coord2d[i, 0], i])
if mean_shift:
dWidth = 3
dHeight = 3
new_coord2d = []
for i in range(coord2d.shape[0]):
x1 = max(coord2d[i, 0] - dWidth, 0)
x2 = min(coord2d[i, 0] + dWidth + 1, s[1])
y1 = max(coord2d[i, 1] - dHeight, 0)
y2 = min(coord2d[i, 1] + dHeight + 1, s[0])
Xmap = np.arange(x1, x2)
Ymap = np.arange(y1, y2)
local_scoremap = scoremaps[y1:y2, x1:x2, i]
gt0 = (local_scoremap > 0)
if gt0.any():
pos_scoremap = gt0 * local_scoremap
xAcc = np.sum(pos_scoremap * Xmap)
yAcc = np.sum(np.transpose(pos_scoremap) * Ymap)
scoreAcc = np.sum(pos_scoremap)
new_coord2d.append([xAcc / scoreAcc, yAcc / scoreAcc])
else:
new_coord2d.append([coord2d[i, 0], coord2d[i, 1]])
coord2d = np.array(new_coord2d, dtype=np.float32)
return coord2d.astype(np.float32), np.array(scores, dtype=np.float32)