diff --git a/mimicmotion/dwpose/.gitignore b/mimicmotion/dwpose/.gitignore new file mode 100644 index 0000000..0d20b64 --- /dev/null +++ b/mimicmotion/dwpose/.gitignore @@ -0,0 +1 @@ +*.pyc diff --git a/mimicmotion/dwpose/__init__.py b/mimicmotion/dwpose/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mimicmotion/dwpose/dwpose_detector.py b/mimicmotion/dwpose/dwpose_detector.py new file mode 100644 index 0000000..5c6241e --- /dev/null +++ b/mimicmotion/dwpose/dwpose_detector.py @@ -0,0 +1,64 @@ +import os + +import numpy as np +import torch + +from .wholebody import Wholebody + +os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +class DWposeDetector: + """ + A pose detect method for image-like data. + + Parameters: + model_det: (str) serialized ONNX format model path, + such as https://huggingface.co/yzd-v/DWPose/blob/main/yolox_l.onnx + model_pose: (str) serialized ONNX format model path, + such as https://huggingface.co/yzd-v/DWPose/blob/main/dw-ll_ucoco_384.onnx + device: (str) 'cpu' or 'cuda:{device_id}' + """ + def __init__(self, model_det, model_pose, device='cpu'): + self.pose_estimation = Wholebody(model_det=model_det, model_pose=model_pose, device=device) + + def __call__(self, oriImg): + oriImg = oriImg.copy() + H, W, C = oriImg.shape + with torch.no_grad(): + candidate, score = self.pose_estimation(oriImg) + nums, _, locs = candidate.shape + candidate[..., 0] /= float(W) + candidate[..., 1] /= float(H) + body = candidate[:, :18].copy() + body = body.reshape(nums * 18, locs) + subset = score[:, :18].copy() + for i in range(len(subset)): + for j in range(len(subset[i])): + if subset[i][j] > 0.3: + subset[i][j] = int(18 * i + j) + else: + subset[i][j] = -1 + + # un_visible = subset < 0.3 + # candidate[un_visible] = -1 + + # foot = candidate[:, 18:24] + + faces = candidate[:, 24:92] + + hands = candidate[:, 92:113] + hands = np.vstack([hands, candidate[:, 113:]]) + + faces_score = score[:, 24:92] + hands_score = np.vstack([score[:, 92:113], score[:, 113:]]) + + bodies = dict(candidate=body, subset=subset, score=score[:, :18]) + pose = dict(bodies=bodies, hands=hands, hands_score=hands_score, faces=faces, faces_score=faces_score) + + return pose + +# dwpose_detector = DWposeDetector( +# model_det="models/DWPose/yolox_l.onnx", +# model_pose="models/DWPose/dw-ll_ucoco_384.onnx", +# device=device) diff --git a/mimicmotion/dwpose/onnxdet.py b/mimicmotion/dwpose/onnxdet.py new file mode 100644 index 0000000..25b1cec --- /dev/null +++ b/mimicmotion/dwpose/onnxdet.py @@ -0,0 +1,145 @@ +import cv2 +import numpy as np + + +def nms(boxes, scores, nms_thr): + """Single class NMS implemented in Numpy. + + Args: + boxes (np.ndarray): shape=(N,4); N is number of boxes + scores (np.ndarray): the score of bboxes + nms_thr (float): the threshold in NMS + + Returns: + List[int]: output bbox ids + """ + x1 = boxes[:, 0] + y1 = boxes[:, 1] + x2 = boxes[:, 2] + y2 = boxes[:, 3] + + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + order = scores.argsort()[::-1] + + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + + w = np.maximum(0.0, xx2 - xx1 + 1) + h = np.maximum(0.0, yy2 - yy1 + 1) + inter = w * h + ovr = inter / (areas[i] + areas[order[1:]] - inter) + + inds = np.where(ovr <= nms_thr)[0] + order = order[inds + 1] + + return keep + +def multiclass_nms(boxes, scores, nms_thr, score_thr): + """Multiclass NMS implemented in Numpy. Class-aware version. + + Args: + boxes (np.ndarray): shape=(N,4); N is number of boxes + scores (np.ndarray): the score of bboxes + nms_thr (float): the threshold in NMS + score_thr (float): the threshold of cls score + + Returns: + np.ndarray: outputs bboxes coordinate + """ + final_dets = [] + num_classes = scores.shape[1] + for cls_ind in range(num_classes): + cls_scores = scores[:, cls_ind] + valid_score_mask = cls_scores > score_thr + if valid_score_mask.sum() == 0: + continue + else: + valid_scores = cls_scores[valid_score_mask] + valid_boxes = boxes[valid_score_mask] + keep = nms(valid_boxes, valid_scores, nms_thr) + if len(keep) > 0: + cls_inds = np.ones((len(keep), 1)) * cls_ind + dets = np.concatenate( + [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1 + ) + final_dets.append(dets) + if len(final_dets) == 0: + return None + return np.concatenate(final_dets, 0) + +def demo_postprocess(outputs, img_size, p6=False): + grids = [] + expanded_strides = [] + strides = [8, 16, 32] if not p6 else [8, 16, 32, 64] + + hsizes = [img_size[0] // stride for stride in strides] + wsizes = [img_size[1] // stride for stride in strides] + + for hsize, wsize, stride in zip(hsizes, wsizes, strides): + xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize)) + grid = np.stack((xv, yv), 2).reshape(1, -1, 2) + grids.append(grid) + shape = grid.shape[:2] + expanded_strides.append(np.full((*shape, 1), stride)) + + grids = np.concatenate(grids, 1) + expanded_strides = np.concatenate(expanded_strides, 1) + outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides + outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides + + return outputs + +def preprocess(img, input_size, swap=(2, 0, 1)): + if len(img.shape) == 3: + padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114 + else: + padded_img = np.ones(input_size, dtype=np.uint8) * 114 + + r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1]) + resized_img = cv2.resize( + img, + (int(img.shape[1] * r), int(img.shape[0] * r)), + interpolation=cv2.INTER_LINEAR, + ).astype(np.uint8) + padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img + + padded_img = padded_img.transpose(swap) + padded_img = np.ascontiguousarray(padded_img, dtype=np.float32) + return padded_img, r + +def inference_detector(session, oriImg): + """run human detect + """ + input_shape = (640,640) + img, ratio = preprocess(oriImg, input_shape) + + ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]} + output = session.run(None, ort_inputs) + predictions = demo_postprocess(output[0], input_shape)[0] + + boxes = predictions[:, :4] + scores = predictions[:, 4:5] * predictions[:, 5:] + + boxes_xyxy = np.ones_like(boxes) + boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2. + boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2. + boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2. + boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2. + boxes_xyxy /= ratio + dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1) + if dets is not None: + final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5] + isscore = final_scores>0.3 + iscat = final_cls_inds == 0 + isbbox = [ i and j for (i, j) in zip(isscore, iscat)] + final_boxes = final_boxes[isbbox] + else: + final_boxes = np.array([]) + + return final_boxes diff --git a/mimicmotion/dwpose/onnxpose.py b/mimicmotion/dwpose/onnxpose.py new file mode 100644 index 0000000..62c5081 --- /dev/null +++ b/mimicmotion/dwpose/onnxpose.py @@ -0,0 +1,375 @@ +from typing import List, Tuple + +import cv2 +import numpy as np +import onnxruntime as ort + +def preprocess( + img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256) +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Do preprocessing for RTMPose model inference. + + Args: + img (np.ndarray): Input image in shape. + input_size (tuple): Input image size in shape (w, h). + + Returns: + tuple: + - resized_img (np.ndarray): Preprocessed image. + - center (np.ndarray): Center of image. + - scale (np.ndarray): Scale of image. + """ + # get shape of image + img_shape = img.shape[:2] + out_img, out_center, out_scale = [], [], [] + if len(out_bbox) == 0: + out_bbox = [[0, 0, img_shape[1], img_shape[0]]] + for i in range(len(out_bbox)): + x0 = out_bbox[i][0] + y0 = out_bbox[i][1] + x1 = out_bbox[i][2] + y1 = out_bbox[i][3] + bbox = np.array([x0, y0, x1, y1]) + + # get center and scale + center, scale = bbox_xyxy2cs(bbox, padding=1.25) + + # do affine transformation + resized_img, scale = top_down_affine(input_size, scale, center, img) + + # normalize image + mean = np.array([123.675, 116.28, 103.53]) + std = np.array([58.395, 57.12, 57.375]) + resized_img = (resized_img - mean) / std + + out_img.append(resized_img) + out_center.append(center) + out_scale.append(scale) + + return out_img, out_center, out_scale + + +def inference(sess: ort.InferenceSession, img: np.ndarray) -> np.ndarray: + """Inference RTMPose model. + + Args: + sess (ort.InferenceSession): ONNXRuntime session. + img (np.ndarray): Input image in shape. + + Returns: + outputs (np.ndarray): Output of RTMPose model. + """ + all_out = [] + # build input + for i in range(len(img)): + input = [img[i].transpose(2, 0, 1)] + + # build output + sess_input = {sess.get_inputs()[0].name: input} + sess_output = [] + for out in sess.get_outputs(): + sess_output.append(out.name) + + # run model + outputs = sess.run(sess_output, sess_input) + all_out.append(outputs) + + return all_out + + +def postprocess(outputs: List[np.ndarray], + model_input_size: Tuple[int, int], + center: Tuple[int, int], + scale: Tuple[int, int], + simcc_split_ratio: float = 2.0 + ) -> Tuple[np.ndarray, np.ndarray]: + """Postprocess for RTMPose model output. + + Args: + outputs (np.ndarray): Output of RTMPose model. + model_input_size (tuple): RTMPose model Input image size. + center (tuple): Center of bbox in shape (x, y). + scale (tuple): Scale of bbox in shape (w, h). + simcc_split_ratio (float): Split ratio of simcc. + + Returns: + tuple: + - keypoints (np.ndarray): Rescaled keypoints. + - scores (np.ndarray): Model predict scores. + """ + all_key = [] + all_score = [] + for i in range(len(outputs)): + # use simcc to decode + simcc_x, simcc_y = outputs[i] + keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio) + + # rescale keypoints + keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2 + all_key.append(keypoints[0]) + all_score.append(scores[0]) + + return np.array(all_key), np.array(all_score) + + +def bbox_xyxy2cs(bbox: np.ndarray, + padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]: + """Transform the bbox format from (x,y,w,h) into (center, scale) + + Args: + bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted + as (left, top, right, bottom) + padding (float): BBox padding factor that will be multilied to scale. + Default: 1.0 + + Returns: + tuple: A tuple containing center and scale. + - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or + (n, 2) + - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or + (n, 2) + """ + # convert single bbox from (4, ) to (1, 4) + dim = bbox.ndim + if dim == 1: + bbox = bbox[None, :] + + # get bbox center and scale + x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3]) + center = np.hstack([x1 + x2, y1 + y2]) * 0.5 + scale = np.hstack([x2 - x1, y2 - y1]) * padding + + if dim == 1: + center = center[0] + scale = scale[0] + + return center, scale + + +def _fix_aspect_ratio(bbox_scale: np.ndarray, + aspect_ratio: float) -> np.ndarray: + """Extend the scale to match the given aspect ratio. + + Args: + scale (np.ndarray): The image scale (w, h) in shape (2, ) + aspect_ratio (float): The ratio of ``w/h`` + + Returns: + np.ndarray: The reshaped image scale in (2, ) + """ + w, h = np.hsplit(bbox_scale, [1]) + bbox_scale = np.where(w > h * aspect_ratio, + np.hstack([w, w / aspect_ratio]), + np.hstack([h * aspect_ratio, h])) + return bbox_scale + + +def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray: + """Rotate a point by an angle. + + Args: + pt (np.ndarray): 2D point coordinates (x, y) in shape (2, ) + angle_rad (float): rotation angle in radian + + Returns: + np.ndarray: Rotated point in shape (2, ) + """ + sn, cs = np.sin(angle_rad), np.cos(angle_rad) + rot_mat = np.array([[cs, -sn], [sn, cs]]) + return rot_mat @ pt + + +def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray: + """To calculate the affine matrix, three pairs of points are required. This + function is used to get the 3rd point, given 2D points a & b. + + The 3rd point is defined by rotating vector `a - b` by 90 degrees + anticlockwise, using b as the rotation center. + + Args: + a (np.ndarray): The 1st point (x,y) in shape (2, ) + b (np.ndarray): The 2nd point (x,y) in shape (2, ) + + Returns: + np.ndarray: The 3rd point. + """ + direction = a - b + c = b + np.r_[-direction[1], direction[0]] + return c + + +def get_warp_matrix(center: np.ndarray, + scale: np.ndarray, + rot: float, + output_size: Tuple[int, int], + shift: Tuple[float, float] = (0., 0.), + inv: bool = False) -> np.ndarray: + """Calculate the affine transformation matrix that can warp the bbox area + in the input image to the output size. + + Args: + center (np.ndarray[2, ]): Center of the bounding box (x, y). + scale (np.ndarray[2, ]): Scale of the bounding box + wrt [width, height]. + rot (float): Rotation angle (degree). + output_size (np.ndarray[2, ] | list(2,)): Size of the + destination heatmaps. + shift (0-100%): Shift translation ratio wrt the width/height. + Default (0., 0.). + inv (bool): Option to inverse the affine transform direction. + (inv=False: src->dst or inv=True: dst->src) + + Returns: + np.ndarray: A 2x3 transformation matrix + """ + shift = np.array(shift) + src_w = scale[0] + dst_w = output_size[0] + dst_h = output_size[1] + + # compute transformation matrix + rot_rad = np.deg2rad(rot) + src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad) + dst_dir = np.array([0., dst_w * -0.5]) + + # get four corners of the src rectangle in the original image + src = np.zeros((3, 2), dtype=np.float32) + src[0, :] = center + scale * shift + src[1, :] = center + src_dir + scale * shift + src[2, :] = _get_3rd_point(src[0, :], src[1, :]) + + # get four corners of the dst rectangle in the input image + dst = np.zeros((3, 2), dtype=np.float32) + dst[0, :] = [dst_w * 0.5, dst_h * 0.5] + dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir + dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :]) + + if inv: + warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src)) + else: + warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst)) + + return warp_mat + + +def top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict, + img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Get the bbox image as the model input by affine transform. + + Args: + input_size (dict): The input size of the model. + bbox_scale (dict): The bbox scale of the img. + bbox_center (dict): The bbox center of the img. + img (np.ndarray): The original image. + + Returns: + tuple: A tuple containing center and scale. + - np.ndarray[float32]: img after affine transform. + - np.ndarray[float32]: bbox scale after affine transform. + """ + w, h = input_size + warp_size = (int(w), int(h)) + + # reshape bbox to fixed aspect ratio + bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h) + + # get the affine matrix + center = bbox_center + scale = bbox_scale + rot = 0 + warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h)) + + # do affine transform + img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR) + + return img, bbox_scale + + +def get_simcc_maximum(simcc_x: np.ndarray, + simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Get maximum response location and value from simcc representations. + + Note: + instance number: N + num_keypoints: K + heatmap height: H + heatmap width: W + + Args: + simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx) + simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy) + + Returns: + tuple: + - locs (np.ndarray): locations of maximum heatmap responses in shape + (K, 2) or (N, K, 2) + - vals (np.ndarray): values of maximum heatmap responses in shape + (K,) or (N, K) + """ + N, K, Wx = simcc_x.shape + simcc_x = simcc_x.reshape(N * K, -1) + simcc_y = simcc_y.reshape(N * K, -1) + + # get maximum value locations + x_locs = np.argmax(simcc_x, axis=1) + y_locs = np.argmax(simcc_y, axis=1) + locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) + max_val_x = np.amax(simcc_x, axis=1) + max_val_y = np.amax(simcc_y, axis=1) + + # get maximum value across x and y axis + mask = max_val_x > max_val_y + max_val_x[mask] = max_val_y[mask] + vals = max_val_x + locs[vals <= 0.] = -1 + + # reshape + locs = locs.reshape(N, K, 2) + vals = vals.reshape(N, K) + + return locs, vals + + +def decode(simcc_x: np.ndarray, simcc_y: np.ndarray, + simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]: + """Modulate simcc distribution with Gaussian. + + Args: + simcc_x (np.ndarray[K, Wx]): model predicted simcc in x. + simcc_y (np.ndarray[K, Wy]): model predicted simcc in y. + simcc_split_ratio (int): The split ratio of simcc. + + Returns: + tuple: A tuple containing center and scale. + - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2) + - np.ndarray[float32]: scores in shape (K,) or (n, K) + """ + keypoints, scores = get_simcc_maximum(simcc_x, simcc_y) + keypoints /= simcc_split_ratio + + return keypoints, scores + + +def inference_pose(session, out_bbox, oriImg): + """run pose detect + + Args: + session (ort.InferenceSession): ONNXRuntime session. + out_bbox (np.ndarray): bbox list + oriImg (np.ndarray): Input image in shape. + + Returns: + tuple: + - keypoints (np.ndarray): Rescaled keypoints. + - scores (np.ndarray): Model predict scores. + """ + h, w = session.get_inputs()[0].shape[2:] + model_input_size = (w, h) + # preprocess for rtm-pose model inference. + resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size) + # run pose estimation for processed img + outputs = inference(session, resized_img) + # postprocess for rtm-pose model output. + keypoints, scores = postprocess(outputs, model_input_size, center, scale) + + return keypoints, scores diff --git a/mimicmotion/dwpose/preprocess.py b/mimicmotion/dwpose/preprocess.py new file mode 100644 index 0000000..a987874 --- /dev/null +++ b/mimicmotion/dwpose/preprocess.py @@ -0,0 +1,70 @@ +import decord +import numpy as np + +from .util import draw_pose +from .dwpose_detector import dwpose_detector as dwprocessor + + +def get_video_pose( + video_path: str, + ref_image: np.ndarray, + sample_stride: int=1): + """preprocess ref image pose and video pose + + Args: + video_path (str): video pose path + ref_image (np.ndarray): reference image + sample_stride (int, optional): Defaults to 1. + + Returns: + np.ndarray: sequence of video pose + """ + # select ref-keypoint from reference pose for pose rescale + ref_pose = dwprocessor(ref_image) + ref_keypoint_id = [0, 1, 2, 5, 8, 11, 14, 15, 16, 17] + ref_keypoint_id = [i for i in ref_keypoint_id \ + if ref_pose['bodies']['score'].shape[0] > 0 and ref_pose['bodies']['score'][0][i] > 0.3] + ref_body = ref_pose['bodies']['candidate'][ref_keypoint_id] + + height, width, _ = ref_image.shape + + # read input video + vr = decord.VideoReader(video_path, ctx=decord.cpu(0)) + sample_stride *= max(1, int(vr.get_avg_fps() / 24)) + + detected_poses = [dwprocessor(frm) for frm in vr.get_batch(list(range(0, len(vr), sample_stride))).asnumpy()] + + detected_bodies = np.stack( + [p['bodies']['candidate'] for p in detected_poses if p['bodies']['candidate'].shape[0] == 18])[:, + ref_keypoint_id] + # compute linear-rescale params + ay, by = np.polyfit(detected_bodies[:, :, 1].flatten(), np.tile(ref_body[:, 1], len(detected_bodies)), 1) + fh, fw, _ = vr[0].shape + ax = ay / (fh / fw / height * width) + bx = np.mean(np.tile(ref_body[:, 0], len(detected_bodies)) - detected_bodies[:, :, 0].flatten() * ax) + a = np.array([ax, ay]) + b = np.array([bx, by]) + output_pose = [] + # pose rescale + for detected_pose in detected_poses: + detected_pose['bodies']['candidate'] = detected_pose['bodies']['candidate'] * a + b + detected_pose['faces'] = detected_pose['faces'] * a + b + detected_pose['hands'] = detected_pose['hands'] * a + b + im = draw_pose(detected_pose, height, width) + output_pose.append(np.array(im)) + return np.stack(output_pose) + + +def get_image_pose(ref_image): + """process image pose + + Args: + ref_image (np.ndarray): reference image pixel value + + Returns: + np.ndarray: pose visual image in RGB-mode + """ + height, width, _ = ref_image.shape + ref_pose = dwprocessor(ref_image) + pose_img = draw_pose(ref_pose, height, width) + return np.array(pose_img) diff --git a/mimicmotion/dwpose/util.py b/mimicmotion/dwpose/util.py new file mode 100644 index 0000000..5ee6b22 --- /dev/null +++ b/mimicmotion/dwpose/util.py @@ -0,0 +1,136 @@ +import math +import numpy as np +import matplotlib +import cv2 + + +eps = 0.01 + +def alpha_blend_color(color, alpha): + """blend color according to point conf + """ + return [int(c * alpha) for c in color] + +def draw_bodypose(canvas, candidate, subset, score): + H, W, C = canvas.shape + candidate = np.array(candidate) + subset = np.array(subset) + + stickwidth = 4 + + limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \ + [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \ + [1, 16], [16, 18], [3, 17], [6, 18]] + + colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \ + [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \ + [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]] + + for i in range(17): + for n in range(len(subset)): + index = subset[n][np.array(limbSeq[i]) - 1] + conf = score[n][np.array(limbSeq[i]) - 1] + if conf[0] < 0.3 or conf[1] < 0.3: + continue + Y = candidate[index.astype(int), 0] * float(W) + X = candidate[index.astype(int), 1] * float(H) + mX = np.mean(X) + mY = np.mean(Y) + length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 + angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) + polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1) + cv2.fillConvexPoly(canvas, polygon, alpha_blend_color(colors[i], conf[0] * conf[1])) + + canvas = (canvas * 0.6).astype(np.uint8) + + for i in range(18): + for n in range(len(subset)): + index = int(subset[n][i]) + if index == -1: + continue + x, y = candidate[index][0:2] + conf = score[n][i] + x = int(x * W) + y = int(y * H) + cv2.circle(canvas, (int(x), int(y)), 4, alpha_blend_color(colors[i], conf), thickness=-1) + + return canvas + +def draw_handpose(canvas, all_hand_peaks, all_hand_scores): + H, W, C = canvas.shape + + edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \ + [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] + + for peaks, scores in zip(all_hand_peaks, all_hand_scores): + + for ie, e in enumerate(edges): + x1, y1 = peaks[e[0]] + x2, y2 = peaks[e[1]] + x1 = int(x1 * W) + y1 = int(y1 * H) + x2 = int(x2 * W) + y2 = int(y2 * H) + score = int(scores[e[0]] * scores[e[1]] * 255) + if x1 > eps and y1 > eps and x2 > eps and y2 > eps: + cv2.line(canvas, (x1, y1), (x2, y2), + matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * score, thickness=2) + + for i, keyponit in enumerate(peaks): + x, y = keyponit + x = int(x * W) + y = int(y * H) + score = int(scores[i] * 255) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), 4, (0, 0, score), thickness=-1) + return canvas + +def draw_facepose(canvas, all_lmks, all_scores): + H, W, C = canvas.shape + for lmks, scores in zip(all_lmks, all_scores): + for lmk, score in zip(lmks, scores): + x, y = lmk + x = int(x * W) + y = int(y * H) + conf = int(score * 255) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), 3, (conf, conf, conf), thickness=-1) + return canvas + +def draw_pose(pose, H, W, include_body, include_hand, include_face, ref_w=2160): + """vis dwpose outputs + + Args: + pose (List): DWposeDetector outputs in dwpose_detector.py + H (int): height + W (int): width + ref_w (int, optional) Defaults to 2160. + + Returns: + np.ndarray: image pixel value in RGB mode + """ + bodies = pose['bodies'] + faces = pose['faces'] + hands = pose['hands'] + candidate = bodies['candidate'] + subset = bodies['subset'] + + sz = min(H, W) + sr = (ref_w / sz) if sz != ref_w else 1 + + ########################################## create zero canvas ################################################## + canvas = np.zeros(shape=(int(H*sr), int(W*sr), 3), dtype=np.uint8) + + ########################################### draw body pose ##################################################### + if include_body: + canvas = draw_bodypose(canvas, candidate, subset, score=bodies['score']) + + ########################################### draw hand pose ##################################################### + if include_hand: + canvas = draw_handpose(canvas, hands, pose['hands_score']) + + ########################################### draw face pose ##################################################### + if include_face: + canvas = draw_facepose(canvas, faces, pose['faces_score']) + + return cv2.cvtColor(cv2.resize(canvas, (W, H)), cv2.COLOR_BGR2RGB).transpose(2, 0, 1) diff --git a/mimicmotion/dwpose/wholebody.py b/mimicmotion/dwpose/wholebody.py new file mode 100644 index 0000000..e416fb6 --- /dev/null +++ b/mimicmotion/dwpose/wholebody.py @@ -0,0 +1,57 @@ +import numpy as np +import onnxruntime as ort + +from .onnxdet import inference_detector +from .onnxpose import inference_pose + + +class Wholebody: + """detect human pose by dwpose + """ + def __init__(self, model_det, model_pose, device="cpu"): + providers = ['CPUExecutionProvider'] if device == 'cpu' else ['CUDAExecutionProvider'] + provider_options = None if device == 'cpu' else [{'device_id': 0}] + + self.session_det = ort.InferenceSession( + path_or_bytes=model_det, providers=providers, provider_options=provider_options + ) + self.session_pose = ort.InferenceSession( + path_or_bytes=model_pose, providers=providers, provider_options=provider_options + ) + + def __call__(self, oriImg): + """call to process dwpose-detect + + Args: + oriImg (np.ndarray): detected image + + """ + det_result = inference_detector(self.session_det, oriImg) + keypoints, scores = inference_pose(self.session_pose, det_result, oriImg) + + keypoints_info = np.concatenate( + (keypoints, scores[..., None]), axis=-1) + # compute neck joint + neck = np.mean(keypoints_info[:, [5, 6]], axis=1) + # neck score when visualizing pred + neck[:, 2:4] = np.logical_and( + keypoints_info[:, 5, 2:4] > 0.3, + keypoints_info[:, 6, 2:4] > 0.3).astype(int) + new_keypoints_info = np.insert( + keypoints_info, 17, neck, axis=1) + mmpose_idx = [ + 17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3 + ] + openpose_idx = [ + 1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17 + ] + new_keypoints_info[:, openpose_idx] = \ + new_keypoints_info[:, mmpose_idx] + keypoints_info = new_keypoints_info + + keypoints, scores = keypoints_info[ + ..., :2], keypoints_info[..., 2] + + return keypoints, scores + + diff --git a/nodes.py b/nodes.py index a2d6282..227f73c 100644 --- a/nodes.py +++ b/nodes.py @@ -3,6 +3,7 @@ from omegaconf import OmegaConf import torch import torch.nn.functional as F import sys +import numpy as np script_directory = os.path.dirname(os.path.abspath(__file__)) sys.path.append(script_directory) @@ -27,7 +28,6 @@ from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection from mimicmotion.modules.unet import UNetSpatioTemporalConditionModel from mimicmotion.modules.pose_net import PoseNet -from mimicmotion.pipelines.pipeline_mimicmotion import MimicMotionPipeline class MimicMotionModel(torch.nn.Module): def __init__(self, base_model_path): @@ -181,14 +181,115 @@ class MimicMotionSampler: print(frames.shape) return frames, - + +class MimicMotionGetPoses: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "ref_image": ("IMAGE",), + "pose_images": ("IMAGE",), + "include_body": ("BOOLEAN", {"default": True}), + "include_hand": ("BOOLEAN", {"default": True}), + "include_face": ("BOOLEAN", {"default": True}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "process" + CATEGORY = "MimicMotionWrapper" + + def process(self, ref_image, pose_images, include_body, include_hand, include_face): + device = mm.get_torch_device() + from mimicmotion.dwpose.util import draw_pose + from mimicmotion.dwpose.dwpose_detector import DWposeDetector + + yolo_model = "yolox_l.onnx" + dw_pose_model = "dw-ll_ucoco_384.onnx" + model_base_path = os.path.join(script_directory, "models", "DWPose") + + model_det=os.path.join(model_base_path, yolo_model) + model_pose=os.path.join(model_base_path, dw_pose_model) + + if not os.path.exists(model_det): + print(f"Downloading yolo model to: {model_base_path}") + from huggingface_hub import snapshot_download + snapshot_download(repo_id="yzd-v/DWPose", + allow_patterns=[f"*{yolo_model}*"], + local_dir=model_base_path, + local_dir_use_symlinks=False) + + if not os.path.exists(model_pose): + print(f"Downloading dwpose model to: {model_base_path}") + from huggingface_hub import snapshot_download + snapshot_download(repo_id="yzd-v/DWPose", + allow_patterns=[f"*{dw_pose_model}*"], + local_dir=model_base_path, + local_dir_use_symlinks=False) + + dwprocessor = DWposeDetector( + model_det=os.path.join(model_base_path, "yolox_l.onnx"), + model_pose=os.path.join(model_base_path, "dw-ll_ucoco_384.onnx"), + device=device) + + ref_image = ref_image.squeeze(0).cpu().numpy() * 255 + + # select ref-keypoint from reference pose for pose rescale + ref_pose = dwprocessor(ref_image) + ref_keypoint_id = [0, 1, 2, 5, 8, 11, 14, 15, 16, 17] + ref_keypoint_id = [i for i in ref_keypoint_id \ + if ref_pose['bodies']['score'].shape[0] > 0 and ref_pose['bodies']['score'][0][i] > 0.3] + ref_body = ref_pose['bodies']['candidate'][ref_keypoint_id] + + + height, width, _ = ref_image.shape + pose_images_np = pose_images.cpu().numpy() * 255 + + # read input video + detected_poses_np_list = [] + for img_np in pose_images_np: + detected_poses_np_list.append(dwprocessor(img_np)) + + detected_bodies = np.stack( + [p['bodies']['candidate'] for p in detected_poses_np_list if p['bodies']['candidate'].shape[0] == 18])[:, + ref_keypoint_id] + # compute linear-rescale params + ay, by = np.polyfit(detected_bodies[:, :, 1].flatten(), np.tile(ref_body[:, 1], len(detected_bodies)), 1) + fh, fw, _ = pose_images_np[0].shape + ax = ay / (fh / fw / height * width) + bx = np.mean(np.tile(ref_body[:, 0], len(detected_bodies)) - detected_bodies[:, :, 0].flatten() * ax) + a = np.array([ax, ay]) + b = np.array([bx, by]) + output_pose = [] + # pose rescale + for detected_pose in detected_poses_np_list: + detected_pose['bodies']['candidate'] = detected_pose['bodies']['candidate'] * a + b + detected_pose['faces'] = detected_pose['faces'] * a + b + detected_pose['hands'] = detected_pose['hands'] * a + b + im = draw_pose(detected_pose, height, width, include_body=include_body, include_hand=include_hand, include_face=include_face) + output_pose.append(np.array(im)) + + output_pose_tensors = [torch.tensor(np.array(im)) for im in output_pose] + output_tensor = torch.stack(output_pose_tensors) / 255 + + ref_pose_img = draw_pose(ref_pose, height, width, include_body=include_body, include_hand=include_hand, include_face=include_face) + ref_pose_tensor = torch.tensor(np.array(ref_pose_img)) / 255 + output_tensor = torch.cat((ref_pose_tensor.unsqueeze(0), output_tensor)) + output_tensor = output_tensor.permute(0, 2, 3, 1).cpu().float() + + return output_tensor, + + + NODE_CLASS_MAPPINGS = { "DownloadAndLoadMimicMotionModel": DownloadAndLoadMimicMotionModel, "MimicMotionSampler": MimicMotionSampler, + "MimicMotionGetPoses": MimicMotionGetPoses } NODE_DISPLAY_NAME_MAPPINGS = { "DownloadAndLoadMimicMotionModel": "DownloadAndLoadMimicMotionModel", "MimicMotionSampler": "MimicMotionSampler", + "MimicMotionGetPoses": "MimicMotionGetPoses" }