Spaces:
Running on Zero
Running on Zero
| # https://github.com/IDEA-Research/DWPose | |
| # Openpose | |
| # Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose | |
| # 2nd Edited by https://github.com/Hzzone/pytorch-openpose | |
| # 3rd Edited by ControlNet | |
| # 4th Edited by ControlNet (added face and correct hands) | |
| import copy | |
| import os | |
| os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from controlnet_aux.util import HWC3, resize_image | |
| from PIL import Image | |
| from . import util | |
| from .wholebody import Wholebody | |
| class DWposeDetector: | |
| def __init__(self, use_batch=False): | |
| self.use_batch = use_batch | |
| pass | |
| def to(self, device): | |
| self.pose_estimation = Wholebody(device, self.use_batch) | |
| return self | |
| def _get_multi_result_from_est(self, candidate, score_result, det_result, H, W): | |
| nums, keys, locs = candidate.shape # n 所有身体关键点数量，坐标 | |
| candidate[..., 0] /= float(W) | |
| candidate[..., 1] /= float(H) | |
| subset_score = score_result[:, :24] # 按照24个骨骼关键点来区分可见位置 | |
| face_score = score_result[:, 24:92] | |
| hand_score = score_result[:, 92:113] | |
| hand_score = np.vstack([hand_score, score_result[:, 113:]]) | |
| body_candidate = candidate[:, :24].copy() # body(n, 24, 2) | |
| for i in range(len(subset_score)): # n 个 | |
| for j in range(len(subset_score[i])): | |
| if subset_score[i][j] > 0.3: | |
| subset_score[i][j] = j # 标注序号，这样后续用的时候可以快速查出可用点 | |
| else: | |
| subset_score[i][j] = -1 # 躯干中去除掉不可见的骨骼 | |
| un_visible = score_result < 0.3 | |
| candidate[un_visible] = -1 # 全部关键点中去掉不可见骨骼 | |
| faces = candidate[:, 24:92] | |
| hands = candidate[:, 92:113] # hands(2*n, 21, 2) | |
| hands = np.vstack([hands, candidate[:, 113:]]) | |
| bodies = dict(candidate=body_candidate, subset=subset_score) | |
| pose = dict(bodies=bodies, hands=hands, faces=faces) | |
| score = dict(body_score=subset_score, hand_score=hand_score, face_score=face_score) | |
| new_det_result = [] | |
| for bbox in det_result: | |
| x1, y1, x2, y2 = bbox | |
| new_x1 = x1 / W | |
| new_y1 = y1 / H | |
| new_x2 = x2 / W | |
| new_y2 = y2 / H | |
| new_bbox = [new_x1, new_y1, new_x2, new_y2] | |
| new_det_result.append(new_bbox) | |
| return pose, score, new_det_result # body_score是原始的躯干骨骼分数 | |
| # def _get_result_from_est(self, input_image, candidate, subset, det_result, image_resolution, output_type, H, W): | |
| # nums, keys, locs = candidate.shape | |
| # candidate[..., 0] /= float(W) | |
| # candidate[..., 1] /= float(H) | |
| # score = subset[:, :18] # 前18个是躯干骨骼 score(n, 18) | |
| # max_ind = np.mean(score, axis=-1).argmax(axis=0) # 返回分数最高的锚框对应的骨骼 | |
| # score = score[[max_ind]] | |
| # body = candidate[:, :18].copy() | |
| # body = body[[max_ind]] | |
| # nums = 1 | |
| # body = body.reshape(nums * 18, locs) # Moore-AA只有一个人体, 0-18表示body | |
| # body_score = copy.deepcopy(score) # 已经去过max_ind | |
| # for i in range(len(score)): | |
| # for j in range(len(score[i])): | |
| # if score[i][j] > 0.3: | |
| # score[i][j] = int(18 * i + j) | |
| # else: | |
| # score[i][j] = -1 # 躯干中去除掉不可见的骨骼 | |
| # un_visible = subset < 0.3 | |
| # candidate[un_visible] = -1 # 全部关键点中去掉不可见骨骼 | |
| # foot = candidate[:, 18:24] | |
| # faces = candidate[[max_ind], 24:92] | |
| # hands = candidate[[max_ind], 92:113] | |
| # hands = np.vstack([hands, candidate[[max_ind], 113:]]) | |
| # bodies = dict(candidate=body, subset=score) | |
| # pose = dict(bodies=bodies, hands=hands, faces=faces) | |
| # return pose, body_score, det_result # body_score是原始的躯干骨骼分数 | |
| def __call__( | |
| self, | |
| input, | |
| **kwargs, | |
| ): | |
| if not self.use_batch: | |
| # PIL要不要颜色反转？ | |
| input = cv2.cvtColor( | |
| np.array(input, dtype=np.uint8), cv2.COLOR_RGB2BGR | |
| ) | |
| input = HWC3(input) | |
| H, W, C = input.shape | |
| with torch.no_grad(): | |
| candidate, subset, det_result = self.pose_estimation(input) # candidate (n, 134, 2) 候选点 / subset (n, 134) 得分 | |
| return self._get_multi_result_from_est(candidate, subset, det_result, H, W) | |
| else: | |
| raise NotImplementedError("DWposeDetector does not support batch mode") | |