| import os | |
| os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE" | |
| import torch | |
| import numpy as np | |
| from . import util | |
| from .body import Body | |
| from .hand import Hand | |
| from huggingface_hub import hf_hub_url, cached_download | |
| REPO_ID = "lllyasviel/ControlNet" | |
| body_estimation = Body(cached_download(hf_hub_url(REPO_ID, 'annotator/ckpts/body_pose_model.pth'))) | |
| hand_estimation = Hand(cached_download(hf_hub_url(REPO_ID,'annotator/ckpts/hand_pose_model.pth'))) | |
| def apply_openpose(oriImg, hand=False): | |
| oriImg = oriImg[:, :, ::-1].copy() | |
| with torch.no_grad(): | |
| candidate, subset = body_estimation(oriImg) | |
| canvas = np.zeros_like(oriImg) | |
| canvas = util.draw_bodypose(canvas, candidate, subset) | |
| if hand: | |
| hands_list = util.handDetect(candidate, subset, oriImg) | |
| all_hand_peaks = [] | |
| for x, y, w, is_left in hands_list: | |
| peaks = hand_estimation(oriImg[y:y+w, x:x+w, :]) | |
| peaks[:, 0] = np.where(peaks[:, 0] == 0, peaks[:, 0], peaks[:, 0] + x) | |
| peaks[:, 1] = np.where(peaks[:, 1] == 0, peaks[:, 1], peaks[:, 1] + y) | |
| all_hand_peaks.append(peaks) | |
| canvas = util.draw_handpose(canvas, all_hand_peaks) | |
| return canvas, dict(candidate=candidate.tolist(), subset=subset.tolist()) | |