File size: 2,268 Bytes
4edb0a5 84707de 4edb0a5 | 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 | 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 annotator.util import annotator_ckpts_path
body_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/body_pose_model.pth"
hand_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/hand_pose_model.pth"
class OpenposeDetector:
def __init__(self):
body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth")
hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth")
if not os.path.exists(annotator_ckpts_path):
os.makedirs(annotator_ckpts_path)
if not os.path.exists(body_modelpath):
from torch.hub import download_url_to_file
print(f"Downloading body_pose_model from {body_model_path}...")
download_url_to_file(body_model_path, body_modelpath)
if not os.path.exists(hand_modelpath):
from torch.hub import download_url_to_file
print(f"Downloading hand_pose_model from {hand_model_path}...")
download_url_to_file(hand_model_path, hand_modelpath)
self.body_estimation = Body(body_modelpath)
self.hand_estimation = Hand(hand_modelpath)
def __call__(self, oriImg, hand=False):
oriImg = oriImg[:, :, ::-1].copy()
with torch.no_grad():
candidate, subset = self.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 = self.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())
|