| 14 | |
| 15 | |
| 16 | class OpenposeDetector: |
| 17 | def __init__(self): |
| 18 | body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth") |
| 19 | hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth") |
| 20 | |
| 21 | if not os.path.exists(hand_modelpath): |
| 22 | from basicsr.utils.download_util import load_file_from_url |
| 23 | load_file_from_url(body_model_path, model_dir=annotator_ckpts_path) |
| 24 | load_file_from_url(hand_model_path, model_dir=annotator_ckpts_path) |
| 25 | |
| 26 | self.body_estimation = Body(body_modelpath) |
| 27 | self.hand_estimation = Hand(hand_modelpath) |
| 28 | |
| 29 | def __call__(self, oriImg, hand=False): |
| 30 | oriImg = oriImg[:, :, ::-1].copy() |
| 31 | with torch.no_grad(): |
| 32 | candidate, subset = self.body_estimation(oriImg) |
| 33 | canvas = np.zeros_like(oriImg) |
| 34 | canvas = util.draw_bodypose(canvas, candidate, subset) |
| 35 | if hand: |
| 36 | hands_list = util.handDetect(candidate, subset, oriImg) |
| 37 | all_hand_peaks = [] |
| 38 | for x, y, w, is_left in hands_list: |
| 39 | peaks = self.hand_estimation(oriImg[y:y+w, x:x+w, :]) |
| 40 | peaks[:, 0] = np.where(peaks[:, 0] == 0, peaks[:, 0], peaks[:, 0] + x) |
| 41 | peaks[:, 1] = np.where(peaks[:, 1] == 0, peaks[:, 1], peaks[:, 1] + y) |
| 42 | all_hand_peaks.append(peaks) |
| 43 | canvas = util.draw_handpose(canvas, all_hand_peaks) |
| 44 | return canvas, dict(candidate=candidate.tolist(), subset=subset.tolist()) |