Last active
March 20, 2018 14:01
-
-
Save berak/8e99f75fac7712443d8efdd1047a5c9e to your computer and use it in GitHub Desktop.
hacked opoenpose.py for tf model
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# cv2.__version__ -> 3.4.0 (from PIP) | |
# | |
# python C:\p\opencv\samples\dnn\openpose.py --model C:\data\mdl\body\graph_opt.pb --input C:\data\img\persons\single2.png --dataset=COCO | |
# | |
# To use Inference Engine backend, specify location of plugins: | |
# export LD_LIBRARY_PATH=/opt/intel/deeplearning_deploymenttoolkit/deployment_tools/external/mklml_lnx/lib:$LD_LIBRARY_PATH | |
import cv2 as cv | |
import numpy as np | |
import argparse | |
parser = argparse.ArgumentParser( | |
description='This script is used to demonstrate OpenPose human pose estimation network ' | |
'from https://github.com/CMU-Perceptual-Computing-Lab/openpose project using OpenCV. ' | |
'The sample and model are simplified and could be used for a single person on the frame.') | |
parser.add_argument('--input', help='Path to image or video. Skip to capture frames from camera') | |
parser.add_argument('--proto', help='Path to .prototxt') | |
parser.add_argument('--model', help='Path to .caffemodel') | |
parser.add_argument('--dataset', help='Specify what kind of model was trained. ' | |
'It could be (COCO, MPI) depends on dataset.') | |
parser.add_argument('--thr', default=0.1, type=float, help='Threshold value for pose parts heat map') | |
parser.add_argument('--width', default=368, type=int, help='Resize input to specific width.') | |
parser.add_argument('--height', default=368, type=int, help='Resize input to specific height.') | |
parser.add_argument('--inf_engine', action='store_true', | |
help='Enable Intel Inference Engine computational backend. ' | |
'Check that plugins folder is in LD_LIBRARY_PATH environment variable') | |
args = parser.parse_args() | |
if args.dataset == 'COCO': | |
BODY_PARTS = { "Nose": 0, "Neck": 1, "RShoulder": 2, "RElbow": 3, "RWrist": 4, | |
"LShoulder": 5, "LElbow": 6, "LWrist": 7, "RHip": 8, "RKnee": 9, | |
"RAnkle": 10, "LHip": 11, "LKnee": 12, "LAnkle": 13, "REye": 14, | |
"LEye": 15, "REar": 16, "LEar": 17, "Background": 18 } | |
POSE_PAIRS = [ ["Neck", "RShoulder"], ["Neck", "LShoulder"], ["RShoulder", "RElbow"], | |
["RElbow", "RWrist"], ["LShoulder", "LElbow"], ["LElbow", "LWrist"], | |
["Neck", "RHip"], ["RHip", "RKnee"], ["RKnee", "RAnkle"], ["Neck", "LHip"], | |
["LHip", "LKnee"], ["LKnee", "LAnkle"], ["Neck", "Nose"], ["Nose", "REye"], | |
["REye", "REar"], ["Nose", "LEye"], ["LEye", "LEar"] ] | |
else: | |
assert(args.dataset == 'MPI') | |
BODY_PARTS = { "Head": 0, "Neck": 1, "RShoulder": 2, "RElbow": 3, "RWrist": 4, | |
"LShoulder": 5, "LElbow": 6, "LWrist": 7, "RHip": 8, "RKnee": 9, | |
"RAnkle": 10, "LHip": 11, "LKnee": 12, "LAnkle": 13, "Chest": 14, | |
"Background": 15 } | |
POSE_PAIRS = [ ["Head", "Neck"], ["Neck", "RShoulder"], ["RShoulder", "RElbow"], | |
["RElbow", "RWrist"], ["Neck", "LShoulder"], ["LShoulder", "LElbow"], | |
["LElbow", "LWrist"], ["Neck", "Chest"], ["Chest", "RHip"], ["RHip", "RKnee"], | |
["RKnee", "RAnkle"], ["Chest", "LHip"], ["LHip", "LKnee"], ["LKnee", "LAnkle"] ] | |
inWidth = args.width | |
inHeight = args.height | |
#net = cv.dnn.readNetFromCaffe(args.proto, args.model) | |
net = cv.dnn.readNetFromTensorflow(args.model) | |
if args.inf_engine: | |
net.setPreferableBackend(cv.dnn.DNN_BACKEND_INFERENCE_ENGINE) | |
cap = cv.VideoCapture(args.input if args.input else 0) | |
while cv.waitKey(1) < 0: | |
hasFrame, frame = cap.read() | |
if not hasFrame: | |
cv.waitKey() | |
break | |
frameWidth = frame.shape[1] | |
frameHeight = frame.shape[0] | |
inp = cv.dnn.blobFromImage(frame, 1.0 , (inWidth, inHeight), | |
(0, 0, 0), swapRB=False, crop=False) | |
net.setInput(inp) | |
out = net.forward() | |
#assert(len(BODY_PARTS) == out.shape[1]) | |
points = [] | |
for i in range(len(BODY_PARTS)): | |
# Slice heatmap of corresponging body's part. | |
heatMap = out[0, i, :, :] | |
# Originally, we try to find all the local maximums. To simplify a sample | |
# we just find a global one. However only a single pose at the same time | |
# could be detected this way. | |
_, conf, _, point = cv.minMaxLoc(heatMap) | |
x = (frameWidth * point[0]) / out.shape[3] | |
y = (frameHeight * point[1]) / out.shape[2] | |
# Add a point if it's confidence is higher than threshold. | |
points.append((int(x), int(y)) if conf > args.thr else None) | |
for pair in POSE_PAIRS: | |
partFrom = pair[0] | |
partTo = pair[1] | |
assert(partFrom in BODY_PARTS) | |
assert(partTo in BODY_PARTS) | |
idFrom = BODY_PARTS[partFrom] | |
idTo = BODY_PARTS[partTo] | |
if points[idFrom] and points[idTo]: | |
cv.line(frame, points[idFrom], points[idTo], (0, 255, 0), 3) | |
cv.ellipse(frame, points[idFrom], (3, 3), 0, 0, 360, (0, 0, 255), cv.FILLED) | |
cv.ellipse(frame, points[idTo], (3, 3), 0, 0, 360, (0, 0, 255), cv.FILLED) | |
t, _ = net.getPerfProfile() | |
freq = cv.getTickFrequency() / 1000 | |
cv.putText(frame, '%.2fms' % (t / freq), (10, 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0)) | |
cv.imshow('OpenPose using OpenCV', frame) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment