coreml-community/ControlNet-v1-1-Annotators-cpu

1# Openpose2# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose3# 2nd Edited by https://github.com/Hzzone/pytorch-openpose4# 3rd Edited by ControlNet5# 4th Edited by ControlNet (added face and correct hands)6 7import os8os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"9 10import torch11import numpy as np12from . import util13from .body import Body14from .hand import Hand15from .face import Face16from annotator.util import annotator_ckpts_path17 18 19body_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/body_pose_model.pth"20hand_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/hand_pose_model.pth"21face_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/facenet.pth"22 23 24def draw_pose(pose, H, W, draw_body=True, draw_hand=True, draw_face=True):25    bodies = pose['bodies']26    faces = pose['faces']27    hands = pose['hands']28    candidate = bodies['candidate']29    subset = bodies['subset']30    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)31 32    if draw_body:33        canvas = util.draw_bodypose(canvas, candidate, subset)34 35    if draw_hand:36        canvas = util.draw_handpose(canvas, hands)37 38    if draw_face:39        canvas = util.draw_facepose(canvas, faces)40 41    return canvas42 43 44class OpenposeDetector:45    def __init__(self):46        body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth")47        hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth")48        face_modelpath = os.path.join(annotator_ckpts_path, "facenet.pth")49 50        if not os.path.exists(body_modelpath):51            from basicsr.utils.download_util import load_file_from_url52            load_file_from_url(body_model_path, model_dir=annotator_ckpts_path)53 54        if not os.path.exists(hand_modelpath):55            from basicsr.utils.download_util import load_file_from_url56            load_file_from_url(hand_model_path, model_dir=annotator_ckpts_path)57 58        if not os.path.exists(face_modelpath):59            from basicsr.utils.download_util import load_file_from_url60            load_file_from_url(face_model_path, model_dir=annotator_ckpts_path)61 62        self.body_estimation = Body(body_modelpath)63        self.hand_estimation = Hand(hand_modelpath)64        self.face_estimation = Face(face_modelpath)65 66    def __call__(self, oriImg, hand_and_face=False, return_is_index=False):67        oriImg = oriImg[:, :, ::-1].copy()68        H, W, C = oriImg.shape69        with torch.no_grad():70            candidate, subset = self.body_estimation(oriImg)71            hands = []72            faces = []73            if hand_and_face:74                # Hand75                hands_list = util.handDetect(candidate, subset, oriImg)76                for x, y, w, is_left in hands_list:77                    peaks = self.hand_estimation(oriImg[y:y+w, x:x+w, :]).astype(np.float32)78                    if peaks.ndim == 2 and peaks.shape[1] == 2:79                        peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)80                        peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)81                        hands.append(peaks.tolist())82                # Face83                faces_list = util.faceDetect(candidate, subset, oriImg)84                for x, y, w in faces_list:85                    heatmaps = self.face_estimation(oriImg[y:y+w, x:x+w, :])86                    peaks = self.face_estimation.compute_peaks_from_heatmaps(heatmaps).astype(np.float32)87                    if peaks.ndim == 2 and peaks.shape[1] == 2:88                        peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)89                        peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)90                        faces.append(peaks.tolist())91            if candidate.ndim == 2 and candidate.shape[1] == 4:92                candidate = candidate[:, :2]93                candidate[:, 0] /= float(W)94                candidate[:, 1] /= float(H)95            bodies = dict(candidate=candidate.tolist(), subset=subset.tolist())96            pose = dict(bodies=bodies, hands=hands, faces=faces)97            if return_is_index:98                return pose99            else:100                return draw_pose(pose, H, W)101