Exemple #1
0
def extract_heatmap_info(heatmap_avg, param_thre1=0.1):
    all_peaks = []
    peak_counter = 0

    for part in range(18):
        map_ori = heatmap_avg[:, :, part]
        map_gau = gaussian_filter(map_ori, sigma=3)

        map_left = np.zeros(map_gau.shape)
        map_left[1:, :] = map_gau[:-1, :]
        map_right = np.zeros(map_gau.shape)
        map_right[:-1, :] = map_gau[1:, :]
        map_up = np.zeros(map_gau.shape)
        map_up[:, 1:] = map_gau[:, :-1]
        map_down = np.zeros(map_gau.shape)
        map_down[:, :-1] = map_gau[:, 1:]

        peaks_binary = np.logical_and.reduce(
            (map_gau >= map_left, map_gau >= map_right, map_gau >= map_up,
             map_gau >= map_down, map_gau > param_thre1))

        peaks = zip(np.nonzero(peaks_binary)[1],
                    np.nonzero(peaks_binary)[0])  # note reverse
        peaks = list(peaks)
        peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
        ids = range(peak_counter, peak_counter + len(peaks))
        peaks_with_score_and_id = [
            peaks_with_score[i] + (ids[i], ) for i in range(len(ids))
        ]

        all_peaks.append(peaks_with_score_and_id)
        peak_counter += len(peaks)

    return all_peaks
def extract_heatmap_info(heatmap_avg, param_thre1=0.1):
    all_peaks = []
    peak_counter = 0

    for part in range(18):
        map_ori = heatmap_avg[:, :, part]
        map_gau = gaussian_filter(map_ori, sigma=3)

        map_left = np.zeros(map_gau.shape)
        map_left[1:, :] = map_gau[:-1, :]
        map_right = np.zeros(map_gau.shape)
        map_right[:-1, :] = map_gau[1:, :]
        map_up = np.zeros(map_gau.shape)
        map_up[:, 1:] = map_gau[:, :-1]
        map_down = np.zeros(map_gau.shape)
        map_down[:, :-1] = map_gau[:, 1:]

        peaks_binary = np.logical_and.reduce(
            (map_gau >= map_left, map_gau >= map_right, map_gau >= map_up,
             map_gau >= map_down, map_gau > param_thre1))

        peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0])  # note reverse
        peaks = list(peaks)
        peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
        ids = range(peak_counter, peak_counter + len(peaks))
        peaks_with_score_and_id = [peaks_with_score[i] + (ids[i],) for i in range(len(ids))]

        all_peaks.append(peaks_with_score_and_id)
        peak_counter += len(peaks)

    return all_peaks
Exemple #3
0
def get_subsets(connection_all, special_k, all_peaks):
    # last number in each row is the total parts number of that person
    # the second last number in each row is the score of the overall configuration
    subset = -1 * np.ones((0, 20))
    candidate = np.array([item for sublist in all_peaks for item in sublist])

    for k in range(len(map_ids)):
        if k not in special_k:
            partAs = connection_all[k][:, 0]
            partBs = connection_all[k][:, 1]
            indexA, indexB = np.array(limb_seq[k]) - 1

            for i in range(len(connection_all[k])):  # = 1:size(temp,1)
                found = 0
                subset_idx = [-1, -1]
                for j in range(len(subset)):  # 1:size(subset,1):
                    if subset[j][indexA] == partAs[i] or subset[j][
                            indexB] == partBs[i]:
                        subset_idx[found] = j
                        found += 1

                if found == 1:
                    j = subset_idx[0]
                    if (subset[j][indexB] != partBs[i]):
                        subset[j][indexB] = partBs[i]
                        subset[j][-1] += 1
                        subset[j][-2] += candidate[partBs[i].astype(int),
                                                   2] + connection_all[k][i][2]
                elif found == 2:  # if found 2 and disjoint, merge them
                    j1, j2 = subset_idx
                    print("found = 2")
                    membership = ((subset[j1] >= 0).astype(int) +
                                  (subset[j2] >= 0).astype(int))[:-2]
                    if len(np.nonzero(membership == 2)[0]) == 0:  # merge
                        subset[j1][:-2] += (subset[j2][:-2] + 1)
                        subset[j1][-2:] += subset[j2][-2:]
                        subset[j1][-2] += connection_all[k][i][2]
                        subset = np.delete(subset, j2, 0)
                    else:  # as like found == 1
                        subset[j1][indexB] = partBs[i]
                        subset[j1][-1] += 1
                        subset[j1][-2] += candidate[
                            partBs[i].astype(int), 2] + connection_all[k][i][2]

                # if find no partA in the subset, create a new subset
                elif not found and k < 17:
                    row = -1 * np.ones(20)
                    row[indexA] = partAs[i]
                    row[indexB] = partBs[i]
                    row[-1] = 2
                    row[-2] = sum(
                        candidate[connection_all[k][i, :2].astype(int),
                                  2]) + connection_all[k][i][2]
                    subset = np.vstack([subset, row])
    return subset, candidate
Exemple #4
0
    def detect(self, minibatch, humans, scales):
        ms = minibatch.shape
        tic = time.time()
        imageToTest = Variable(T.transpose(
            T.transpose((torch.from_numpy(minibatch).float() / 256.0) - 0.5, 2,
                        3), 1, 2),
                               volatile=True).cuda()

        output1, output2 = self.model(imageToTest)

        heatmap_avg = nn.UpsamplingBilinear2d((ms[1], ms[2])).cuda()(output2)

        paf_avg = nn.UpsamplingBilinear2d((ms[1], ms[2])).cuda()(output1)

        heatmap_avg = T.transpose(heatmap_avg, 1, 2)
        heatmap_avg = T.transpose(heatmap_avg, 2, 3)

        heatmap_avg = heatmap_avg.cpu().data.numpy()
        #print 'heatmap shape : ' , heatmap_avg.shape

        paf_avg = paf_avg.cpu().data.numpy()

        heatmap_avg = heatmap_avg[:, :, :, :-1]

        map = gaussian_filter1d(heatmap_avg, sigma=3, axis=2)
        map = gaussian_filter1d(map, sigma=3, axis=1)

        map_left = np.zeros(map.shape)
        map_left[:, 1:, :, :] = map[:, :-1, :, :]
        map_right = np.zeros(map.shape)
        map_right[:, :-1, :, :] = map[:, 1:, :, :]
        map_up = np.zeros(map.shape)
        map_up[:, :, 1:, :] = map[:, :, :-1, :]
        map_down = np.zeros(map.shape)
        map_down[:, :, :-1, :] = map[:, :, 1:, :]

        peaks_binary = np.logical_and.reduce(
            (map >= map_left, map >= map_right, map >= map_up, map >= map_down,
             map > self.param_['thre1']))
        peaks = np.nonzero(peaks_binary)

        for i in range(peaks[0].shape[0]):
            human_idx = int(peaks[0][i])

            x = int(peaks[1][i] / scales[human_idx])
            y = int(peaks[2][i] / scales[human_idx])
            #x = peaks[1][i]
            #y = peaks[2][i]
            joint = peaks[3][i]
            humans[human_idx].joints[joint * 2] = x
            humans[human_idx].joints[joint * 2 + 1] = y

        return humans
def extract_paf_info(img_raw, paf_avg, all_peaks, param_thre2=0.05, param_thre3=0.5):
    connection_all = []
    special_k = []
    mid_num = 10

    for k in range(len(map_ids)):
        score_mid = paf_avg[:, :, [x - 19 for x in map_ids[k]]]
        candA = all_peaks[limb_seq[k][0] - 1]
        candB = all_peaks[limb_seq[k][1] - 1]
        nA = len(candA)
        nB = len(candB)
        if nA != 0 and nB != 0:
            connection_candidate = []
            for i in range(nA):
                for j in range(nB):
                    vec = np.subtract(candB[j][:2], candA[i][:2])
                    norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
                    vec = np.divide(vec, norm)

                    startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num),
                                   np.linspace(candA[i][1], candB[j][1], num=mid_num))
                    startend = list(startend)

                    vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0]
                                      for I in range(len(startend))])
                    vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1]
                                      for I in range(len(startend))])

                    score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
                    score_with_dist_prior = sum(score_midpts) / len(score_midpts)
                    score_with_dist_prior += min(0.5 * img_raw.shape[0] / norm - 1, 0)

                    criterion1 = len(np.nonzero(score_midpts > param_thre2)[0]) > 0.8 * len(score_midpts)
                    criterion2 = score_with_dist_prior > 0
                    if criterion1 and criterion2:
                        connection_candidate.append(
                            [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])

            connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
            connection = np.zeros((0, 5))
            for c in range(len(connection_candidate)):
                i, j, s = connection_candidate[c][0:3]
                if i not in connection[:, 3] and j not in connection[:, 4]:
                    connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
                    if len(connection) >= min(nA, nB):
                        break

            connection_all.append(connection)
        else:
            special_k.append(k)
            connection_all.append([])

    return special_k, connection_all
