otmanheddouch/house_design

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 6import os7os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"8 9import torch10import numpy as np11from . import util12from .body import Body13from .hand import Hand14from annotator.util import annotator_ckpts_path15 16 17body_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/body_pose_model.pth"18hand_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/hand_pose_model.pth"19 20 21class OpenposeDetector:22    def __init__(self):23        body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth")24        hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth")25 26        if not os.path.exists(hand_modelpath):27            from basicsr.utils.download_util import load_file_from_url28            load_file_from_url(body_model_path, model_dir=annotator_ckpts_path)29            load_file_from_url(hand_model_path, model_dir=annotator_ckpts_path)30 31        self.body_estimation = Body(body_modelpath)32        self.hand_estimation = Hand(hand_modelpath)33 34    def __call__(self, oriImg, hand=False):35        oriImg = oriImg[:, :, ::-1].copy()36        with torch.no_grad():37            candidate, subset = self.body_estimation(oriImg)38            canvas = np.zeros_like(oriImg)39            canvas = util.draw_bodypose(canvas, candidate, subset)40            if hand:41                hands_list = util.handDetect(candidate, subset, oriImg)42                all_hand_peaks = []43                for x, y, w, is_left in hands_list:44                    peaks = self.hand_estimation(oriImg[y:y+w, x:x+w, :])45                    peaks[:, 0] = np.where(peaks[:, 0] == 0, peaks[:, 0], peaks[:, 0] + x)46                    peaks[:, 1] = np.where(peaks[:, 1] == 0, peaks[:, 1], peaks[:, 1] + y)47                    all_hand_peaks.append(peaks)48                canvas = util.draw_handpose(canvas, all_hand_peaks)49            return canvas, dict(candidate=candidate.tolist(), subset=subset.tolist())50