Spaces:
Runtime error
Runtime error
| 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)) | |