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
    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])
                    # MODIFIED.
                    try:
                        score_with_dist_prior = sum(score_midpts) / len(
                            score_midpts) + min(
Esempio n. 4
0
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])
                    if (norm == 0 or len(score_midpts) == 0):
                        continue
                    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)

    # for i in range 17:
    #     for j in range(len(all_peaks[10])):

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

    # heat_map = heat_map_circle(heat_map, center)
    # cv2.circle(canvas, center, 4, colors[0], thickness=-1)
    # cv2.circle(canvas, all_peaks[10][j][0:2], 4, colors[2], thickness=-1)
    #
    # cv2.circle(canvas, all_peaks[13][j][0:2], 4, colors[3], thickness=-1)
    stickwidth = 4
    track_list = []
    # return canvas, heat_map
    # while True:
    #     bbox = cv2.selectROI("tracker", canvas)
    #     print bbox

    for n in range(len(subset)):
        person_points = []
        for i in range(19):
            # for i in [7,8,10,11]:
            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]
            mX = np.mean(X)
            mY = np.mean(Y)
            cur_canvas = canvas.copy()
            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)
            person_points.extend([(Y[0], X[0]), (Y[1], X[1])])
        bounding_box = find_person_bounding_box(person_points)
        track_list.append(tracker(TRACK_TYPE, bounding_box, canvas))
    # for i in range(19):
    #     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, track_list
Esempio n. 5
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)
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] +
                                     0.0000000000001)  # +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
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]])