def detect_keypoints2d_PAF()

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)