dwpose node

mimicmotion/dwpose/.gitignore

+*.pyc

mimicmotion/dwpose/dwpose_detector.py

+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)

mimicmotion/dwpose/onnxdet.py

+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

mimicmotion/dwpose/preprocess.py

+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)

mimicmotion/dwpose/util.py

+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)

mimicmotion/dwpose/wholebody.py

+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
+
+

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"
 }