def get_subsets(connection_all, special_k, all_peaks):
    # last number in each row is the total parts number of that person
    # the second last number in each row is the score of the overall configuration
    subset = -1 * np.ones((0, 20))
    candidate = np.array([item for sublist in all_peaks for item in sublist])

    for k in range(len(map_ids)):
        if k not in special_k:
            partAs = connection_all[k][:, 0]
            partBs = connection_all[k][:, 1]
            indexA, indexB = np.array(limb_seq[k]) - 1

            for i in range(len(connection_all[k])):  # = 1:size(temp,1)
                found = 0
                subset_idx = [-1, -1]
                for j in range(len(subset)):  # 1:size(subset,1):
                    if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
                        subset_idx[found] = j
                        found += 1

                if found == 1:
                    j = subset_idx[0]
                    if (subset[j][indexB] != partBs[i]):
                        subset[j][indexB] = partBs[i]
                        subset[j][-1] += 1
                        subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
                elif found == 2:  # if found 2 and disjoint, merge them
                    j1, j2 = subset_idx
                    print("found = 2")
                    membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
                    if len(np.nonzero(membership == 2)[0]) == 0:  # merge
                        subset[j1][:-2] += (subset[j2][:-2] + 1)
                        subset[j1][-2:] += subset[j2][-2:]
                        subset[j1][-2] += connection_all[k][i][2]
                        subset = np.delete(subset, j2, 0)
                    else:  # as like found == 1
                        subset[j1][indexB] = partBs[i]
                        subset[j1][-1] += 1
                        subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]

                # if find no partA in the subset, create a new subset
                elif not found and k < 17:
                    row = -1 * np.ones(20)
                    row[indexA] = partAs[i]
                    row[indexB] = partBs[i]
                    row[-1] = 2
                    row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
                    subset = np.vstack([subset, row])
    return subset, candidate
Exemple #7
0
        map_left = np.zeros(map.shape)
        map_left[1:, :] = map[:-1, :]
        map_right = np.zeros(map.shape)
        map_right[:-1, :] = map[1:, :]
        map_up = np.zeros(map.shape)
        map_up[:, 1:] = map[:, :-1]
        map_down = np.zeros(map.shape)
        map_down[:, :-1] = map[:, 1:]

        peaks_binary = np.logical_and.reduce(
            (map >= map_left, map >= map_right, map >= map_up, map >= map_down,
             map > param_['thre1']))
        #    peaks_binary = T.eq(
        #    peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])

        peaks = zip(np.nonzero(peaks_binary)[1],
                    np.nonzero(peaks_binary)[0])  # note reverse

        peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
        id = range(peak_counter, peak_counter + len(peaks))
        peaks_with_score_and_id = [
            peaks_with_score[i] + (id[i], ) for i in range(len(id))
        ]

        all_peaks.append(peaks_with_score_and_id)
        peak_counter += len(peaks)

    raw_key_points = []
    for item in all_peaks:
        if item != []:
            raw_key_points.append(item[0][0:2])
Exemple #8
0
    map_left = np.zeros(map.shape)
    map_left[1:, :] = map[:-1, :]
    map_right = np.zeros(map.shape)
    map_right[:-1, :] = map[1:, :]
    map_up = np.zeros(map.shape)
    map_up[:, 1:] = map[:, :-1]
    map_down = np.zeros(map.shape)
    map_down[:, :-1] = map[:, 1:]

    peaks_binary = np.logical_and.reduce(
        (map >= map_left, map >= map_right, map >= map_up, map >= map_down,
         map > param_['thre1']))
    #    peaks_binary = T.eq(
    #    peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])

    peaks = zip(np.nonzero(peaks_binary)[1],
                np.nonzero(peaks_binary)[0])  # note reverse

    peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
    id = range(peak_counter, peak_counter + len(peaks))
    peaks_with_score_and_id = [
        peaks_with_score[i] + (id[i], ) for i in range(len(id))
    ]

    all_peaks.append(peaks_with_score_and_id)
    peak_counter += len(peaks)

connection_all = []
special_k = []
mid_num = 10
Exemple #9
0
def extract_paf_info(img_raw,
                     paf_avg,
                     all_peaks,
                     param_thre2=0.05,
                     param_thre3=0.5):
    connection_all = []
    special_k = []
    mid_num = 10

    for k in range(len(map_ids)):
        score_mid = paf_avg[:, :, [x - 19 for x in map_ids[k]]]
        candA = all_peaks[limb_seq[k][0] - 1]
        candB = all_peaks[limb_seq[k][1] - 1]
        nA = len(candA)
        nB = len(candB)
        if nA != 0 and nB != 0:
            connection_candidate = []
            for i in range(nA):
                for j in range(nB):
                    vec = np.subtract(candB[j][:2], candA[i][:2])
                    norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
                    vec = np.divide(vec, norm)

                    startend = zip(
                        np.linspace(candA[i][0], candB[j][0], num=mid_num),
                        np.linspace(candA[i][1], candB[j][1], num=mid_num))
                    startend = list(startend)

                    vec_x = np.array([
                        score_mid[int(round(startend[I][1])),
                                  int(round(startend[I][0])), 0]
                        for I in range(len(startend))
                    ])
                    vec_y = np.array([
                        score_mid[int(round(startend[I][1])),
                                  int(round(startend[I][0])), 1]
                        for I in range(len(startend))
                    ])

                    score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
                        vec_y, vec[1])
                    score_with_dist_prior = sum(score_midpts) / len(
                        score_midpts)
                    score_with_dist_prior += min(
                        0.5 * img_raw.shape[0] / norm - 1, 0)

                    criterion1 = len(
                        np.nonzero(score_midpts > param_thre2)
                        [0]) > 0.8 * len(score_midpts)
                    criterion2 = score_with_dist_prior > 0
                    if criterion1 and criterion2:
                        connection_candidate.append([
                            i, j, score_with_dist_prior,
                            score_with_dist_prior + candA[i][2] + candB[j][2]
                        ])

            connection_candidate = sorted(connection_candidate,
                                          key=lambda x: x[2],
                                          reverse=True)
            connection = np.zeros((0, 5))
            for c in range(len(connection_candidate)):
                i, j, s = connection_candidate[c][0:3]
                if i not in connection[:, 3] and j not in connection[:, 4]:
                    connection = np.vstack(
                        [connection, [candA[i][3], candB[j][3], s, i, j]])
                    if len(connection) >= min(nA, nB):
                        break

            connection_all.append(connection)
        else:
            special_k.append(k)
            connection_all.append([])

    return special_k, connection_all
