| 12 | |
| 13 | |
| 14 | class general_pose_model(object): |
| 15 | def __init__(self, modelpath): |
| 16 | self.num_points = 22 |
| 17 | self.point_pairs = [[0,1],[1,2],[2,3],[3,4], |
| 18 | [0,5],[5,6],[6,7],[7,8], |
| 19 | [0,9],[9,10],[10,11],[11,12], |
| 20 | [0,13],[13,14],[14,15],[15,16], |
| 21 | [0,17],[17,18],[18,19],[19,20]] |
| 22 | # self.inWidth = 368 |
| 23 | self.inHeight = 368 |
| 24 | self.threshold = 0.1 |
| 25 | self.hand_net = self.get_hand_model(modelpath) |
| 26 | |
| 27 | |
| 28 | def get_hand_model(self, modelpath): |
| 29 | |
| 30 | prototxt = os.path.join(modelpath, "pose_deploy.prototxt") |
| 31 | caffemodel = os.path.join(modelpath, "../pose_iter_102000.caffemodel") |
| 32 | hand_model = cv2.dnn.readNetFromCaffe(prototxt, caffemodel) |
| 33 | |
| 34 | return hand_model |
| 35 | |
| 36 | |
| 37 | def predict(self, imgfile): |
| 38 | img_cv2 = cv2.imread(imgfile) |
| 39 | img_height, img_width, _ = img_cv2.shape |
| 40 | aspect_ratio = img_width / img_height |
| 41 | |
| 42 | inWidth = int(((aspect_ratio * self.inHeight) * 8) // 8) |
| 43 | inpBlob = cv2.dnn.blobFromImage(img_cv2, 1.0 / 255, (inWidth, self.inHeight), (0, 0, 0), swapRB=False, crop=False) |
| 44 | |
| 45 | self.hand_net.setInput(inpBlob) |
| 46 | |
| 47 | output = self.hand_net.forward() |
| 48 | |
| 49 | # vis heatmaps |
| 50 | self.vis_heatmaps(imgfile, output) |
| 51 | |
| 52 | # |
| 53 | points = [] |
| 54 | for idx in range(self.num_points): |
| 55 | probMap = output[0, idx, :, :] # confidence map. |
| 56 | probMap = cv2.resize(probMap, (img_width, img_height)) |
| 57 | |
| 58 | # Find global maxima of the probMap. |
| 59 | minVal, prob, minLoc, point = cv2.minMaxLoc(probMap) |
| 60 | |
| 61 | if prob > self.threshold: |
| 62 | points.append((int(point[0]), int(point[1]))) |
| 63 | else: |
| 64 | points.append(None) |
| 65 | |
| 66 | return points |
| 67 | |
| 68 | |
| 69 | def vis_heatmaps(self, imgfile, net_outputs): |
| 70 | img_cv2 = cv2.imread(imgfile) |
| 71 | plt.figure(figsize=[10, 10]) |