Spaces:
Runtime error
Runtime error
File size: 3,599 Bytes
4a285f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import cv2
import numpy as np
import os
import json
class general_pose_model(object):
def __init__(self, modelpath):
# Specify the model to be used
# Body25: 25 points
# COCO: 18 points
# MPI: 15 points
self.inWidth = 368
self.inHeight = 368
self.threshold = 0.05
self.pose_net = self.general_coco_model(modelpath)
def general_coco_model(self, modelpath):
self.points_name = {
"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}
self.num_points = 18
self.point_pairs = [[1, 0], [1, 2], [1, 5],
[2, 3], [3, 4], [5, 6],
[6, 7], [1, 8], [8, 9],
[9, 10], [1, 11], [11, 12],
[12, 13], [0, 14], [0, 15],
[14, 16], [15, 17]]
prototxt = os.path.join(
modelpath,
'pose_deploy_linevec.prototxt')
caffemodel = os.path.join(
modelpath,
'pose_iter_440000.caffemodel')
coco_model = cv2.dnn.readNetFromCaffe(prototxt, caffemodel)
return coco_model
def predict(self, imgfile):
img_cv2 = cv2.imread(imgfile)
img_height, img_width, _ = img_cv2.shape
inpBlob = cv2.dnn.blobFromImage(img_cv2,
1.0 / 255,
(self.inWidth, self.inHeight),
(0, 0, 0),
swapRB=False,
crop=False)
self.pose_net.setInput(inpBlob)
self.pose_net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
self.pose_net.setPreferableTarget(cv2.dnn.DNN_TARGET_OPENCL)
output = self.pose_net.forward()
H = output.shape[2]
W = output.shape[3]
points = []
for idx in range(self.num_points):
probMap = output[0, idx, :, :] # confidence map.
# Find global maxima of the probMap.
minVal, prob, minLoc, point = cv2.minMaxLoc(probMap)
# Scale the point to fit on the original image
x = (img_width * point[0]) / W
y = (img_height * point[1]) / H
if prob > self.threshold:
points.append(x)
points.append(y)
points.append(prob)
else:
points.append(0)
points.append(0)
points.append(0)
return points
def generate_pose_keypoints(img_file, pose_file):
modelpath = 'pose'
pose_model = general_pose_model(modelpath)
res_points = pose_model.predict(img_file)
pose_data = {"version": 1,
"people": [
{"pose_keypoints": res_points}
]
}
pose_keypoints_path = pose_file
json_object = json.dumps(pose_data, indent = 4)
# Writing to sample.json
with open(pose_keypoints_path, "w") as outfile:
outfile.write(json_object)
print('File saved at {}'.format(pose_keypoints_path))
|