multi-person-openpose.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
import cv2
import time
import numpy as np
from random import randint
import argparse

parser = argparse.ArgumentParser(description='Run keypoint detection')
parser.add_argument('--device', default='cpu',
                    help='Device to inference on')
parser.add_argument('--image_file', default='group.jpg',
                    help='Input image')

args = parser.parse_args()

image1 = cv2.imread(args.image_file)

protoFile = 'pose/coco/pose_deploy_linevec.prototxt'
weightsFile = 'pose/coco/pose_iter_440000.caffemodel'
nPoints = 18

# COCO Output Format

keypointsMapping = [
    'Nose',
    'Neck',
    'R-Sho',
    'R-Elb',
    'R-Wr',
    'L-Sho',
    'L-Elb',
    'L-Wr',
    'R-Hip',
    'R-Knee',
    'R-Ank',
    'L-Hip',
    'L-Knee',
    'L-Ank',
    'R-Eye',
    'L-Eye',
    'R-Ear',
    'L-Ear',
    ]

POSE_PAIRS = [
    [1, 2],
    [1, 5],
    [2, 3],
    [3, 4],
    [5, 6],
    [6, 7],
    [1, 8],
    [8, 9],
    [9, 10],
    [1, 11],
    [11, 12],
    [12, 13],
    [1, 0],
    [0, 14],
    [14, 16],
    [0, 15],
    [15, 17],
    [2, 17],
    [5, 16],
    ]


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],
    ]

colors = [
    [0, 100, 255],
    [0, 100, 255],
    [0, 255, 255],
    [0, 100, 255],
    [0, 255, 255],
    [0, 100, 255],
    [0, 255, 0],
    [255, 200, 100],
    [255, 0, 255],
    [0, 255, 0],
    [255, 200, 100],
    [255, 0, 255],
    [0, 0, 255],
    [255, 0, 0],
    [200, 200, 0],
    [255, 0, 0],
    [200, 200, 0],
    [0, 0, 0],
    ]


def getKeypoints(probMap, threshold=0.1):

    mapSmooth = cv2.GaussianBlur(probMap, (3, 3), 0, 0)

    mapMask = np.uint8(mapSmooth > threshold)
    keypoints = []


    (contours, _) = cv2.findContours(mapMask, cv2.RETR_TREE,
            cv2.CHAIN_APPROX_SIMPLE)


    for cnt in contours:
        blobMask = np.zeros(mapMask.shape)
        blobMask = cv2.fillConvexPoly(blobMask, cnt, 1)
        maskedProbMap = mapSmooth * blobMask
        (_, maxVal, _, maxLoc) = cv2.minMaxLoc(maskedProbMap)
        keypoints.append(maxLoc + (probMap[maxLoc[1], maxLoc[0]], ))

    return keypoints



def getValidPairs(output):
    valid_pairs = []
    invalid_pairs = []
    n_interp_samples = 10
    paf_score_th = 0.1
    conf_th = 0.7


    for k in range(len(mapIdx)):


        pafA = output[0, mapIdx[k][0], :, :]
        pafB = output[0, mapIdx[k][1], :, :]
        pafA = cv2.resize(pafA, (frameWidth, frameHeight))
        pafB = cv2.resize(pafB, (frameWidth, frameHeight))


        candA = detected_keypoints[POSE_PAIRS[k][0]]
        candB = detected_keypoints[POSE_PAIRS[k][1]]
        nA = len(candA)
        nB = len(candB)


        if nA != 0 and nB != 0:
            valid_pair = np.zeros((0, 3))
            for i in range(nA):
                max_j = -1
                maxScore = -1
                found = 0
                for j in range(nB):


                    d_ij = np.subtract((candB[j])[:2], (candA[i])[:2])
                    norm = np.linalg.norm(d_ij)
                    if norm:
                        d_ij = d_ij / norm
                    else:
                        continue


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


                    paf_interp = []
                    for k in range(len(interp_coord)):
                        paf_interp.append([pafA[int(round(interp_coord[k][1])),
                                int(round(interp_coord[k][0]))],
                                pafB[int(round(interp_coord[k][1])),
                                int(round(interp_coord[k][0]))]])


                    paf_scores = np.dot(paf_interp, d_ij)
                    avg_paf_score = sum(paf_scores) / len(paf_scores)


                    if len(np.where(paf_scores > paf_score_th)[0]) \
                        / n_interp_samples > conf_th:
                        if avg_paf_score > maxScore:
                            max_j = j
                            maxScore = avg_paf_score
                            found = 1


                if found:
                    valid_pair = np.append(valid_pair, [[candA[i][3],
                            candB[max_j][3], maxScore]], axis=0)


            valid_pairs.append(valid_pair)
        else:


            print('No Connection : k = {}'.format(k))
            invalid_pairs.append(k)
            valid_pairs.append([])
    return (valid_pairs, invalid_pairs)



def getPersonwiseKeypoints(valid_pairs, invalid_pairs):


    personwiseKeypoints = -1 * np.ones((0, 19))

    for k in range(len(mapIdx)):
        if k not in invalid_pairs:
            partAs = valid_pairs[k][:, 0]
            partBs = valid_pairs[k][:, 1]
            (indexA, indexB) = np.array(POSE_PAIRS[k])

            for i in range(len(valid_pairs[k])):
                found = 0
                person_idx = -1
                for j in range(len(personwiseKeypoints)):
                    if personwiseKeypoints[j][indexA] == partAs[i]:
                        person_idx = j
                        found = 1
                        break

                if found:
                    personwiseKeypoints[person_idx][indexB] = partBs[i]
                    personwiseKeypoints[person_idx][-1] += \
                        keypoints_list[partBs[i].astype(int), 2] \
                        + valid_pairs[k][i][2]
                elif not found and k < 17:


                    row = -1 * np.ones(19)
                    row[indexA] = partAs[i]
                    row[indexB] = partBs[i]


                    row[-1] = sum(keypoints_list[valid_pairs[k][i, :
                                  2].astype(int), 2]) \
                        + valid_pairs[k][i][2]
                    personwiseKeypoints = \
                        np.vstack([personwiseKeypoints, row])
    return personwiseKeypoints


frameWidth = image1.shape[1]
frameHeight = image1.shape[0]

t = time.time()
net = cv2.dnn.readNetFromCaffe(protoFile, weightsFile)
if args.device == 'cpu':
    net.setPreferableBackend(cv2.dnn.DNN_TARGET_CPU)
    print('Using CPU device')
elif args.device == 'gpu':
    net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
    net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
    print('Using GPU device')


inHeight = 368
inWidth = int(inHeight / frameHeight * frameWidth)

inpBlob = cv2.dnn.blobFromImage(
    image1,
    1.0 / 255,
    (inWidth, inHeight),
    (0, 0, 0),
    swapRB=False,
    crop=False,
    )

net.setInput(inpBlob)
output = net.forward()
print('Time Taken in forward pass = {}'.format(time.time() - t))

detected_keypoints = []
keypoints_list = np.zeros((0, 3))
keypoint_id = 0
threshold = 0.1

for part in range(nPoints):
    probMap = output[0, part, :, :]
    probMap = cv2.resize(probMap, (image1.shape[1], image1.shape[0]))
    keypoints = getKeypoints(probMap, threshold)
    print ('Keypoints - {} : {}'.format(keypointsMapping[part],keypoints))
    keypoints_with_id = []
    for i in range(len(keypoints)):
        keypoints_with_id.append(keypoints[i] + (keypoint_id, ))
        keypoints_list = np.vstack([keypoints_list, keypoints[i]])
        keypoint_id += 1

    detected_keypoints.append(keypoints_with_id)

frameClone = image1.copy()
for i in range(nPoints):
    for j in range(len(detected_keypoints[i])):
        cv2.circle(
            frameClone,
            (detected_keypoints[i][j])[0:2],
            5,
            colors[i],
            -1,
            cv2.LINE_AA,
            )
cv2.imshow('Keypoints', frameClone)

(valid_pairs, invalid_pairs) = getValidPairs(output)
personwiseKeypoints = getPersonwiseKeypoints(valid_pairs, invalid_pairs)

for i in range(17):
    for n in range(len(personwiseKeypoints)):
        index = personwiseKeypoints[n][np.array(POSE_PAIRS[i])]
        if -1 in index:
            continue
        B = np.int32(keypoints_list[index.astype(int), 0])
        A = np.int32(keypoints_list[index.astype(int), 1])
        cv2.line(
            frameClone,
            (B[0], A[0]),
            (B[1], A[1]),
            colors[i],
            3,
            cv2.LINE_AA,
            )

cv2.imshow('Detected Pose', frameClone)
cv2.waitKey(0)