def handle_one(oriImg):

    # for visualize
    canvas = np.copy(oriImg)
    imageToTest = Variable(T.transpose(
        T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(), 0), 2, 3), 1,
        2),
                           volatile=True).cuda()
    print oriImg.shape
    scale = model_['boxsize'] / float(oriImg.shape[0])
    print scale
    h = int(oriImg.shape[0] * scale)
    w = int(oriImg.shape[1] * scale)
    pad_h = 0 if (h % model_['stride']
                  == 0) else model_['stride'] - (h % model_['stride'])
    pad_w = 0 if (w % model_['stride']
                  == 0) else model_['stride'] - (w % model_['stride'])
    new_h = h + pad_h
    new_w = w + pad_w

    imageToTest = cv2.resize(oriImg, (0, 0),
                             fx=scale,
                             fy=scale,
                             interpolation=cv2.INTER_CUBIC)
    imageToTest_padded, pad = util.padRightDownCorner(imageToTest,
                                                      model_['stride'],
                                                      model_['padValue'])
    imageToTest_padded = np.transpose(
        np.float32(imageToTest_padded[:, :, :, np.newaxis]),
        (3, 2, 0, 1)) / 256 - 0.5

    feed = Variable(T.from_numpy(imageToTest_padded)).cuda()

    output1, output2 = model(feed)

    heatmap = nn.UpsamplingBilinear2d(
        (oriImg.shape[0], oriImg.shape[1])).cuda()(output2)

    paf = nn.UpsamplingBilinear2d(
        (oriImg.shape[0], oriImg.shape[1])).cuda()(output1)

    print heatmap.size()
    print paf.size()
    print type(heatmap)
    heatmap_avg = T.transpose(T.transpose(heatmap[0], 0, 1), 1,
                              2).data.cpu().numpy()
    paf_avg = T.transpose(T.transpose(paf[0], 0, 1), 1, 2).data.cpu().numpy()

    all_peaks = []
    peak_counter = 0

    #maps =
    for part in range(18):
        map_ori = heatmap_avg[:, :, part]
        map = gaussian_filter(map_ori, sigma=3)

        map_left = np.zeros(map.shape)
        map_left[1:, :] = map[:-1, :]
        map_right = np.zeros(map.shape)
        map_right[:-1, :] = map[1:, :]
        map_up = np.zeros(map.shape)
        map_up[:, 1:] = map[:, :-1]
        map_down = np.zeros(map.shape)
        map_down[:, :-1] = map[:, 1:]

        peaks_binary = np.logical_and.reduce(
            (map >= map_left, map >= map_right, map >= map_up, map >= map_down,
             map > param_['thre1']))
        #    peaks_binary = T.eq(
        #    peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])

        peaks = zip(np.nonzero(peaks_binary)[1],
                    np.nonzero(peaks_binary)[0])  # note reverse

        peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
        id = range(peak_counter, peak_counter + len(peaks))
        peaks_with_score_and_id = [
            peaks_with_score[i] + (id[i], ) for i in range(len(id))
        ]

        all_peaks.append(peaks_with_score_and_id)
        peak_counter += len(peaks)

    connection_all = []
    special_k = []
    mid_num = 10

    for k in range(len(mapIdx)):
        score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
        candA = all_peaks[limbSeq[k][0] - 1]
        candB = all_peaks[limbSeq[k][1] - 1]
        nA = len(candA)
        nB = len(candB)
        indexA, indexB = limbSeq[k]
        if (nA != 0 and nB != 0):
            connection_candidate = []
            for i in range(nA):
                for j in range(nB):
                    vec = np.subtract(candB[j][:2], candA[i][:2])
                    norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
                    vec = np.divide(vec, norm)

                    startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
                                   np.linspace(candA[i][1], candB[j][1], num=mid_num))

                    vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
                                      for I in range(len(startend))])
                    vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
                                      for I in range(len(startend))])

                    score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
                        vec_y, vec[1])
                    score_with_dist_prior = sum(
                        score_midpts) / len(score_midpts) + min(
                            0.5 * oriImg.shape[0] / norm - 1, 0)
                    criterion1 = len(
                        np.nonzero(score_midpts > param_['thre2'])
                        [0]) > 0.8 * len(score_midpts)
                    criterion2 = score_with_dist_prior > 0
                    if criterion1 and criterion2:
                        connection_candidate.append([
                            i, j, score_with_dist_prior,
                            score_with_dist_prior + candA[i][2] + candB[j][2]
                        ])

            connection_candidate = sorted(connection_candidate,
                                          key=lambda x: x[2],
                                          reverse=True)
            connection = np.zeros((0, 5))
            for c in range(len(connection_candidate)):
                i, j, s = connection_candidate[c][0:3]
                if (i not in connection[:, 3] and j not in connection[:, 4]):
                    connection = np.vstack(
                        [connection, [candA[i][3], candB[j][3], s, i, j]])
                    if (len(connection) >= min(nA, nB)):
                        break

            connection_all.append(connection)
        else:
            special_k.append(k)
            connection_all.append([])

    # last number in each row is the total parts number of that person
    # the second last number in each row is the score of the overall configuration
    subset = -1 * np.ones((0, 20))
    candidate = np.array([item for sublist in all_peaks for item in sublist])

    for k in range(len(mapIdx)):
        if k not in special_k:
            partAs = connection_all[k][:, 0]
            partBs = connection_all[k][:, 1]
            indexA, indexB = np.array(limbSeq[k]) - 1

            for i in range(len(connection_all[k])):  #= 1:size(temp,1)
                found = 0
                subset_idx = [-1, -1]
                for j in range(len(subset)):  #1:size(subset,1):
                    if subset[j][indexA] == partAs[i] or subset[j][
                            indexB] == partBs[i]:
                        subset_idx[found] = j
                        found += 1

                if found == 1:
                    j = subset_idx[0]
                    if (subset[j][indexB] != partBs[i]):
                        subset[j][indexB] = partBs[i]
                        subset[j][-1] += 1
                        subset[j][-2] += candidate[partBs[i].astype(int),
                                                   2] + connection_all[k][i][2]
                elif found == 2:  # if found 2 and disjoint, merge them
                    j1, j2 = subset_idx
                    print "found = 2"
                    membership = ((subset[j1] >= 0).astype(int) +
                                  (subset[j2] >= 0).astype(int))[:-2]
                    if len(np.nonzero(membership == 2)[0]) == 0:  #merge
                        subset[j1][:-2] += (subset[j2][:-2] + 1)
                        subset[j1][-2:] += subset[j2][-2:]
                        subset[j1][-2] += connection_all[k][i][2]
                        subset = np.delete(subset, j2, 0)
                    else:  # as like found == 1
                        subset[j1][indexB] = partBs[i]
                        subset[j1][-1] += 1
                        subset[j1][-2] += candidate[
                            partBs[i].astype(int), 2] + connection_all[k][i][2]

                # if find no partA in the subset, create a new subset
                elif not found and k < 17:
                    row = -1 * np.ones(20)
                    row[indexA] = partAs[i]
                    row[indexB] = partBs[i]
                    row[-1] = 2
                    row[-2] = sum(
                        candidate[connection_all[k][i, :2].astype(int),
                                  2]) + connection_all[k][i][2]
                    subset = np.vstack([subset, row])

    # delete some rows of subset which has few parts occur
    deleteIdx = []
    for i in range(len(subset)):
        if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
            deleteIdx.append(i)
    subset = np.delete(subset, deleteIdx, axis=0)

    #    canvas = cv2.imread(test_image) # B,G,R order
    for i in range(18):
        for j in range(len(all_peaks[i])):
            cv2.circle(canvas,
                       all_peaks[i][j][0:2],
                       4,
                       colors[i],
                       thickness=-1)

    stickwidth = 4

    for i in range(17):
        for n in range(len(subset)):
            index = subset[n][np.array(limbSeq[i]) - 1]
            if -1 in index:
                continue
            cur_canvas = canvas.copy()
            Y = candidate[index.astype(int), 0]
            X = candidate[index.astype(int), 1]
            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(cur_canvas, polygon, colors[i])
            canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)

    return canvas
Exemple #11
0
def post_process(oriImg, canvas, paf_avg, all_peak_idxs, nonzero_vals,
                 dont_draw):

    all_peaks = [[] for x in range(N_JOINTS)]

    for i in xrange(all_peak_idxs.shape[0]):
        all_peaks[all_peak_idxs[i, 0]].append(
            (all_peak_idxs[i, 2], all_peak_idxs[i, 1], nonzero_vals[i], i))

    PEAKT = time.time()

    connection_all = []
    special_k = []
    mid_num = 10

    KTTT = time.time()

    # don't really know how this works
    for k in xrange(len(mapIdx)):
        idxs = []
        for x in mapIdx[k]:
            idxs.append(x - 19)

        score_mid = paf_avg[idxs, :, :]
        candA = all_peaks[limbSeq[k][0] - 1]
        candB = all_peaks[limbSeq[k][1] - 1]
        nA = len(candA)
        nB = len(candB)

        indexA, indexB = limbSeq[k]
        if (nA != 0 and nB != 0):
            connection_candidate = []
            for i in range(nA):
                for j in range(nB):
                    vec = np.subtract(candB[j][:2], candA[i][:2])
                    norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
                    if norm == 0:
                        continue
                    vec = np.divide(vec, norm)

                    startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
                                   np.linspace(candA[i][1], candB[j][1], num=mid_num))

                    vec_x = np.zeros(len(startend))
                    vec_y = np.zeros(len(startend))

                    for I in xrange(len(startend)):
                        vec_x[I] = score_mid[0,
                                             int(round(startend[I][1])),
                                             int(round(startend[I][0]))]
                        vec_y[I] = score_mid[1,
                                             int(round(startend[I][1])),
                                             int(round(startend[I][0]))]

                    score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
                        vec_y, vec[1])
                    score_with_dist_prior = np.sum(
                        score_midpts) / len(score_midpts) + min(
                            0.5 * oriImg.shape[0] / norm - 1, 0)
                    criterion1 = len(
                        np.nonzero(score_midpts > param_['thre2'])
                        [0]) > 0.8 * len(score_midpts)
                    criterion2 = score_with_dist_prior > 0
                    if criterion1 and criterion2:
                        connection_candidate.append([
                            i, j, score_with_dist_prior,
                            score_with_dist_prior + candA[i][2] + candB[j][2]
                        ])

            connection_candidate = sorted(connection_candidate,
                                          key=get_third_item,
                                          reverse=True)
            connection = np.zeros((0, 5))
            for c in range(len(connection_candidate)):
                i, j, s = connection_candidate[c][0:3]
                if (i not in connection[:, 3] and j not in connection[:, 4]):
                    connection = np.vstack(
                        [connection, [candA[i][3], candB[j][3], s, i, j]])
                    if (len(connection) >= min(nA, nB)):
                        break

            connection_all.append(connection)
        else:
            special_k.append(k)
            connection_all.append([])

    # last number in each row is the total parts number of that person
    # the second last number in each row is the score of the overall configuration
    subset = -1 * np.ones((0, 20))
    cand_tmp = []
    for sublist in all_peaks:
        for item in sublist:
            cand_tmp.append(item)

    candidate = np.asarray(cand_tmp)

    # don't really know how this works either
    for k in xrange(len(mapIdx)):
        if k not in special_k:
            partAs = connection_all[k][:, 0]
            partBs = connection_all[k][:, 1]
            indexA, indexB = np.array(limbSeq[k]) - 1

            for i in range(len(connection_all[k])):  #= 1:size(temp,1)
                found = 0
                subset_idx = [-1, -1]
                for j in range(len(subset)):  #1:size(subset,1):
                    if subset[j][indexA] == partAs[i] or subset[j][
                            indexB] == partBs[i]:
                        subset_idx[found] = j
                        found += 1

                if found == 1:
                    j = subset_idx[0]
                    if (subset[j][indexB] != partBs[i]):
                        subset[j][indexB] = partBs[i]
                        subset[j][-1] += 1
                        subset[j][-2] += candidate[partBs[i].astype(int),
                                                   2] + connection_all[k][i][2]
                elif found == 2:  # if found 2 and disjoint, merge them
                    j1, j2 = subset_idx
                    #print "found = 2"
                    membership = ((subset[j1] >= 0).astype(int) +
                                  (subset[j2] >= 0).astype(int))[:-2]
                    if len(np.nonzero(membership == 2)[0]) == 0:  #merge
                        subset[j1][:-2] += (subset[j2][:-2] + 1)
                        subset[j1][-2:] += subset[j2][-2:]
                        subset[j1][-2] += connection_all[k][i][2]
                        subset = np.delete(subset, j2, 0)
                    else:  # as like found == 1
                        subset[j1][indexB] = partBs[i]
                        subset[j1][-1] += 1
                        subset[j1][-2] += candidate[
                            partBs[i].astype(int), 2] + connection_all[k][i][2]

                # if find no partA in the subset, create a new subset
                elif not found and k < 17:
                    row = -1 * np.ones(20)
                    row[indexA] = partAs[i]
                    row[indexB] = partBs[i]
                    row[-1] = 2
                    row[-2] = np.sum(
                        candidate[connection_all[k][i, :2].astype(int),
                                  2]) + connection_all[k][i][2]
                    subset = np.vstack([subset, row])

    #print subset.shape
    # delete some rows of subset which has few parts occur
    deleteIdx = []
    for i in range(len(subset)):
        if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
            deleteIdx.append(i)
    subset = np.delete(subset, deleteIdx, axis=0)
    N_BONES = N_JOINTS - 1

    if dont_draw:
        res = {'found_ppl': []}
        for n in range(len(subset)):
            skel = {}
            for i in range(N_BONES):
                index = subset[n][np.array(limbSeq[i]) - 1]
                if -1 in index:
                    continue
                Y = candidate[index.astype(int), 0]
                X = candidate[index.astype(int), 1]
                bone_name = bone_map[i]
                skel[bone_name] = ((X[0], Y[0]), (X[1], Y[1]))
            res['found_ppl'].append(skel)
        return res

    p = time.time()
    canvas = canvas.copy()

    for i in range(N_JOINTS):
        for j in range(len(all_peaks[i])):
            cv2.circle(canvas,
                       all_peaks[i][j][0:2],
                       4,
                       colors[i],
                       thickness=-1)

    stickwidth = 4

    N_BONES = N_JOINTS - 1

    for i in range(N_BONES):
        for n in range(len(subset)):
            index = subset[n][np.array(limbSeq[i]) - 1]
            if -1 in index:
                continue
            cur_canvas = canvas.copy()
            Y = candidate[index.astype(int), 0]
            X = candidate[index.astype(int), 1]
            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(cur_canvas, polygon, colors[i])
            cv2.putText(canvas, str(i), (int(mY), int(mX)),
                        cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 3)
            cv2.putText(canvas, str(i), (int(mY), int(mX)),
                        cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 0), 2)
            canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)

    print('drawing took: ', time.time() - p)
    return canvas
def handle_one(oriImg):
    
    # for visualize
    canvas = np.copy(oriImg)
    imageToTest = Variable(T.transpose(T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(),0),2,3),1,2),volatile=True).cuda()
    print oriImg.shape
    scale = model_['boxsize'] / float(oriImg.shape[0])
    print scale
    h = int(oriImg.shape[0]*scale)
    w = int(oriImg.shape[1]*scale)
    pad_h = 0 if (h%model_['stride']==0) else model_['stride'] - (h % model_['stride']) 
    pad_w = 0 if (w%model_['stride']==0) else model_['stride'] - (w % model_['stride'])
    new_h = h+pad_h
    new_w = w+pad_w

    imageToTest = cv2.resize(oriImg, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
    imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_['stride'], model_['padValue'])
    imageToTest_padded = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,2,0,1))/256 - 0.5

    feed = Variable(T.from_numpy(imageToTest_padded)).cuda()      

    output1,output2 = model(feed)

    heatmap = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output2)

    paf = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output1)       
    
    print heatmap.size()
    print paf.size()
    print type(heatmap)
    heatmap_avg = T.transpose(T.transpose(heatmap[0],0,1),1,2).data.cpu().numpy()
    paf_avg = T.transpose(T.transpose(paf[0],0,1),1,2).data.cpu().numpy()
        
    all_peaks = []
    peak_counter = 0

    #maps = 
    for part in range(18):
        map_ori = heatmap_avg[:,:,part]
        map = gaussian_filter(map_ori, sigma=3)
        
        map_left = np.zeros(map.shape)
        map_left[1:,:] = map[:-1,:]
        map_right = np.zeros(map.shape)
        map_right[:-1,:] = map[1:,:]
        map_up = np.zeros(map.shape)
        map_up[:,1:] = map[:,:-1]
        map_down = np.zeros(map.shape)
        map_down[:,:-1] = map[:,1:]
        
        peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param_['thre1']))
    #    peaks_binary = T.eq(
    #    peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
        
        peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse
        
        peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]
        id = range(peak_counter, peak_counter + len(peaks))
        peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))]

        all_peaks.append(peaks_with_score_and_id)
        peak_counter += len(peaks)
        
        
        
        
        
    connection_all = []
    special_k = []
    mid_num = 10

    for k in range(len(mapIdx)):
        score_mid = paf_avg[:,:,[x-19 for x in mapIdx[k]]]
        candA = all_peaks[limbSeq[k][0]-1]
        candB = all_peaks[limbSeq[k][1]-1]
        nA = len(candA)
        nB = len(candB)
        indexA, indexB = limbSeq[k]
        if(nA != 0 and nB != 0):
            connection_candidate = []
            for i in range(nA):
                for j in range(nB):
                    vec = np.subtract(candB[j][:2], candA[i][:2])
                    norm = math.sqrt(vec[0]*vec[0] + vec[1]*vec[1])
                    vec = np.divide(vec, norm)
                    
                    startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
                                   np.linspace(candA[i][1], candB[j][1], num=mid_num))
                    
                    vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
                                      for I in range(len(startend))])
                    vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
                                      for I in range(len(startend))])

                    score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
                    score_with_dist_prior = sum(score_midpts)/len(score_midpts) + min(0.5*oriImg.shape[0]/norm-1, 0)
                    criterion1 = len(np.nonzero(score_midpts > param_['thre2'])[0]) > 0.8 * len(score_midpts)
                    criterion2 = score_with_dist_prior > 0
                    if criterion1 and criterion2:
                        connection_candidate.append([i, j, score_with_dist_prior, score_with_dist_prior+candA[i][2]+candB[j][2]])

            connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
            connection = np.zeros((0,5))
            for c in range(len(connection_candidate)):
                i,j,s = connection_candidate[c][0:3]
                if(i not in connection[:,3] and j not in connection[:,4]):
                    connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
                    if(len(connection) >= min(nA, nB)):
                        break

            connection_all.append(connection)
        else:
            special_k.append(k)
            connection_all.append([])

    # last number in each row is the total parts number of that person
    # the second last number in each row is the score of the overall configuration
    subset = -1 * np.ones((0, 20))
    candidate = np.array([item for sublist in all_peaks for item in sublist])

    for k in range(len(mapIdx)):
        if k not in special_k:
            partAs = connection_all[k][:,0]
            partBs = connection_all[k][:,1]
            indexA, indexB = np.array(limbSeq[k]) - 1

            for i in range(len(connection_all[k])): #= 1:size(temp,1)
                found = 0
                subset_idx = [-1, -1]
                for j in range(len(subset)): #1:size(subset,1):
                    if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
                        subset_idx[found] = j
                        found += 1
                
                if found == 1:
                    j = subset_idx[0]
                    if(subset[j][indexB] != partBs[i]):
                        subset[j][indexB] = partBs[i]
                        subset[j][-1] += 1
                        subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
                elif found == 2: # if found 2 and disjoint, merge them
                    j1, j2 = subset_idx
                    print "found = 2"
                    membership = ((subset[j1]>=0).astype(int) + (subset[j2]>=0).astype(int))[:-2]
                    if len(np.nonzero(membership == 2)[0]) == 0: #merge
                        subset[j1][:-2] += (subset[j2][:-2] + 1)
                        subset[j1][-2:] += subset[j2][-2:]
                        subset[j1][-2] += connection_all[k][i][2]
                        subset = np.delete(subset, j2, 0)
                    else: # as like found == 1
                        subset[j1][indexB] = partBs[i]
                        subset[j1][-1] += 1
                        subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]

                # if find no partA in the subset, create a new subset
                elif not found and k < 17:
                    row = -1 * np.ones(20)
                    row[indexA] = partAs[i]
                    row[indexB] = partBs[i]
                    row[-1] = 2
                    row[-2] = sum(candidate[connection_all[k][i,:2].astype(int), 2]) + connection_all[k][i][2]
                    subset = np.vstack([subset, row])

    # delete some rows of subset which has few parts occur
    deleteIdx = [];
    for i in range(len(subset)):
        if subset[i][-1] < 4 or subset[i][-2]/subset[i][-1] < 0.4:
            deleteIdx.append(i)
    subset = np.delete(subset, deleteIdx, axis=0)

#    canvas = cv2.imread(test_image) # B,G,R order
    for i in range(18):
        for j in range(len(all_peaks[i])):
            cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)

    stickwidth = 4

    for i in range(17):
        for n in range(len(subset)):
            index = subset[n][np.array(limbSeq[i])-1]
            if -1 in index:
                continue
            cur_canvas = canvas.copy()
            Y = candidate[index.astype(int), 0]
            X = candidate[index.astype(int), 1]
            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(cur_canvas, polygon, colors[i])
            canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)

    return canvas
def process_image(oriImg, model, model_params, jjac_info):
    # for visualize
    t0 = time.time()
    canvas = np.copy(oriImg)
    #imageToTest = Variable(T.transpose(T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(), 0), 2, 3), 1, 2),
    #                       volatile=True).cuda()
    #print oriImg.shape
    tic = time.time()
    scale = model_params['model_']['boxsize'] / float(oriImg.shape[0])
    #print scale
    h = int(oriImg.shape[0] * scale)
    w = int(oriImg.shape[1] * scale)
    pad_h = 0 if (h % model_params['model_']['stride']
                  == 0) else model_params['model_']['stride'] - (
                      h % model_params['model_']['stride'])
    pad_w = 0 if (w % model_params['model_']['stride']
                  == 0) else model_params['model_']['stride'] - (
                      w % model_params['model_']['stride'])
    new_h = h + pad_h
    new_w = w + pad_w
    #print 'scaled width and height ({}, {})'.format(h, w)

    imageToTest = cv2.resize(oriImg, (0, 0),
                             fx=scale,
                             fy=scale,
                             interpolation=cv2.INTER_CUBIC)
    imageToTest_padded, pad = util.padRightDownCorner(
        imageToTest, model_params['model_']['stride'],
        model_params['model_']['padValue'])
    imageToTest_padded = np.transpose(
        np.float32(imageToTest_padded[:, :, :, np.newaxis]),
        (3, 2, 0, 1)) / 256 - 0.5

    feed = Variable(T.from_numpy(imageToTest_padded)).cuda()

    output1, output2 = model(feed)

    heatmap = nn.UpsamplingBilinear2d(
        (oriImg.shape[0], oriImg.shape[1])).cuda()(output2)

    paf = nn.UpsamplingBilinear2d(
        (oriImg.shape[0], oriImg.shape[1])).cuda()(output1)

    toc = time.time()
    model_running_time = toc - tic
    #print heatmap.size() # (360, 640, 3)
    #print paf.size()
    #print type(heatmap)

    tic = time.time()
    heatmap_avg = T.transpose(T.transpose(heatmap[0], 0, 1), 1,
                              2).data.cpu().numpy()
    #paf_avg = T.transpose(T.transpose(paf[0], 0, 1), 1, 2).data.cpu().numpy()

    all_peaks = []
    peak_counter = 0
    # 13-17 are head
    # 10, 11, 12, 13 are legs
    # hand_parts = [5, 6, 7, 8, 15, 16]

    # 4 is right hand

    # todo is it actually left elbow below because inverse? confirm
    # 5 was left elbow
    # 6 was left shoulder
    # 7 is left hand. Right hand only?
    # 15 is head. right part of head
    body_part_index_to_name = {
        '4': 'Right Hand',
        '5': 'Left Elbow',
        '6': 'Left Shoulder',
        '7': 'Left Hand',
        '15': 'Left Eye',
        '14': '1',
        '16': '2',
        '17': '3'
    }
    hand_parts = [4, 7, 15, 16, 17]
    # hand_parts = list(range(18))
    # for part in range(18):
    for part in hand_parts:
        map_ori = heatmap_avg[:, :, part]
        map = gaussian_filter(map_ori, sigma=3)

        map_left = np.zeros(map.shape)
        map_left[1:, :] = map[:-1, :]
        map_right = np.zeros(map.shape)
        map_right[:-1, :] = map[1:, :]
        map_up = np.zeros(map.shape)
        map_up[:, 1:] = map[:, :-1]
        map_down = np.zeros(map.shape)
        map_down[:, :-1] = map[:, 1:]

        peaks_binary = np.logical_and.reduce(
            (map >= map_left, map >= map_right, map >= map_up, map >= map_down,
             map > model_params['param_']['thre1']))
        #    peaks_binary = T.eq(
        #    peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])

        peaks = zip(np.nonzero(peaks_binary)[1],
                    np.nonzero(peaks_binary)[0])  # note reverse

        peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
        id = range(peak_counter, peak_counter + len(peaks))
        peaks_with_score_and_id = [
            peaks_with_score[i] + (id[i], ) for i in range(len(id))
        ]

        all_peaks.append(peaks_with_score_and_id)
        # print peaks_with_score_and_id
        peak_counter += len(peaks)

    #print 'heat map peak stuff time is %.5f' % (time.time() - tic)

    # connection_all = []
    # special_k = []
    # mid_num = 10

    # for k in range(len(mapIdx)):
    #     score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
    #     candA = all_peaks[limbSeq[k][0] - 1]
    #     candB = all_peaks[limbSeq[k][1] - 1]
    #     nA = len(candA)
    #     nB = len(candB)
    #     indexA, indexB = limbSeq[k]
    #     if (nA != 0 and nB != 0):
    #         connection_candidate = []
    #         for i in range(nA):
    #             for j in range(nB):
    #                 vec = np.subtract(candB[j][:2], candA[i][:2])
    #                 norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
    #                 vec = np.divide(vec, norm)
    #
    #                 startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
    #                                np.linspace(candA[i][1], candB[j][1], num=mid_num))
    #
    #                 vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
    #                                   for I in range(len(startend))])
    #                 vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
    #                                   for I in range(len(startend))])
    #
    #                 score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
    #                 score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
    #                     0.5 * oriImg.shape[0] / norm - 1, 0)
    #                 criterion1 = len(np.nonzero(score_midpts > param_['thre2'])[0]) > 0.8 * len(score_midpts)
    #                 criterion2 = score_with_dist_prior > 0
    #                 if criterion1 and criterion2:
    #                     connection_candidate.append(
    #                         [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
    #
    #         connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
    #         connection = np.zeros((0, 5))
    #         for c in range(len(connection_candidate)):
    #             i, j, s = connection_candidate[c][0:3]
    #             if (i not in connection[:, 3] and j not in connection[:, 4]):
    #                 connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
    #                 if (len(connection) >= min(nA, nB)):
    #                     break
    #
    #         connection_all.append(connection)
    #     else:
    #         special_k.append(k)
    #         connection_all.append([])

    # last number in each row is the total parts number of that person
    # the second last number in each row is the score of the overall configuration
    # subset = -1 * np.ones((0, 20))
    # candidate = np.array([item for sublist in all_peaks for item in sublist])
    #
    # for k in range(len(mapIdx)):
    #     if k not in special_k:
    #         partAs = connection_all[k][:, 0]
    #         partBs = connection_all[k][:, 1]
    #         indexA, indexB = np.array(limbSeq[k]) - 1
    #
    #         for i in range(len(connection_all[k])):  # = 1:size(temp,1)
    #             found = 0
    #             subset_idx = [-1, -1]
    #             for j in range(len(subset)):  # 1:size(subset,1):
    #                 if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
    #                     subset_idx[found] = j
    #                     found += 1
    #
    #             if found == 1:
    #                 j = subset_idx[0]
    #                 if (subset[j][indexB] != partBs[i]):
    #                     subset[j][indexB] = partBs[i]
    #                     subset[j][-1] += 1
    #                     subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
    #             elif found == 2:  # if found 2 and disjoint, merge them
    #                 j1, j2 = subset_idx
    #                 print "found = 2"
    #                 membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
    #                 if len(np.nonzero(membership == 2)[0]) == 0:  # merge
    #                     subset[j1][:-2] += (subset[j2][:-2] + 1)
    #                     subset[j1][-2:] += subset[j2][-2:]
    #                     subset[j1][-2] += connection_all[k][i][2]
    #                     subset = np.delete(subset, j2, 0)
    #                 else:  # as like found == 1
    #                     subset[j1][indexB] = partBs[i]
    #                     subset[j1][-1] += 1
    #                     subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
    #
    #             # if find no partA in the subset, create a new subset
    #             elif not found and k < 17:
    #                 row = -1 * np.ones(20)
    #                 row[indexA] = partAs[i]
    #                 row[indexB] = partBs[i]
    #                 row[-1] = 2
    #                 row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
    #                 subset = np.vstack([subset, row])
    #
    # # delete some rows of subset which has few parts occur
    # deleteIdx = [];
    # for i in range(len(subset)):
    #     if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
    #         deleteIdx.append(i)
    # subset = np.delete(subset, deleteIdx, axis=0)

    #    canvas = cv2.imread(test_image) # B,G,R order
    found_hand = False
    found_head = False
    found_2_hands = False

    tic = time.time()
    # for i in range(18):
    for i in range(len(hand_parts)):
        for j in range(len(all_peaks[i])):
            #cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
            # Right hand is first in loop. so if it fails we have backup
            if body_part_index_to_name[str(hand_parts[i])] == 'Right Hand':  #
                # BGR order. So blue is below
                # cv2.circle(canvas, all_peaks[i][j][0:2], 4, (255, 0, 0), thickness=-1)
                #print 'Hand position:', all_peaks[i][j][0:2]
                found_hand = True
                hand_position = all_peaks[i][j][0:2]
            elif body_part_index_to_name[str(hand_parts[i])] == 'Left Hand':
                # cv2.circle(canvas, all_peaks[i][j][0:2], 4, (0, 255, 0), thickness=-1)
                if not found_hand:
                    print 'Did not find Right hand but found left hand!'
                    found_hand = True
                    hand_position = all_peaks[i][j][0:2]
                else:
                    print 'Found two hands'
                    found_2_hands = True
                    second_hand_position = all_peaks[i][j][0:2]
                #hand_position = all_peaks[i][j][0:2]
            elif body_part_index_to_name[str(hand_parts[i])] == 'Left Eye':
                # BGR order. So red is below
                # cv2.circle(canvas, all_peaks[i][j][0:2], 4, (0, 0, 255), thickness=-1)
                #print 'Head position:', all_peaks[i][j][0:2]
                found_head = True
                head_position = all_peaks[i][j][0:2]
                jjac_info['last_x_head_pos'].append(head_position[1])
                if len(jjac_info['last_x_head_pos']) > 10:
                    jjac_info['last_x_head_pos'].pop()
                    biggest_diff = max(jjac_info['last_x_head_pos']) - min(
                        jjac_info['last_x_head_pos'])
                    if biggest_diff > jjac_info['biggest_diff']:
                        jjac_info['biggest_diff'] = biggest_diff
                        # espeak_command = 'Largest y range from last_x_head_pos: {}'.format(biggest_diff)
                        espeak_command = 'Largest range from last positions: {}'.format(
                            biggest_diff)
                        print espeak_command
                        # os.system(espeak_command)
                # jjac_info['head_y_range']
            elif body_part_index_to_name[str(hand_parts[i])] == '1':
                cv2.circle(canvas,
                           all_peaks[i][j][0:2],
                           4, (0, 0, 255),
                           thickness=-1)
            elif body_part_index_to_name[str(hand_parts[i])] == '2':
                cv2.circle(canvas,
                           all_peaks[i][j][0:2],
                           4, (0, 255, 255),
                           thickness=-1)
            elif body_part_index_to_name[str(hand_parts[i])] == '3':
                cv2.circle(canvas,
                           all_peaks[i][j][0:2],
                           4, (255, 255, 255),
                           thickness=-1)
            else:
                # cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
                cv2.circle(canvas,
                           all_peaks[i][j][0:2],
                           4, (128, 128, 128),
                           thickness=-1)

    #global how_many_times_hands_went_over_head, hands_over_head

    if found_hand and found_head:
        # both hands over head or 1 hand over head
        if (found_2_hands and hand_position[1] < head_position[1]
                and second_hand_position[1] < head_position[1]
            ) or hand_position[1] < head_position[1]:

            print 'Hand is higher than head {} < {}'.format(
                hand_position[1], head_position[1])
            # if first time hand over head, then increase jumping jack count
            if not jjac_info['hands_over_head']:
                jjac_info['hands_over_head'] = True
                jjac_info['num_jumping_jacks'] += 1
                print 'Number of jumping jacks:', jjac_info[
                    'num_jumping_jacks']
                # speak speech command example
                # espeak '1 Jumping Jack'
                espeak_command = "espeak {}".format(
                    '\'{} Jumping jack{}\''.format(
                        jjac_info['num_jumping_jacks'],
                        's' if jjac_info['num_jumping_jacks'] > 1 else ''))
                # os.system(espeak_command) # slows it down?
        else:
            jjac_info['hands_over_head'] = False
            print 'hand is lower than head {} > {}'.format(
                hand_position[1], head_position[1])
        #print 'Drawing circles time is %.5f' % (time.time() - tic)

    cv2.putText(
        canvas,
        "No. Jumping Jacks: {} ".format(jjac_info['num_jumping_jacks']),
        (0, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, 255)
    cv2.putText(
        canvas, "Hands over head: {}".format(
            'Yes' if jjac_info['hands_over_head'] else 'No'), (0, 45),
        cv2.FONT_HERSHEY_SIMPLEX, 0.6, 255)

    # stickwidth = 4

    # for i in range(17):
    #     for n in range(len(subset)):
    #         index = subset[n][np.array(limbSeq[i]) - 1]
    #         if -1 in index:
    #             continue
    #         cur_canvas = canvas.copy()
    #         Y = candidate[index.astype(int), 0]
    #         X = candidate[index.astype(int), 1]
    #         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(cur_canvas, polygon, colors[i])
    #         canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)

    print 'cv and print: %.5f, Model time: %.5f, Full processing time is %.5f.' % (
        time.time() - tic, model_running_time, time.time() - t0)
    return canvas
Exemple #14
0
def get_pose(image, gpu):
    T.set_num_threads(T.get_num_threads())
    weight_name = './model/pose_model.pth'
    (y, x, _) = image.shape
    blocks = {}

    # find connection in the specified sequence, center 29 is in the position 15
    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]]

    # the middle joints heatmap correpondence
    mapIdx = [[31,32], [39,40], [33,34], [35,36], [41,42], [43,44], [19,20], [21,22], \
        [23,24], [25,26], [27,28], [29,30], [47,48], [49,50], [53,54], [51,52], \
        [55,56], [37,38], [45,46]]

    # the joints between corresponding limbs
    joints = [
        [12, 0],
        [12, 1],
        [12, 6],
        [12, 9],
        [12, 13],
        [12, 15],  #connected to the Nose-Neck
        [1, 4],
        [4, 5],  #connect LShoulder-LElbow-LWrist
        [0, 2],
        [2, 3],  #connect RShoulder-RElbow-RWrist
        [6, 7],
        [7, 8],  #connect RHip-RKnee-RAnkle
        [9, 10],
        [10, 11],  #connect LHip-LKnee-LAnkle
        [13, 14],  #connect REye-REar
        [15, 16]  #connect LEye-LEar
    ]

    # visualize
    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]]

    block0 = [{
        'conv1_1': [3, 64, 3, 1, 1]
    }, {
        'conv1_2': [64, 64, 3, 1, 1]
    }, {
        'pool1_stage1': [2, 2, 0]
    }, {
        'conv2_1': [64, 128, 3, 1, 1]
    }, {
        'conv2_2': [128, 128, 3, 1, 1]
    }, {
        'pool2_stage1': [2, 2, 0]
    }, {
        'conv3_1': [128, 256, 3, 1, 1]
    }, {
        'conv3_2': [256, 256, 3, 1, 1]
    }, {
        'conv3_3': [256, 256, 3, 1, 1]
    }, {
        'conv3_4': [256, 256, 3, 1, 1]
    }, {
        'pool3_stage1': [2, 2, 0]
    }, {
        'conv4_1': [256, 512, 3, 1, 1]
    }, {
        'conv4_2': [512, 512, 3, 1, 1]
    }, {
        'conv4_3_CPM': [512, 256, 3, 1, 1]
    }, {
        'conv4_4_CPM': [256, 128, 3, 1, 1]
    }]

    blocks['block1_1'] = [{
        'conv5_1_CPM_L1': [128, 128, 3, 1, 1]
    }, {
        'conv5_2_CPM_L1': [128, 128, 3, 1, 1]
    }, {
        'conv5_3_CPM_L1': [128, 128, 3, 1, 1]
    }, {
        'conv5_4_CPM_L1': [128, 512, 1, 1, 0]
    }, {
        'conv5_5_CPM_L1': [512, 38, 1, 1, 0]
    }]

    blocks['block1_2'] = [{
        'conv5_1_CPM_L2': [128, 128, 3, 1, 1]
    }, {
        'conv5_2_CPM_L2': [128, 128, 3, 1, 1]
    }, {
        'conv5_3_CPM_L2': [128, 128, 3, 1, 1]
    }, {
        'conv5_4_CPM_L2': [128, 512, 1, 1, 0]
    }, {
        'conv5_5_CPM_L2': [512, 19, 1, 1, 0]
    }]

    for i in range(2, 7):
        blocks['block%d_1' % i] = [{
            'Mconv1_stage%d_L1' % i: [185, 128, 7, 1, 3]
        }, {
            'Mconv2_stage%d_L1' % i: [128, 128, 7, 1, 3]
        }, {
            'Mconv3_stage%d_L1' % i: [128, 128, 7, 1, 3]
        }, {
            'Mconv4_stage%d_L1' % i: [128, 128, 7, 1, 3]
        }, {
            'Mconv5_stage%d_L1' % i: [128, 128, 7, 1, 3]
        }, {
            'Mconv6_stage%d_L1' % i: [128, 128, 1, 1, 0]
        }, {
            'Mconv7_stage%d_L1' % i: [128, 38, 1, 1, 0]
        }]
        blocks['block%d_2' % i] = [{
            'Mconv1_stage%d_L2' % i: [185, 128, 7, 1, 3]
        }, {
            'Mconv2_stage%d_L2' % i: [128, 128, 7, 1, 3]
        }, {
            'Mconv3_stage%d_L2' % i: [128, 128, 7, 1, 3]
        }, {
            'Mconv4_stage%d_L2' % i: [128, 128, 7, 1, 3]
        }, {
            'Mconv5_stage%d_L2' % i: [128, 128, 7, 1, 3]
        }, {
            'Mconv6_stage%d_L2' % i: [128, 128, 1, 1, 0]
        }, {
            'Mconv7_stage%d_L2' % i: [128, 19, 1, 1, 0]
        }]

    def make_layers(cfg_dict):
        layers = []
        for i in range(len(cfg_dict) - 1):
            one_ = cfg_dict[i]
            for k, v in one_.items():
                if 'pool' in k:
                    layers += [
                        nn.MaxPool2d(kernel_size=v[0],
                                     stride=v[1],
                                     padding=v[2])
                    ]
                else:
                    conv2d = nn.Conv2d(in_channels=v[0],
                                       out_channels=v[1],
                                       kernel_size=v[2],
                                       stride=v[3],
                                       padding=v[4])
                    layers += [conv2d, nn.ReLU(inplace=True)]
        one_ = cfg_dict[-1].keys()
        k = list(one_)[0]
        v = cfg_dict[-1][k]
        conv2d = nn.Conv2d(in_channels=v[0],
                           out_channels=v[1],
                           kernel_size=v[2],
                           stride=v[3],
                           padding=v[4])
        layers += [conv2d]
        return nn.Sequential(*layers)

    layers = []
    for i in range(len(block0)):
        one_ = block0[i]
        #for k,v in one_.iteritems():
        for k, v in one_.items():
            if 'pool' in k:
                layers += [
                    nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2])
                ]
            else:
                conv2d = nn.Conv2d(in_channels=v[0],
                                   out_channels=v[1],
                                   kernel_size=v[2],
                                   stride=v[3],
                                   padding=v[4])
                layers += [conv2d, nn.ReLU(inplace=True)]

    models = {}
    models['block0'] = nn.Sequential(*layers)

    #for k,v in blocks.iteritems():
    for k, v in blocks.items():
        models[k] = make_layers(v)

    model = pose_model(models)
    model.load_state_dict(torch.load(weight_name))
    model.cuda(gpu)
    model.float()
    model.eval()

    param_, model_ = config_reader()

    #torch.nn.functional.pad(img pad, mode='constant', value=model_['padValue'])
    tic = time.time()
    oriImg = image
    imageToTest = Variable(T.transpose(
        T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(), 0), 2, 3), 1,
        2),
                           volatile=True).cuda(gpu)

    multiplier = [
        x * model_['boxsize'] / oriImg.shape[0] for x in param_['scale_search']
    ]

    heatmap_avg = torch.zeros(
        (len(multiplier), 19, oriImg.shape[0], oriImg.shape[1])).cuda(gpu)
    paf_avg = torch.zeros(
        (len(multiplier), 38, oriImg.shape[0], oriImg.shape[1])).cuda(gpu)
    #print(heatmap_avg.size())

    toc = time.time()
    #print('time to load model is %.5f'%(toc-tic))
    tic = time.time()
    for m in range(len(multiplier)):
        scale = multiplier[m]
        h = int(oriImg.shape[0] * scale)
        w = int(oriImg.shape[1] * scale)
        pad_h = 0 if (h % model_['stride']
                      == 0) else model_['stride'] - (h % model_['stride'])
        pad_w = 0 if (w % model_['stride']
                      == 0) else model_['stride'] - (w % model_['stride'])
        new_h = h + pad_h
        new_w = w + pad_w

        imageToTest = cv2.resize(oriImg, (0, 0),
                                 fx=scale,
                                 fy=scale,
                                 interpolation=cv2.INTER_CUBIC)
        imageToTest_padded, pad = util.padRightDownCorner(
            imageToTest, model_['stride'], model_['padValue'])
        imageToTest_padded = np.transpose(
            np.float32(imageToTest_padded[:, :, :, np.newaxis]),
            (3, 2, 0, 1)) / 256 - 0.5

        feed = Variable(T.from_numpy(imageToTest_padded)).cuda(gpu)
        output1, output2 = model(feed)
        heatmap = nn.UpsamplingBilinear2d(
            (oriImg.shape[0], oriImg.shape[1])).cuda(gpu)(output2)
        paf = nn.UpsamplingBilinear2d(
            (oriImg.shape[0], oriImg.shape[1])).cuda(gpu)(output1)

        heatmap_avg[m] = heatmap[0].data
        paf_avg[m] = paf[0].data

    toc = time.time()
    #print('time to forward pass is %.5f'%(toc-tic))
    tic = time.time()

    heatmap_avg = T.transpose(
        T.transpose(T.squeeze(T.mean(heatmap_avg, 0)), 0, 1), 1, 2)
    paf_avg = T.transpose(T.transpose(T.squeeze(T.mean(paf_avg, 0)), 0, 1), 1,
                          2)
    heatmap_avg = heatmap_avg.cpu().numpy()
    paf_avg = paf_avg.cpu().numpy()
    toc = time.time()
    #print('time to take averages is %.5f'%(toc-tic))
    tic = time.time()

    all_peaks = []
    peak_counter = 0

    #maps =
    for part in range(18):
        map_ori = heatmap_avg[:, :, part]
        map_ak = gaussian_filter(map_ori, sigma=3)

        map_left = np.zeros(map_ak.shape)
        map_left[1:, :] = map_ak[:-1, :]
        map_right = np.zeros(map_ak.shape)
        map_right[:-1, :] = map_ak[1:, :]
        map_up = np.zeros(map_ak.shape)
        map_up[:, 1:] = map_ak[:, :-1]
        map_down = np.zeros(map_ak.shape)
        map_down[:, :-1] = map_ak[:, 1:]

        peaks_binary = np.logical_and.reduce(
            (map_ak >= map_left, map_ak >= map_right, map_ak >= map_up,
             map_ak >= map_down, map_ak > param_['thre1']))
        #    peaks_binary = T.eq(
        #    peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])

        peaks = list(
            zip(np.nonzero(peaks_binary)[1],
                np.nonzero(peaks_binary)[0]))  # note reverse
        peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
        some_id = range(peak_counter, peak_counter + len(list(peaks)))
        peaks_with_score_and_id = [
            peaks_with_score[i] + (some_id[i], ) for i in range(len(some_id))
        ]

        all_peaks.append(peaks_with_score_and_id)
        peak_counter += len(peaks)

    connection_all = []
    special_k = []
    mid_num = 10

    for k in range(len(mapIdx)):
        score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
        candA = all_peaks[limbSeq[k][0] - 1]
        candB = all_peaks[limbSeq[k][1] - 1]
        nA = len(candA)
        nB = len(candB)
        indexA, indexB = limbSeq[k]
        if (nA != 0 and nB != 0):
            connection_candidate = []
            for i in range(nA):
                for j in range(nB):
                    vec = np.subtract(candB[j][:2], candA[i][:2])
                    norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
                    if norm != 0:
                        vec = np.divide(vec, norm)

                    startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
                          np.linspace(candA[i][1], candB[j][1], num=mid_num)))

                    vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
                          for I in range(len(startend))])
                    vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
                          for I in range(len(startend))])

                    score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
                        vec_y, vec[1])
                    if norm != 0:
                        score_with_dist_prior = sum(score_midpts) / len(
                            score_midpts) + min(
                                0.5 * oriImg.shape[0] / norm - 1, 0)
                    else:
                        score_with_dist_prior = sum(score_midpts) / len(
                            score_midpts) + min(0.5 * oriImg.shape[0] - 1, 0)
                    criterion1 = len(
                        np.nonzero(score_midpts > param_['thre2'])
                        [0]) > 0.8 * len(score_midpts)
                    criterion2 = score_with_dist_prior > 0
                    if criterion1 and criterion2:
                        connection_candidate.append([
                            i, j, score_with_dist_prior,
                            score_with_dist_prior + candA[i][2] + candB[j][2]
                        ])

            connection_candidate = sorted(connection_candidate,
                                          key=lambda x: x[2],
                                          reverse=True)
            connection = np.zeros((0, 5))
            for c in range(len(connection_candidate)):
                i, j, s = connection_candidate[c][0:3]
                if (i not in connection[:, 3] and j not in connection[:, 4]):
                    connection = np.vstack(
                        [connection, [candA[i][3], candB[j][3], s, i, j]])
                    if (len(connection) >= min(nA, nB)):
                        break

            connection_all.append(connection)
        else:
            special_k.append(k)
            connection_all.append([])

    # last number in each row is the total parts number of that person
    # the second last number in each row is the score of the overall configuration
    subset = -1 * np.ones((0, 20))
    candidate = np.array([item for sublist in all_peaks for item in sublist])

    for k in range(len(mapIdx)):
        if k not in special_k:
            partAs = connection_all[k][:, 0]
            partBs = connection_all[k][:, 1]
            indexA, indexB = np.array(limbSeq[k]) - 1

            for i in range(len(connection_all[k])):  #= 1:size(temp,1)
                found = 0
                subset_idx = [-1, -1]
                for j in range(len(subset)):  #1:size(subset,1):
                    if subset[j][indexA] == partAs[i] or subset[j][
                            indexB] == partBs[i]:
                        subset_idx[found] = j
                        found += 1

                if found == 1:
                    j = subset_idx[0]
                    if (subset[j][indexB] != partBs[i]):
                        subset[j][indexB] = partBs[i]
                        subset[j][-1] += 1
                        subset[j][-2] += candidate[partBs[i].astype(int),
                                                   2] + connection_all[k][i][2]
                    if (subset[j][indexA] != partAs[i]):
                        subset[j][indexA] = partAs[i]
                        subset[j][-1] += 1
                        subset[j][-2] += candidate[partAs[i].astype(int),
                                                   2] + connection_all[k][i][2]
                elif found == 2:  # if found 2 and disjoint, merge them
                    j1, j2 = subset_idx
                    #print("found = 2")
                    membership = ((subset[j1] >= 0).astype(int) +
                                  (subset[j2] >= 0).astype(int))[:-2]
                    if len(np.nonzero(membership == 2)[0]) == 0:  #merge
                        subset[j1][:-2] += (subset[j2][:-2] + 1)
                        subset[j1][-2:] += subset[j2][-2:]
                        subset[j1][-2] += connection_all[k][i][2]
                        subset = np.delete(subset, j2, 0)
                    else:  # as like found == 1
                        subset[j1][indexB] = partBs[i]
                        subset[j1][-1] += 1
                        subset[j1][-2] += candidate[
                            partBs[i].astype(int), 2] + connection_all[k][i][2]

                # if find no partA in the subset, create a new subset
                elif not found and k < 17:
                    row = -1 * np.ones(20)
                    row[indexA] = partAs[i]
                    row[indexB] = partBs[i]
                    row[-1] = 2
                    row[-2] = sum(
                        candidate[connection_all[k][i, :2].astype(int),
                                  2]) + connection_all[k][i][2]
                    subset = np.vstack([subset, row])

    # delete some rows of subset which has few parts occur
    deleteIdx = []
    for i in range(len(subset)):
        if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
            deleteIdx.append(i)
    subset = np.delete(subset, deleteIdx, axis=0)

    #canvas = cv2.imread(test_image) # B,G,R order
    # for i in range(18):
    #     for j in range(len(all_peaks[i])):
    #         cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
    results = []
    stickwidth = 5
    for n in range(len(subset)):
        limbs = {}
        angles = []
        position = []
        for i in range(17):
            index = subset[n][np.array(limbSeq[i]) - 1]
            if -1 in index:
                continue
            # cur_canvas = canvas.copy()
            Y = candidate[index.astype(int), 0]
            X = candidate[index.astype(int), 1]

            limb = ((X[0] - X[1], Y[0] - Y[1]))
            limbs[tuple(limbSeq[i])] = limb
            #if i == 12:
            #position.append(np.mean(X)/x)
            #position.append(np.mean(Y)/y)
            # 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(cur_canvas, polygon, colors[i])
            # canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
        results.append(extract.get_joints(limbs, joints, limbSeq))

    #Parallel(n_jobs=1)(delayed(handle_one)(i) for i in range(18))
    toc = time.time()
    #torch.cuda.empty_cache()
    #print('time to feature extract is %.5f'%(toc-tic))
    return (results)
    map = gaussian_filter(map_ori, sigma=3)
    
    map_left = np.zeros(map.shape)
    map_left[1:,:] = map[:-1,:]
    map_right = np.zeros(map.shape)
    map_right[:-1,:] = map[1:,:]
    map_up = np.zeros(map.shape)
    map_up[:,1:] = map[:,:-1]
    map_down = np.zeros(map.shape)
    map_down[:,:-1] = map[:,1:]
    
    peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param_['thre1']))
#    peaks_binary = T.eq(
#    peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
    
    peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse
    
    peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]
    id = range(peak_counter, peak_counter + len(peaks))
    peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))]

    all_peaks.append(peaks_with_score_and_id)
    peak_counter += len(peaks)
    
    
    
    
    
connection_all = []
special_k = []
mid_num = 10