Update render_kps.py

README.md

@@ -0,0 +1,123 @@
+# OpenPose_Remapping for ComfyUI(骨骼重映射节点)
+
+> ✨ A pose remapping node with support for joint locking, motion-aware scaling, and optional easing – perfect for animation refinement and mech rig control.
+
+一个用于 **骨骼重映射（Skeleton Remapping）** 的 ComfyUI 节点，支持锁定躯干、末端关键点比例调节、帧间缓动过渡等功能，可用于高质量的人体姿态变换、动画平滑处理与机甲刚体控制。
+
+---
+
+# Pose Remapping 节点更新声明（7.25）
+
+我们对以下节点进行了功能增强与优化，适配更灵活的人体姿态重建需求：
+
+## ✅ 更新节点列表与说明
+
+### 1. `ApplyPoseScalesToFrames`
+- 支持按骨骼段进行 **比例缩放**
+- 支持以参考帧为锚点进行 **全局位置对齐**
+- 内置 **帧间平滑（EMA）**，减少动画抖动
+
+### 2. `RenderKps`
+- 支持批量渲染 **多帧姿态图像**
+- 输出为 ComfyUI 兼容图像列表，可用于后续处理或预览
+
+### 3. `ApplyPoseDiff`
+- 应用两帧姿态间的 **差值位移**
+- 可选择性缩放末端关节（腕、踝等）
+- 支持缓动过渡，适合生成流畅的骨骼动画序列
+
+---
+
+这些节点可协同使用，完成从关键点解析、比例适配、动态迁移到图像渲染的完整流程。
+
+更新时间：2025-07-25  
+作者：SSsnap
+
+---
+## 🔗 相关资源
+
+- 🧠 **模型下载地址**：  
+  [Wan2.1-Fun-14B-Control（Hugging Face）](https://huggingface.co/alibaba-pai/Wan2.1-Fun-14B-Control/tree/main)
+
+- 🧩 **ComfyUI 配套节点**：  
+  [ComfyUI OpenPose editor（GitHub）](https://github.com/huchenlei/ComfyUI-openpose-editor)
+
+🧠 Features | 功能特点
+---
+
+✅ Skeleton Remapping – Automatically remap pose differences from the first frame to the entire animation
+
+✅ 骨骼重映射：基于第一帧姿态差异，自动映射全帧姿态
+
+✅ Motion-Aware Scaling – Dynamically scale movement on end-effectors (hands, feet) based on motion magnitude
+
+✅ 支持比例缩放：末端关键点（手、脚）动作幅度可自适应
+
+✅ Joint Locking – Keep shoulders and hips fixed to preserve torso stability
+
+✅ 关节锁定：肩膀与髋部始终固定，保持身体稳定性
+
+✅ Easing to Target Poses – Supports optional scaled_pose_diff + t to smoothly blend into global targets
+
+✅ 动作缓动过渡：支持 optional scaled_pose_diff + t，实现向任意姿态过渡
+
+✅ COCO‑18 & BODY‑25 Compatible – Auto-detects keypoint format
+
+✅ 兼容 COCO‑18 与 BODY‑25：自动判断关键点格式
+
+✅ Native ComfyUI Integration – Seamlessly works with ControlNet-style pose pipelines
+
+✅ ComfyUI 原生节点，支持与 ControlNet pose pipeline 无缝集成
+
+🚀 Use Cases | 应用场景
+---
+🔄 Smooth animation retargeting for pose keypoints
+
+姿态动画过渡平滑化
+
+🤖 Mecha / Robot control with rigid torso preservation
+
+机器人 / 机甲 动作调整（保持躯干刚性）
+
+
+---
+
+## 🔧 Node Inputs | 节点输入说明
+
+| 输入名          | 类型              | 描述                             |
+| ------------ | --------------- | ------------------------------ |
+| `frame_list` | `POSE_KEYPOINT` | 多帧姿态数据（待缩放与对齐）                 |
+| `scales`     | `SCALES`        | 每条骨骼的缩放比例（如左臂、右腿）              |
+| `ref_pose`   | `POSE_KEYPOINT` | 用于初始对齐的参考姿态帧（例如原始 OpenPose 输出） |
+| `kps`         | `POSE_KEYPOINT` | 多帧关键点数据         |
+| `render_body` | `BOOLEAN`       | 是否绘制身体主骨架（默认启用） |
+
+
+
+## display
+
+Skeletal Remapping +  3D Composition 🦴
+
+
+
+https://github.com/user-attachments/assets/2134ba60-9526-40a6-a653-072945422101
+
+
+
+
+Multi-Camera AI Cinematography 📹
+
+
+
+
+
+
+https://github.com/user-attachments/assets/24875b91-3c67-46d1-898a-2f374f9b94da
+
+
+
+
+
+
+
+

__init__.py

@@ -1,13 +1,22 @@
-from .apply_diss_pose import ApplyPoseDiff
-from .cha import PoseDiffCalculator
+from .ratio_cau import CalculatePoseScales
+from .frame_cau import ApplyPoseScalesToFrames
+from .convertPoseToStandardFormat import ConvertPoseToStandardFormat
+from .render_kps import RenderKps
+# 把下面一行指向同目录下的 show.py
+
 
 NODE_CLASS_MAPPINGS = {
-    "ApplyPoseDiff": ApplyPoseDiff,
-    "PoseDiffCalculator" : PoseDiffCalculator
+    "CalculatePoseScales": CalculatePoseScales,
+    "ApplyPoseScalesToFrames": ApplyPoseScalesToFrames,
+    "ConvertPoseToStandardFormat": ConvertPoseToStandardFormat,
+    "RenderKps": RenderKps,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "ApplyPoseDiff": "SSsnap Apply Pose Diff ✂️",
-    "PoseDiffCalculator" : "SSsnap Pose Diff Calculator 🛠️"
+    "CalculatePoseScales": "C&scales🧮",
+    "ApplyPoseScalesToFrames" : "A&scales🔧",
+    "ConvertPoseToStandardFormat" : "ConvertPose ✂️",
+    "RenderKps" : "RenderKps 📹"
 }
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

apply_diss_pose.py

@@ -1,83 +0,0 @@
-import json
-import numpy as np
-
-class ApplyPoseDiff:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "pose_diff": ("POSE_KEYPOINT", {}),          # 差值数据
-                "original_frame_pose": ("POSE_KEYPOINT", {}),# 原始关键帧数据
-                "canvas_width": ("INT", {"default": 1440}),   # 画布宽度
-                "canvas_height": ("INT", {"default": 1440})   # 画布高度
-            }
-        }
-
-    # 输出类型为 POSE_KEYPOINT，直接对应 RenderPeopleKps 的输入
-    RETURN_TYPES = ("POSE_KEYPOINT",)
-    RETURN_NAMES = ("adjusted_frame_pose",)
-    FUNCTION = "apply_diff"
-    CATEGORY = "Snap Processing"
-
-    def extract_pose_array(self, data):
-        """
-        安全地提取 people[0].pose_keypoints_2d，返回 numpy 数组。
-        data 可能是 [ { ... } ] 或 { ... } 两种情况。
-        """
-        # 若外层是 list，就取第 0 个
-        if isinstance(data, list):
-            data = data[0]
-        try:
-            pose_list = data["people"][0]["pose_keypoints_2d"]
-            return np.array(pose_list)
-        except (IndexError, KeyError, TypeError):
-            raise ValueError("Input pose data is invalid or malformed.")
-
-    def extract_other_keypoints(self, data):
-        """
-        从原始数据中提取 face、手等信息，用于在输出里保持它们。
-        若不需要保留，可注释掉。
-        """
-        if isinstance(data, list):
-            data = data[0]
-        # 为防止找不到对应字段，使用 get
-        person = data["people"][0]
-        return {
-            "face_keypoints_2d": person.get("face_keypoints_2d", []),
-            "hand_left_keypoints_2d": person.get("hand_left_keypoints_2d", []),
-            "hand_right_keypoints_2d": person.get("hand_right_keypoints_2d", [])
-        }
-
-    def apply_diff(self, pose_diff, original_frame_pose, canvas_width, canvas_height):
-        # 1) 提取原始 pose & 差值 pose
-        original_array = self.extract_pose_array(original_frame_pose)
-        diff_array = self.extract_pose_array(pose_diff)
-
-        # 2) 检查一致性
-        if original_array.size == 0 or diff_array.size == 0 or original_array.size != diff_array.size:
-            raise ValueError("Input pose data is invalid or keypoints are mismatched.")
-
-        # 3) 应用差值
-        adjusted_array = (original_array + diff_array).tolist()
-
-        # 4) 若想保留原始 hand/face 关键点，提取并合并
-        other_kps = self.extract_other_keypoints(original_frame_pose)
-
-        # 5) 构建新的关键帧数据 (Python 对象)
-        #    与 RenderPeopleKps 中 decode_json_as_poses(kps) 兼容
-        adjusted_frame = [{
-            "people": [
-                {
-                    "pose_keypoints_2d": adjusted_array,
-                    "hand_left_keypoints_2d": other_kps["hand_left_keypoints_2d"],
-                    "hand_right_keypoints_2d": other_kps["hand_right_keypoints_2d"],
-                    "face_keypoints_2d": other_kps["face_keypoints_2d"]
-                }
-            ],
-            "canvas_width": canvas_width,
-            "canvas_height": canvas_height
-        }]
-
-        # 返回一个 tuple，里面只有一个值，供 ComfyUI 使用
-        return (adjusted_frame,)
-

cha.py

@@ -1,47 +0,0 @@
-import json
-import numpy as np
-
-class PoseDiffCalculator:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "original_frame_pose": ("POSE_KEYPOINT", {}),
-                "adjusted_frame_pose": ("POSE_KEYPOINT", {})
-            }
-        }
-
-    RETURN_TYPES = ("POSE_KEYPOINT",)
-    RETURN_NAMES = ("pose_diff",)
-    FUNCTION = "calculate_diff"
-    CATEGORY = "Snap Processing"
-
-    def extract_pose(self, data):
-        if isinstance(data, list):
-            data = data[0]
-        try:
-            return np.array(data['people'][0]['pose_keypoints_2d'])
-        except (IndexError, KeyError, TypeError):
-            raise ValueError("Input pose data is invalid or malformed.")
-
-    def calculate_diff(self, original_frame_pose, adjusted_frame_pose):
-        original_pose = self.extract_pose(original_frame_pose)
-        adjusted_pose = self.extract_pose(adjusted_frame_pose)
-
-        # 检查数据有效性
-        if original_pose.size == 0 or adjusted_pose.size == 0 or original_pose.size != adjusted_pose.size:
-            raise ValueError("Input pose data is invalid or keypoints are mismatched.")
-
-        # 计算差值
-        pose_diff = (adjusted_pose - original_pose).tolist()
-
-        # 构建差值数据
-        diff_result = [{
-            "people": [
-                {
-                    "pose_keypoints_2d": pose_diff
-                }
-            ]
-        }]
-
-        return (diff_result,)

convertPoseToStandardFormat.py

@@ -0,0 +1,78 @@
+import json
+
+class ConvertPoseToStandardFormat:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "raw_pose_data": ("POSE_KEYPOINT",),  
+            }
+        }
+
+    RETURN_TYPES = ("POSE_KEYPOINT",)
+    RETURN_NAMES = ("standardized_pose_data",)
+    FUNCTION = "convert"
+    CATEGORY = "Snap Processing"
+
+    def convert(self, raw_pose_data):
+       
+        frames = []
+
+        if isinstance(raw_pose_data, str):
+            try:
+                parsed = json.loads(raw_pose_data)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"[ConvertPoseToStandardFormat] JSON 解析失败: {e}")
+            if isinstance(parsed, dict):
+                frames = [parsed]
+            elif isinstance(parsed, list):
+                frames = parsed
+            else:
+                raise ValueError(f"[ConvertPoseToStandardFormat] 解析后不是 dict/list: {type(parsed)}")
+
+        elif isinstance(raw_pose_data, dict):
+            frames = [raw_pose_data]
+
+        elif isinstance(raw_pose_data, list):
+            for idx, entry in enumerate(raw_pose_data):
+                if isinstance(entry, dict):
+                    frames.append(entry)
+                elif isinstance(entry, str):
+                    try:
+                        e = json.loads(entry)
+                    except json.JSONDecodeError:
+                        raise ValueError(f"[ConvertPoseToStandardFormat] 第 {idx} 项 JSON 解析失败：{entry[:100]}…")
+                    if isinstance(e, dict):
+                        frames.append(e)
+                    else:
+                        raise ValueError(f"[ConvertPoseToStandardFormat] 第 {idx} 项解析后不是 dict: {type(e)}")
+                else:
+                    raise ValueError(f"[ConvertPoseToStandardFormat] 第 {idx} 项类型不支持: {type(entry)}")
+        else:
+            raise ValueError(f"[ConvertPoseToStandardFormat] 不支持的 raw_pose_data 类型: {type(raw_pose_data)}")
+
+        output = []
+        for frame in frames:
+            ppl = frame.get("people")
+            if not isinstance(ppl, list) or len(ppl) == 0:
+                
+                print(f"[ConvertPoseToStandardFormat] 跳过无效 frame：{frame}")
+                continue
+
+            new_people = []
+            for person in ppl:
+                new_people.append({
+                    "pose_keypoints_2d": person.get("pose_keypoints_2d", []),
+                    "face_keypoints_2d": [],
+                    "hand_left_keypoints_2d": [],
+                    "hand_right_keypoints_2d": [],
+                })
+
+            output.append({"people": new_people})
+
+        return (output,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "ConvertPoseToStandardFormat": ConvertPoseToStandardFormat,
+}

dwpose/LICENSE

@@ -0,0 +1,108 @@
+OPENPOSE: MULTIPERSON KEYPOINT DETECTION
+SOFTWARE LICENSE AGREEMENT
+ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY
+
+BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS LICENSE AGREEMENT.  IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR DOWNLOAD THE SOFTWARE.
+
+This is a license agreement ("Agreement") between your academic institution or non-profit organization or self (called "Licensee" or "You" in this Agreement) and Carnegie Mellon University (called "Licensor" in this Agreement).  All rights not specifically granted to you in this Agreement are reserved for Licensor. 
+
+RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: 
+Licensor retains exclusive ownership of any copy of the Software (as defined below) licensed under this Agreement and hereby grants to Licensee a personal, non-exclusive, 
+non-transferable license to use the Software for noncommercial research purposes, without the right to sublicense, pursuant to the terms and conditions of this Agreement.  As used in this Agreement, the term "Software" means (i) the actual copy of all or any portion of code for program routines made accessible to Licensee by Licensor pursuant to this Agreement, inclusive of backups, updates, and/or merged copies permitted hereunder or subsequently supplied by Licensor,  including all or any file structures, programming instructions, user interfaces and screen formats and sequences as well as any and all documentation and instructions related to it, and (ii) all or any derivatives and/or modifications created or made by You to any of the items specified in (i).
+
+CONFIDENTIALITY: Licensee acknowledges that the Software is proprietary to Licensor, and as such, Licensee agrees to receive all such materials in confidence and use the Software only in accordance with the terms of this Agreement.  Licensee agrees to use reasonable effort to protect the Software from unauthorized use, reproduction, distribution, or publication.
+
+COPYRIGHT: The Software is owned by Licensor and is protected by United 
+States copyright laws and applicable international treaties and/or conventions.
+
+PERMITTED USES:  The Software may be used for your own noncommercial internal research purposes. You understand and agree that Licensor is not obligated to implement any suggestions and/or feedback you might provide regarding the Software, but to the extent Licensor does so, you are not entitled to any compensation related thereto.
+
+DERIVATIVES: You may create derivatives of or make modifications to the Software, however, You agree that all and any such derivatives and modifications will be owned by Licensor and become a part of the Software licensed to You under this Agreement.  You may only use such derivatives and modifications for your own noncommercial internal research purposes, and you may not otherwise use, distribute or copy such derivatives and modifications in violation of this Agreement.
+
+BACKUPS:  If Licensee is an organization, it may make that number of copies of the Software necessary for internal noncommercial use at a single site within its organization provided that all information appearing in or on the original labels, including the copyright and trademark notices are copied onto the labels of the copies.
+
+USES NOT PERMITTED:  You may not distribute, copy or use the Software except as explicitly permitted herein. Licensee has not been granted any trademark license as part of this Agreement and may not use the name or mark “OpenPose", "Carnegie Mellon" or any renditions thereof without the prior written permission of Licensor.
+
+You may not sell, rent, lease, sublicense, lend, time-share or transfer, in whole or in part, or provide third parties access to prior or present versions (or any parts thereof) of the Software.
+
+ASSIGNMENT: You may not assign this Agreement or your rights hereunder without the prior written consent of Licensor. Any attempted assignment without such consent shall be null and void.
+
+TERM: The term of the license granted by this Agreement is from Licensee's acceptance of this Agreement by downloading the Software or by using the Software until terminated as provided below.
+
+The Agreement automatically terminates without notice if you fail to comply with any provision of this Agreement.  Licensee may terminate this Agreement by ceasing using the Software.  Upon any termination of this Agreement, Licensee will delete any and all copies of the Software. You agree that all provisions which operate to protect the proprietary rights of Licensor shall remain in force should breach occur and that the obligation of confidentiality described in this Agreement is binding in perpetuity and, as such, survives the term of the Agreement.
+
+FEE: Provided Licensee abides completely by the terms and conditions of this Agreement, there is no fee due to Licensor for Licensee's use of the Software in accordance with this Agreement.
+
+DISCLAIMER OF WARRANTIES:  THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON-INFRINGEMENT.  LICENSEE BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND RELATED MATERIALS.
+
+SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is provided as part of this Agreement.  
+
+EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent permitted under applicable law, Licensor shall not be liable for direct, indirect, special, incidental, or consequential damages or lost profits related to Licensee's use of and/or inability to use the Software, even if Licensor is advised of the possibility of such damage.
+
+EXPORT REGULATION: Licensee agrees to comply with any and all applicable 
+U.S. export control laws, regulations, and/or other laws related to embargoes and sanction programs administered by the Office of Foreign Assets Control.
+
+SEVERABILITY: If any provision(s) of this Agreement shall be held to be invalid, illegal, or unenforceable by a court or other tribunal of competent jurisdiction, the validity, legality and enforceability of the remaining provisions shall not in any way be affected or impaired thereby.
+
+NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right or remedy under this Agreement shall be construed as a waiver of any future or other exercise of such right or remedy by Licensor.
+
+GOVERNING LAW: This Agreement shall be construed and enforced in accordance with the laws of the Commonwealth of Pennsylvania without reference to conflict of laws principles.  You consent to the personal jurisdiction of the courts of this County and waive their rights to venue outside of Allegheny County, Pennsylvania.
+
+ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and entire agreement between Licensee and Licensor as to the matter set forth herein and supersedes any previous agreements, understandings, and arrangements between the parties relating hereto.
+
+
+
+************************************************************************
+
+THIRD-PARTY SOFTWARE NOTICES AND INFORMATION
+
+This project incorporates material from the project(s) listed below (collectively, "Third Party Code").  This Third Party Code is licensed to you under their original license terms set forth below.  We reserves all other rights not expressly granted, whether by implication, estoppel or otherwise.
+ 
+1.	Caffe, version 1.0.0, (https://github.com/BVLC/caffe/)
+
+COPYRIGHT
+
+All contributions by the University of California:
+Copyright (c) 2014-2017 The Regents of the University of California (Regents)
+All rights reserved.
+
+All other contributions:
+Copyright (c) 2014-2017, the respective contributors
+All rights reserved.
+
+Caffe uses a shared copyright model: each contributor holds copyright over
+their contributions to Caffe. The project versioning records all such
+contribution and copyright details. If a contributor wants to further mark
+their specific copyright on a particular contribution, they should indicate
+their copyright solely in the commit message of the change when it is
+committed.
+
+LICENSE
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met: 
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer. 
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution. 
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+CONTRIBUTION AGREEMENT
+
+By contributing to the BVLC/caffe repository through pull-request, comment,
+or otherwise, the contributor releases their content to the
+license and copyright terms herein.
+
+************END OF THIRD-PARTY SOFTWARE NOTICES AND INFORMATION**********

dwpose/__init__.py

@@ -0,0 +1,307 @@
+# 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)
+# 5th Edited by ControlNet (Improved JSON serialization/deserialization, and lots of bug fixs)
+# This preprocessor is licensed by CMU for non-commercial use only.
+
+import os
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+import json
+import torch
+import numpy as np
+from . import util
+from .body import Body, BodyResult, Keypoint
+from .hand import Hand
+from .face import Face
+from .types import PoseResult, HandResult, FaceResult, AnimalPoseResult
+from huggingface_hub import hf_hub_download
+from .wholebody import Wholebody
+import warnings
+from ..tools import HWC3, resize_image_with_pad, common_input_validate, custom_hf_download
+import cv2
+from PIL import Image
+from .animalpose import AnimalPoseImage
+
+from typing import Tuple, List, Callable, Union, Optional
+
+
+def draw_animalposes(animals: list[list[Keypoint]], H: int, W: int) -> np.ndarray:
+    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
+    for animal_pose in animals:
+        canvas = draw_animalpose(canvas, animal_pose)
+    return canvas
+
+
+def draw_animalpose(canvas: np.ndarray, keypoints: list[Keypoint]) -> np.ndarray:
+    # order of the keypoints for AP10k and a standardized list of colors for limbs
+    keypointPairsList = [
+        (1, 2),
+        (2, 3),
+        (1, 3),
+        (3, 4),
+        (4, 9),
+        (9, 10),
+        (10, 11),
+        (4, 6),
+        (6, 7),
+        (7, 8),
+        (4, 5),
+        (5, 15),
+        (15, 16),
+        (16, 17),
+        (5, 12),
+        (12, 13),
+        (13, 14),
+    ]
+    colorsList = [
+        (255, 255, 255),
+        (100, 255, 100),
+        (150, 255, 255),
+        (100, 50, 255),
+        (50, 150, 200),
+        (0, 255, 255),
+        (0, 150, 0),
+        (0, 0, 255),
+        (0, 0, 150),
+        (255, 50, 255),
+        (255, 0, 255),
+        (255, 0, 0),
+        (150, 0, 0),
+        (255, 255, 100),
+        (0, 150, 0),
+        (255, 255, 0),
+        (150, 150, 150),
+    ]  # 16 colors needed
+
+    for ind, (i, j) in enumerate(keypointPairsList):
+        p1 = keypoints[i - 1]
+        p2 = keypoints[j - 1]
+
+        if p1 is not None and p2 is not None:
+            cv2.line(
+                canvas,
+                (int(p1.x), int(p1.y)),
+                (int(p2.x), int(p2.y)),
+                colorsList[ind],
+                5,
+            )
+    return canvas
+
+
+def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, draw_face=True, xinsr_stick_scaling=False):
+
+    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
+
+    for pose in poses:
+        if draw_body:
+            canvas = util.draw_bodypose(canvas, pose.body.keypoints, xinsr_stick_scaling)
+
+        if draw_hand:
+            canvas = util.draw_handpose(canvas, pose.left_hand)
+            canvas = util.draw_handpose(canvas, pose.right_hand)
+
+        if draw_face:
+            canvas = util.draw_facepose(canvas, pose.face)
+
+    return canvas
+
+
+def decode_json_as_poses(
+    pose_json: dict,
+) -> Tuple[List[PoseResult], List[AnimalPoseResult], int, int]:
+
+    height = pose_json["canvas_height"]
+    width = pose_json["canvas_width"]
+
+    def chunks(lst, n):
+        """Yield successive n-sized chunks from lst."""
+        for i in range(0, len(lst), n):
+            yield lst[i : i + n]
+
+    def decompress_keypoints(
+        numbers: Optional[List[float]],
+    ) -> Optional[List[Optional[Keypoint]]]:
+        if not numbers:
+            return None
+
+        assert len(numbers) % 3 == 0
+
+        def create_keypoint(x, y, c):
+            if c < 1.0:
+                return None
+            keypoint = Keypoint(x, y)
+            return keypoint
+
+        return [create_keypoint(x, y, c) for x, y, c in chunks(numbers, n=3)]
+
+    return (
+        [
+            PoseResult(
+                body=BodyResult(
+                    keypoints=decompress_keypoints(pose.get("pose_keypoints_2d"))
+                ),
+                left_hand=decompress_keypoints(pose.get("hand_left_keypoints_2d")),
+                right_hand=decompress_keypoints(pose.get("hand_right_keypoints_2d")),
+                face=decompress_keypoints(pose.get("face_keypoints_2d")),
+            )
+            for pose in pose_json.get("people", [])
+        ],
+        [decompress_keypoints(pose) for pose in pose_json.get("animals", [])],
+        height,
+        width,
+    )
+
+
+def encode_poses_as_dict(poses: List[PoseResult], canvas_height: int, canvas_width: int) -> str:
+    """ Encode the pose as a dict following openpose JSON output format:
+    https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/02_output.md
+    """
+    def compress_keypoints(keypoints: Union[List[Keypoint], None]) -> Union[List[float], None]:
+        if not keypoints:
+            return None
+        
+        return [
+            value
+            for keypoint in keypoints
+            for value in (
+                [float(keypoint.x), float(keypoint.y), 1.0]
+                if keypoint is not None
+                else [0.0, 0.0, 0.0]
+            )
+        ]
+
+    return {
+        'people': [
+            {
+                'pose_keypoints_2d': compress_keypoints(pose.body.keypoints),
+                "face_keypoints_2d": compress_keypoints(pose.face),
+                "hand_left_keypoints_2d": compress_keypoints(pose.left_hand),
+                "hand_right_keypoints_2d":compress_keypoints(pose.right_hand),
+            }
+            for pose in poses
+        ],
+        'canvas_height': canvas_height,
+        'canvas_width': canvas_width,
+    }
+
+global_cached_dwpose = Wholebody()
+
+class DwposeDetector:
+    """
+    A class for detecting human poses in images using the Dwpose model.
+
+    Attributes:
+        model_dir (str): Path to the directory where the pose models are stored.
+    """
+    def __init__(self, dw_pose_estimation):
+        self.dw_pose_estimation = dw_pose_estimation
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path, pretrained_det_model_or_path=None, det_filename=None, pose_filename=None, torchscript_device="cuda"):
+        global global_cached_dwpose
+        pretrained_det_model_or_path = pretrained_det_model_or_path or pretrained_model_or_path
+
+        pose_filename = pose_filename or "dw-ll_ucoco_384.onnx"
+        
+        det_model_path = None
+        if det_filename is not None:
+            det_model_path = custom_hf_download(pretrained_det_model_or_path, det_filename)
+        pose_model_path = custom_hf_download(pretrained_model_or_path, pose_filename)
+        
+        print(f"\nDWPose: Using {det_filename} for bbox detection and {pose_filename} for pose estimation")
+        if global_cached_dwpose.det is None or global_cached_dwpose.det_filename != det_filename:
+            t = Wholebody(det_model_path, None, torchscript_device=torchscript_device)
+            t.pose = global_cached_dwpose.pose
+            t.pose_filename = global_cached_dwpose.pose
+            global_cached_dwpose = t
+        
+        if global_cached_dwpose.pose is None or global_cached_dwpose.pose_filename != pose_filename:
+            t = Wholebody(None, pose_model_path, torchscript_device=torchscript_device)
+            t.det = global_cached_dwpose.det
+            t.det_filename = global_cached_dwpose.det_filename
+            global_cached_dwpose = t
+        return cls(global_cached_dwpose)
+
+    def detect_poses(self, oriImg) -> List[PoseResult]:
+        with torch.no_grad():
+            keypoints_info = self.dw_pose_estimation(oriImg.copy())
+            return Wholebody.format_result(keypoints_info)
+    
+    def __call__(self, input_image, detect_resolution=512, include_body=True, include_hand=False, include_face=False, hand_and_face=None, output_type="pil", image_and_json=False, upscale_method="INTER_CUBIC", xinsr_stick_scaling=False, **kwargs):
+        if hand_and_face is not None:
+            warnings.warn("hand_and_face is deprecated. Use include_hand and include_face instead.", DeprecationWarning)
+            include_hand = hand_and_face
+            include_face = hand_and_face
+
+        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
+        input_image, _ = resize_image_with_pad(input_image, 0, upscale_method)
+        poses = self.detect_poses(input_image)
+        
+        canvas = draw_poses(poses, input_image.shape[0], input_image.shape[1], draw_body=include_body, draw_hand=include_hand, draw_face=include_face, xinsr_stick_scaling=xinsr_stick_scaling)
+        canvas, remove_pad = resize_image_with_pad(canvas, detect_resolution, upscale_method)
+        detected_map = HWC3(remove_pad(canvas))
+
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
+        
+        if image_and_json:
+            return (detected_map, encode_poses_as_dict(poses, input_image.shape[0], input_image.shape[1]))
+        
+        return detected_map
+
+global_cached_animalpose = AnimalPoseImage()
+class AnimalposeDetector:
+    """
+    A class for detecting animal poses in images using the RTMPose AP10k model.
+
+    Attributes:
+        model_dir (str): Path to the directory where the pose models are stored.
+    """
+    def __init__(self, animal_pose_estimation):
+        self.animal_pose_estimation = animal_pose_estimation
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path, pretrained_det_model_or_path=None, det_filename="yolox_l.onnx", pose_filename="dw-ll_ucoco_384.onnx", torchscript_device="cuda"):
+        global global_cached_animalpose
+        det_model_path = custom_hf_download(pretrained_det_model_or_path, det_filename)
+        pose_model_path = custom_hf_download(pretrained_model_or_path, pose_filename)
+        
+        print(f"\nAnimalPose: Using {det_filename} for bbox detection and {pose_filename} for pose estimation")
+        if global_cached_animalpose.det is None or global_cached_animalpose.det_filename != det_filename:
+            t = AnimalPoseImage(det_model_path, None, torchscript_device=torchscript_device)
+            t.pose = global_cached_animalpose.pose
+            t.pose_filename = global_cached_animalpose.pose
+            global_cached_animalpose = t
+        
+        if global_cached_animalpose.pose is None or global_cached_animalpose.pose_filename != pose_filename:
+            t = AnimalPoseImage(None, pose_model_path, torchscript_device=torchscript_device)
+            t.det = global_cached_animalpose.det
+            t.det_filename = global_cached_animalpose.det_filename
+            global_cached_animalpose = t
+        return cls(global_cached_animalpose)
+    
+    def __call__(self, input_image, detect_resolution=512, output_type="pil", image_and_json=False, upscale_method="INTER_CUBIC", **kwargs):
+        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
+        input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
+        result = self.animal_pose_estimation(input_image)
+        if result is None:
+            detected_map = np.zeros_like(input_image)
+            openpose_dict = {
+                'version': 'ap10k',
+                'animals': [],
+                'canvas_height': input_image.shape[0],
+                'canvas_width': input_image.shape[1]
+            }
+        else:
+            detected_map, openpose_dict = result
+        detected_map = remove_pad(detected_map)
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
+        
+        if image_and_json:
+            return (detected_map, openpose_dict)
+
+        return detected_map

dwpose/animalpose.py

@@ -0,0 +1,271 @@
+import numpy as np
+import cv2
+import os
+import cv2
+from .dw_onnx.cv_ox_det import inference_detector as inference_onnx_yolox
+from .dw_onnx.cv_ox_yolo_nas import inference_detector as inference_onnx_yolo_nas
+from .dw_onnx.cv_ox_pose import inference_pose as inference_onnx_pose
+
+from .dw_torchscript.jit_det import inference_detector as inference_jit_yolox
+from .dw_torchscript.jit_pose import inference_pose as inference_jit_pose
+from typing import List, Optional
+from .types import PoseResult, BodyResult, Keypoint
+from custom_controlnet_aux.dwpose.util import guess_onnx_input_shape_dtype, get_ort_providers, get_model_type, is_model_torchscript
+from timeit import default_timer
+import torch
+
+def drawBetweenKeypoints(pose_img, keypoints, indexes, color, scaleFactor):
+    ind0 = indexes[0] - 1
+    ind1 = indexes[1] - 1
+    
+    point1 = (keypoints[ind0][0], keypoints[ind0][1])
+    point2 = (keypoints[ind1][0], keypoints[ind1][1])
+
+    thickness = int(5 // scaleFactor)
+
+
+    cv2.line(pose_img, (int(point1[0]), int(point1[1])), (int(point2[0]), int(point2[1])), color, thickness)
+
+
+def drawBetweenKeypointsList(pose_img, keypoints, keypointPairsList, colorsList, scaleFactor):
+    for ind, keypointPair in enumerate(keypointPairsList):
+        drawBetweenKeypoints(pose_img, keypoints, keypointPair, colorsList[ind], scaleFactor)
+
+def drawBetweenSetofKeypointLists(pose_img, keypoints_set, keypointPairsList, colorsList, scaleFactor):
+    for keypoints in keypoints_set:
+        drawBetweenKeypointsList(pose_img, keypoints, keypointPairsList, colorsList, scaleFactor)
+
+
+def padImg(img, size, blackBorder=True):
+    left, right, top, bottom = 0, 0, 0, 0
+
+    # pad x
+    if img.shape[1] < size[1]:
+        sidePadding = int((size[1] - img.shape[1]) // 2)
+        left = sidePadding
+        right = sidePadding
+
+        # pad extra on right if padding needed is an odd number
+        if img.shape[1] % 2 == 1:
+            right += 1
+
+    # pad y
+    if img.shape[0] < size[0]:
+        topBottomPadding = int((size[0] - img.shape[0]) // 2)
+        top = topBottomPadding
+        bottom = topBottomPadding
+        
+        # pad extra on bottom if padding needed is an odd number
+        if img.shape[0] % 2 == 1:
+            bottom += 1
+
+    if blackBorder:
+        paddedImg = cv2.copyMakeBorder(src=img, top=top, bottom=bottom, left=left, right=right, borderType=cv2.BORDER_CONSTANT, value=(0,0,0))
+    else:
+        paddedImg = cv2.copyMakeBorder(src=img, top=top, bottom=bottom, left=left, right=right, borderType=cv2.BORDER_REPLICATE)
+
+    return paddedImg
+
+def smartCrop(img, size, center):
+
+    width = img.shape[1]
+    height = img.shape[0]
+    xSize = size[1]
+    ySize = size[0]
+    xCenter = center[0]
+    yCenter = center[1]
+
+    if img.shape[0] > size[0] or img.shape[1] > size[1]:
+
+
+        leftMargin = xCenter - xSize//2
+        rightMargin = xCenter + xSize//2
+        upMargin = yCenter - ySize//2
+        downMargin = yCenter + ySize//2
+
+
+        if(leftMargin < 0):
+            xCenter += (-leftMargin)
+        if(rightMargin > width):
+            xCenter -= (rightMargin - width)
+
+        if(upMargin < 0):
+            yCenter -= -upMargin
+        if(downMargin > height):
+            yCenter -= (downMargin - height)
+
+
+        img = cv2.getRectSubPix(img, size, (xCenter, yCenter))
+
+
+
+    return img
+
+
+
+def calculateScaleFactor(img, size, poseSpanX, poseSpanY):
+
+    poseSpanX = max(poseSpanX, size[0])
+
+    scaleFactorX = 1
+
+
+    if poseSpanX > size[0]:
+        scaleFactorX = size[0] / poseSpanX
+
+    scaleFactorY = 1
+    if poseSpanY > size[1]:
+        scaleFactorY = size[1] / poseSpanY
+
+    scaleFactor = min(scaleFactorX, scaleFactorY)
+
+
+    return scaleFactor
+
+
+
+def scaleImg(img, size, poseSpanX, poseSpanY, scaleFactor):
+    scaledImg = img
+
+    scaledImg = cv2.resize(img, (0, 0), fx=scaleFactor, fy=scaleFactor) 
+
+    return scaledImg, scaleFactor
+
+class AnimalPoseImage:
+    def __init__(self, det_model_path: Optional[str] = None, pose_model_path: Optional[str] = None, torchscript_device="cuda"):
+        self.det_filename = det_model_path and os.path.basename(det_model_path)
+        self.pose_filename = pose_model_path and os.path.basename(pose_model_path)
+        self.det, self.pose = None, None
+        # return type: None ort cv2 torchscript
+        self.det_model_type = get_model_type("AnimalPose",self.det_filename)
+        self.pose_model_type = get_model_type("AnimalPose",self.pose_filename)
+        # Always loads to CPU to avoid building OpenCV.
+        cv2_device = 'cpu'
+        cv2_backend = cv2.dnn.DNN_BACKEND_OPENCV if cv2_device == 'cpu' else cv2.dnn.DNN_BACKEND_CUDA
+        # You need to manually build OpenCV through cmake to work with your GPU.
+        cv2_providers = cv2.dnn.DNN_TARGET_CPU if cv2_device == 'cpu' else cv2.dnn.DNN_TARGET_CUDA
+        ort_providers = get_ort_providers()
+
+        if self.det_model_type is None:
+            pass
+        elif self.det_model_type == "ort":
+            try:
+                import onnxruntime as ort
+                self.det = ort.InferenceSession(det_model_path, providers=ort_providers)
+            except:
+                print(f"Failed to load onnxruntime with {self.det.get_providers()}.\nPlease change EP_list in the config.yaml and restart ComfyUI")
+                self.det = ort.InferenceSession(det_model_path, providers=["CPUExecutionProvider"])
+        elif self.det_model_type == "cv2":
+            try:
+                self.det = cv2.dnn.readNetFromONNX(det_model_path)
+                self.det.setPreferableBackend(cv2_backend)
+                self.det.setPreferableTarget(cv2_providers)
+            except:
+                print("TopK operators may not work on your OpenCV, try use onnxruntime with CPUExecutionProvider")
+                try:
+                    import onnxruntime as ort
+                    self.det = ort.InferenceSession(det_model_path, providers=["CPUExecutionProvider"])
+                except:
+                    print(f"Failed to load {det_model_path}, you can use other models instead")
+        else:
+            self.det = torch.jit.load(det_model_path)
+            self.det.to(torchscript_device)
+
+        if self.pose_model_type is None:
+            pass
+        elif self.pose_model_type == "ort":
+            try:
+                import onnxruntime as ort
+                self.pose = ort.InferenceSession(pose_model_path, providers=ort_providers)
+            except:
+                print(f"Failed to load onnxruntime with {self.pose.get_providers()}.\nPlease change EP_list in the config.yaml and restart ComfyUI")
+                self.pose = ort.InferenceSession(pose_model_path, providers=["CPUExecutionProvider"])
+        elif self.pose_model_type == "cv2":
+            self.pose = cv2.dnn.readNetFromONNX(pose_model_path)
+            self.pose.setPreferableBackend(cv2_backend)
+            self.pose.setPreferableTarget(cv2_providers)
+        else:
+            self.pose = torch.jit.load(pose_model_path)
+            self.pose.to(torchscript_device)
+
+        if self.pose_filename is not None:
+            self.pose_input_size, _ = guess_onnx_input_shape_dtype(self.pose_filename)
+    
+    def __call__(self, oriImg):
+        detect_classes = list(range(14, 23 + 1)) #https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco.yaml
+        
+        #Sacrifice accurate time measurement for compatibility 
+        det_start = default_timer()
+        if is_model_torchscript(self.det):
+            det_result = inference_jit_yolox(self.det, oriImg, detect_classes=detect_classes)
+        else:
+            det_start = default_timer()
+            det_onnx_dtype = np.float32 if "yolox" in self.det_filename else np.uint8
+            if "yolox" in self.det_filename:
+                det_result = inference_onnx_yolox(self.det, oriImg, detect_classes=detect_classes, dtype=det_onnx_dtype)
+            else:
+                #FP16 and INT8 YOLO NAS accept uint8 input
+                det_result = inference_onnx_yolo_nas(self.det, oriImg, detect_classes=detect_classes, dtype=det_onnx_dtype)
+        print(f"AnimalPose: Bbox {((default_timer() - det_start) * 1000):.2f}ms")
+        
+        if (det_result is None) or (det_result.shape[0] == 0):
+            openpose_dict = {
+                'version': 'ap10k',
+                'animals': [],
+                'canvas_height': oriImg.shape[0],
+                'canvas_width': oriImg.shape[1]
+            }
+            return np.zeros_like(oriImg), openpose_dict
+        
+        pose_start = default_timer()
+        if is_model_torchscript(self.pose):
+            keypoint_sets, scores = inference_jit_pose(self.pose, det_result, oriImg, self.pose_input_size)
+        else:
+            pose_start = default_timer()
+            _, pose_onnx_dtype = guess_onnx_input_shape_dtype(self.pose_filename)
+            keypoint_sets, scores = inference_onnx_pose(self.pose, det_result, oriImg, self.pose_input_size, dtype=pose_onnx_dtype)
+        print(f"AnimalPose: Pose {((default_timer() - pose_start) * 1000):.2f}ms on {det_result.shape[0]} animals\n")
+
+        animal_kps_scores = []
+        pose_img = np.zeros((oriImg.shape[0], oriImg.shape[1], 3), dtype = np.uint8)
+        for (idx, keypoints) in enumerate(keypoint_sets):
+            # don't use keypoints that go outside the frame in calculations for the center
+            interorKeypoints = keypoints[((keypoints[:,0] > 0) & (keypoints[:,0] < oriImg.shape[1])) & ((keypoints[:,1] > 0) & (keypoints[:,1] < oriImg.shape[0]))]
+
+            xVals = interorKeypoints[:,0]
+            yVals = interorKeypoints[:,1]
+
+            minX = np.amin(xVals)
+            minY = np.amin(yVals)
+            maxX = np.amax(xVals)
+            maxY = np.amax(yVals)
+
+            poseSpanX = maxX - minX
+            poseSpanY = maxY - minY
+
+            # find mean center
+
+            xSum = np.sum(xVals)
+            ySum = np.sum(yVals)
+
+            xCenter = xSum // xVals.shape[0]
+            yCenter = ySum // yVals.shape[0]
+            center_of_keypoints = (xCenter,yCenter)
+
+            # order of the keypoints for AP10k and a standardized list of colors for limbs
+            keypointPairsList = [(1,2), (2,3), (1,3), (3,4), (4,9), (9,10), (10,11), (4,6), (6,7), (7,8), (4,5), (5,15), (15,16), (16,17), (5,12), (12,13), (13,14)]
+            colorsList = [(255,255,255), (100,255,100), (150,255,255), (100,50,255), (50,150,200), (0,255,255), (0,150,0), (0,0,255), (0,0,150), (255,50,255), (255,0,255), (255,0,0), (150,0,0), (255,255,100), (0,150,0), (255,255,0), (150,150,150)] # 16 colors needed
+
+            drawBetweenKeypointsList(pose_img, keypoints, keypointPairsList, colorsList, scaleFactor=1.0)
+            score = scores[idx, ..., None]
+            score[score > 1.0] = 1.0
+            score[score < 0.0] = 0.0
+            animal_kps_scores.append(np.concatenate((keypoints, score), axis=-1))
+        
+        openpose_dict = {
+            'version': 'ap10k',
+            'animals': [keypoints.tolist() for keypoints in animal_kps_scores],
+            'canvas_height': oriImg.shape[0],
+            'canvas_width': oriImg.shape[1]
+        }
+        return pose_img, openpose_dict

dwpose/body.py

@@ -0,0 +1,261 @@
+import cv2
+import numpy as np
+import math
+import time
+from scipy.ndimage.filters import gaussian_filter
+import matplotlib.pyplot as plt
+import matplotlib
+import torch
+from torchvision import transforms
+from typing import NamedTuple, List, Union
+
+from . import util
+from .model import bodypose_model
+from .types import Keypoint, BodyResult
+
+class Body(object):
+    def __init__(self, model_path):
+        self.model = bodypose_model()
+        # if torch.cuda.is_available():
+        #     self.model = self.model.cuda()
+            # print('cuda')
+        model_dict = util.transfer(self.model, torch.load(model_path))
+        self.model.load_state_dict(model_dict)
+        self.model.eval()
+
+    def __call__(self, oriImg):
+        # scale_search = [0.5, 1.0, 1.5, 2.0]
+        scale_search = [0.5]
+        boxsize = 368
+        stride = 8
+        padValue = 128
+        thre1 = 0.1
+        thre2 = 0.05
+        multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
+        heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
+        paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
+
+        for m in range(len(multiplier)):
+            scale = multiplier[m]
+            imageToTest = util.smart_resize_k(oriImg, fx=scale, fy=scale)
+            imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
+            im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
+            im = np.ascontiguousarray(im)
+
+            data = torch.from_numpy(im).float()
+            if torch.cuda.is_available():
+                data = data.cuda()
+            # data = data.permute([2, 0, 1]).unsqueeze(0).float()
+            with torch.no_grad():
+                data = data.to(self.cn_device)
+                Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
+            Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
+            Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
+
+            # extract outputs, resize, and remove padding
+            # heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data), (1, 2, 0))  # output 1 is heatmaps
+            heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0))  # output 1 is heatmaps
+            heatmap = util.smart_resize_k(heatmap, fx=stride, fy=stride)
+            heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
+            heatmap = util.smart_resize(heatmap, (oriImg.shape[0], oriImg.shape[1]))
+
+            # paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0))  # output 0 is PAFs
+            paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0))  # output 0 is PAFs
+            paf = util.smart_resize_k(paf, fx=stride, fy=stride)
+            paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
+            paf = util.smart_resize(paf, (oriImg.shape[0], oriImg.shape[1]))
+
+            heatmap_avg += heatmap_avg + heatmap / len(multiplier)
+            paf_avg += + paf / len(multiplier)
+
+        all_peaks = []
+        peak_counter = 0
+
+        for part in range(18):
+            map_ori = heatmap_avg[:, :, part]
+            one_heatmap = gaussian_filter(map_ori, sigma=3)
+
+            map_left = np.zeros(one_heatmap.shape)
+            map_left[1:, :] = one_heatmap[:-1, :]
+            map_right = np.zeros(one_heatmap.shape)
+            map_right[:-1, :] = one_heatmap[1:, :]
+            map_up = np.zeros(one_heatmap.shape)
+            map_up[:, 1:] = one_heatmap[:, :-1]
+            map_down = np.zeros(one_heatmap.shape)
+            map_down[:, :-1] = one_heatmap[:, 1:]
+
+            peaks_binary = np.logical_and.reduce(
+                (one_heatmap >= map_left, one_heatmap >= map_right, one_heatmap >= map_up, one_heatmap >= map_down, one_heatmap > thre1))
+            peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]))  # note reverse
+            peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
+            peak_id = range(peak_counter, peak_counter + len(peaks))
+            peaks_with_score_and_id = [peaks_with_score[i] + (peak_id[i],) for i in range(len(peak_id))]
+
+            all_peaks.append(peaks_with_score_and_id)
+            peak_counter += len(peaks)
+
+        # find connection in the specified sequence, center 29 is in the position 15
+        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]]
+        # the middle joints heatmap correpondence
+        mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], \
+                  [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], \
+                  [55, 56], [37, 38], [45, 46]]
+
+        connection_all = []
+        special_k = []
+        mid_num = 10
+
+        for k in range(len(mapIdx)):
+            score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
+            candA = all_peaks[limbSeq[k][0] - 1]
+            candB = all_peaks[limbSeq[k][1] - 1]
+            nA = len(candA)
+            nB = len(candB)
+            indexA, indexB = limbSeq[k]
+            if (nA != 0 and nB != 0):
+                connection_candidate = []
+                for i in range(nA):
+                    for j in range(nB):
+                        vec = np.subtract(candB[j][:2], candA[i][:2])
+                        norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
+                        norm = max(0.001, norm)
+                        vec = np.divide(vec, norm)
+
+                        startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
+                                            np.linspace(candA[i][1], candB[j][1], num=mid_num)))
+
+                        vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
+                                          for I in range(len(startend))])
+                        vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
+                                          for I in range(len(startend))])
+
+                        score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
+                        score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
+                            0.5 * oriImg.shape[0] / norm - 1, 0)
+                        criterion1 = len(np.nonzero(score_midpts > thre2)[0]) > 0.8 * len(score_midpts)
+                        criterion2 = score_with_dist_prior > 0
+                        if criterion1 and criterion2:
+                            connection_candidate.append(
+                                [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
+
+                connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
+                connection = np.zeros((0, 5))
+                for c in range(len(connection_candidate)):
+                    i, j, s = connection_candidate[c][0:3]
+                    if (i not in connection[:, 3] and j not in connection[:, 4]):
+                        connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
+                        if (len(connection) >= min(nA, nB)):
+                            break
+
+                connection_all.append(connection)
+            else:
+                special_k.append(k)
+                connection_all.append([])
+
+        # last number in each row is the total parts number of that person
+        # the second last number in each row is the score of the overall configuration
+        subset = -1 * np.ones((0, 20))
+        candidate = np.array([item for sublist in all_peaks for item in sublist])
+
+        for k in range(len(mapIdx)):
+            if k not in special_k:
+                partAs = connection_all[k][:, 0]
+                partBs = connection_all[k][:, 1]
+                indexA, indexB = np.array(limbSeq[k]) - 1
+
+                for i in range(len(connection_all[k])):  # = 1:size(temp,1)
+                    found = 0
+                    subset_idx = [-1, -1]
+                    for j in range(len(subset)):  # 1:size(subset,1):
+                        if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
+                            subset_idx[found] = j
+                            found += 1
+
+                    if found == 1:
+                        j = subset_idx[0]
+                        if subset[j][indexB] != partBs[i]:
+                            subset[j][indexB] = partBs[i]
+                            subset[j][-1] += 1
+                            subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
+                    elif found == 2:  # if found 2 and disjoint, merge them
+                        j1, j2 = subset_idx
+                        membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
+                        if len(np.nonzero(membership == 2)[0]) == 0:  # merge
+                            subset[j1][:-2] += (subset[j2][:-2] + 1)
+                            subset[j1][-2:] += subset[j2][-2:]
+                            subset[j1][-2] += connection_all[k][i][2]
+                            subset = np.delete(subset, j2, 0)
+                        else:  # as like found == 1
+                            subset[j1][indexB] = partBs[i]
+                            subset[j1][-1] += 1
+                            subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
+
+                    # if find no partA in the subset, create a new subset
+                    elif not found and k < 17:
+                        row = -1 * np.ones(20)
+                        row[indexA] = partAs[i]
+                        row[indexB] = partBs[i]
+                        row[-1] = 2
+                        row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
+                        subset = np.vstack([subset, row])
+        # delete some rows of subset which has few parts occur
+        deleteIdx = []
+        for i in range(len(subset)):
+            if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
+                deleteIdx.append(i)
+        subset = np.delete(subset, deleteIdx, axis=0)
+
+        # subset: n*20 array, 0-17 is the index in candidate, 18 is the total score, 19 is the total parts
+        # candidate: x, y, score, id
+        return candidate, subset
+    
+    @staticmethod
+    def format_body_result(candidate: np.ndarray, subset: np.ndarray) -> List[BodyResult]:
+        """
+        Format the body results from the candidate and subset arrays into a list of BodyResult objects.
+        
+        Args:
+            candidate (np.ndarray): An array of candidates containing the x, y coordinates, score, and id
+                for each body part.
+            subset (np.ndarray): An array of subsets containing indices to the candidate array for each
+                person detected. The last two columns of each row hold the total score and total parts
+                of the person.
+
+        Returns:
+            List[BodyResult]: A list of BodyResult objects, where each object represents a person with
+                detected keypoints, total score, and total parts.
+        """
+        return [
+            BodyResult(
+                keypoints=[
+                    Keypoint(
+                        x=candidate[candidate_index][0],
+                        y=candidate[candidate_index][1],
+                        score=candidate[candidate_index][2],
+                        id=candidate[candidate_index][3]
+                    ) if candidate_index != -1 else None
+                    for candidate_index in person[:18].astype(int)
+                ],
+                total_score=person[18],
+                total_parts=person[19]
+            )
+            for person in subset
+        ]
+    
+
+if __name__ == "__main__":
+    body_estimation = Body('../model/body_pose_model.pth')
+
+    test_image = '../images/ski.jpg'
+    oriImg = cv2.imread(test_image)  # B,G,R order
+    candidate, subset = body_estimation(oriImg)
+    bodies = body_estimation.format_body_result(candidate, subset)
+
+    canvas = oriImg
+    for body in bodies:
+        canvas = util.draw_bodypose(canvas, body)
+        
+    plt.imshow(canvas[:, :, [2, 1, 0]])
+    plt.show()

dwpose/dw_onnx/__init__.py

@@ -0,0 +1 @@
+#Dummy file ensuring this package will be recognized

dwpose/dw_onnx/cv_ox_det.py

@@ -0,0 +1,129 @@
+import cv2
+import numpy as np
+
+def nms(boxes, scores, nms_thr):
+    """Single class NMS implemented in Numpy."""
+    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."""
+    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, detect_classes=[0], dtype=np.float32):
+    input_shape = (640,640)
+    img, ratio = preprocess(oriImg, input_shape)
+
+    input = img[None, :, :, :]
+    input = input.astype(dtype)
+    if "InferenceSession" in type(session).__name__:
+        input_name = session.get_inputs()[0].name
+        output = session.run(None, {input_name: input})
+    else:
+        outNames = session.getUnconnectedOutLayersNames()
+        session.setInput(input)
+        output = session.forward(outNames)
+
+    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 None:
+        return None
+    final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
+    isscore = final_scores>0.3
+    iscat = np.isin(final_cls_inds, detect_classes)
+    isbbox = [ i and j for (i, j) in zip(isscore, iscat)]
+    final_boxes = final_boxes[isbbox]
+    return final_boxes

dwpose/dw_onnx/cv_ox_pose.py

@@ -0,0 +1,363 @@
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+
+def preprocess(
+    img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Do preprocessing for DWPose 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, img, dtype=np.float32):
+    """Inference DWPose model. Processing all image segments at once to take advantage of GPU's parallelism ability if onnxruntime is installed
+
+    Args:
+        sess : ONNXRuntime session.
+        img : Input image in shape.
+
+    Returns:
+        outputs : Output of DWPose model.
+    """
+    all_out = []
+    # build input
+    input = np.stack(img, axis=0).transpose(0, 3, 1, 2)
+    input = input.astype(dtype)
+    if "InferenceSession" in type(sess).__name__:
+        input_name = sess.get_inputs()[0].name
+        all_outputs = sess.run(None, {input_name: input})
+        for batch_idx in range(len(all_outputs[0])):
+            outputs = [all_outputs[i][batch_idx:batch_idx+1,...] for i in range(len(all_outputs))]
+            all_out.append(outputs)
+        return all_out
+
+    #OpenCV doesn't support batch processing sadly
+    for i in range(len(img)):
+        input = img[i].transpose(2, 0, 1)
+        input = input[None, :, :, :]
+
+        outNames = sess.getUnconnectedOutLayersNames()
+        sess.setInput(input)
+        outputs = sess.forward(outNames)
+        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 DWPose 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, model_input_size=(288, 384), dtype=np.float32):
+    resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
+    outputs = inference(session, resized_img, dtype)
+    keypoints, scores = postprocess(outputs, model_input_size, center, scale)
+
+    return keypoints, scores

dwpose/dw_onnx/cv_ox_yolo_nas.py

@@ -0,0 +1,60 @@
+# Source: https://github.com/Hyuto/yolo-nas-onnx/tree/master/yolo-nas-py
+# Inspired from: https://github.com/Deci-AI/super-gradients/blob/3.1.1/src/super_gradients/training/processing/processing.py
+
+import numpy as np
+import cv2
+
+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, detect_classes=[0], dtype=np.uint8):
+    """
+    This function is only compatible with onnx models exported from the new API with built-in NMS
+    ```py
+    from super_gradients.conversion.conversion_enums import ExportQuantizationMode
+    from super_gradients.common.object_names import Models
+    from super_gradients.training import models
+
+    model = models.get(Models.YOLO_NAS_L, pretrained_weights="coco")
+
+    export_result = model.export(
+        "yolo_nas/yolo_nas_l_fp16.onnx",
+        quantization_mode=ExportQuantizationMode.FP16,
+        device="cuda"
+    )
+    ```
+    """
+    input_shape = (640,640)
+    img, ratio = preprocess(oriImg, input_shape)
+    input = img[None, :, :, :]
+    input = input.astype(dtype)
+    if "InferenceSession" in type(session).__name__:
+        input_name = session.get_inputs()[0].name
+        output = session.run(None, {input_name: input})
+    else:
+        outNames = session.getUnconnectedOutLayersNames()
+        session.setInput(input)
+        output = session.forward(outNames)
+    num_preds, pred_boxes, pred_scores, pred_classes = output
+    num_preds = num_preds[0,0]
+    if num_preds == 0:
+        return None
+    idxs = np.where((np.isin(pred_classes[0, :num_preds], detect_classes)) & (pred_scores[0, :num_preds] > 0.3))
+    if (len(idxs) == 0) or (idxs[0].size == 0):
+        return None
+    return pred_boxes[0, idxs].squeeze(axis=0) / ratio

dwpose/dw_torchscript/__init__.py

@@ -0,0 +1 @@
+#Dummy file ensuring this package will be recognized

dwpose/dw_torchscript/jit_det.py

@@ -0,0 +1,125 @@
+import cv2
+import numpy as np
+import torch
+
+def nms(boxes, scores, nms_thr):
+    """Single class NMS implemented in Numpy."""
+    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."""
+    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(model, oriImg, detect_classes=[0]):
+    input_shape = (640,640)
+    img, ratio = preprocess(oriImg, input_shape)
+
+    device, dtype = next(model.parameters()).device, next(model.parameters()).dtype
+    input = img[None, :, :, :]
+    input = torch.from_numpy(input).to(device, dtype)
+
+    output = model(input).float().cpu().detach().numpy()
+    predictions = demo_postprocess(output[0], input_shape)
+
+    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 None:
+        return None
+    final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
+    isscore = final_scores>0.3
+    iscat = np.isin(final_cls_inds, detect_classes)
+    isbbox = [ i and j for (i, j) in zip(isscore, iscat)]
+    final_boxes = final_boxes[isbbox]
+    return final_boxes

dwpose/dw_torchscript/jit_pose.py

@@ -0,0 +1,363 @@
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+import torch
+
+def preprocess(
+    img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Do preprocessing for DWPose 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(model, img, bs=5):
+    """Inference DWPose model implemented in TorchScript.
+
+    Args:
+        model : TorchScript Model.
+        img : Input image in shape.
+
+    Returns:
+        outputs : Output of DWPose model.
+    """
+    all_out = []
+    # build input
+    orig_img_count = len(img)
+    #Pad zeros to fit batch size
+    for _ in range(bs - (orig_img_count % bs)):
+        img.append(np.zeros_like(img[0]))
+    input = np.stack(img, axis=0).transpose(0, 3, 1, 2)
+    device, dtype = next(model.parameters()).device, next(model.parameters()).dtype
+    input = torch.from_numpy(input).to(device, dtype)
+
+    out1, out2 = [], []
+    for i in range(input.shape[0] // bs):
+        curr_batch_output = model(input[i*bs:(i+1)*bs])
+        out1.append(curr_batch_output[0].float())
+        out2.append(curr_batch_output[1].float())
+    out1, out2 = torch.cat(out1, dim=0)[:orig_img_count], torch.cat(out2, dim=0)[:orig_img_count]
+    out1, out2 = out1.float().cpu().detach().numpy(), out2.float().cpu().detach().numpy()
+    all_outputs = out1, out2
+
+    for batch_idx in range(len(all_outputs[0])):
+        outputs = [all_outputs[i][batch_idx:batch_idx+1,...] for i in range(len(all_outputs))]
+        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 DWPose 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(model, out_bbox, oriImg, model_input_size=(288, 384)):
+    resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
+    #outputs = inference(session, resized_img, dtype)
+    outputs = inference(model, resized_img)
+
+    keypoints, scores = postprocess(outputs, model_input_size, center, scale)
+
+    return keypoints, scores

dwpose/face.py

@@ -0,0 +1,362 @@
+import logging
+import numpy as np
+from torchvision.transforms import ToTensor, ToPILImage
+import torch
+import torch.nn.functional as F
+import cv2
+
+from . import util
+from torch.nn import Conv2d, Module, ReLU, MaxPool2d, init
+
+
+class FaceNet(Module):
+    """Model the cascading heatmaps. """
+    def __init__(self):
+        super(FaceNet, self).__init__()
+        # cnn to make feature map
+        self.relu = ReLU()
+        self.max_pooling_2d = MaxPool2d(kernel_size=2, stride=2)
+        self.conv1_1 = Conv2d(in_channels=3, out_channels=64,
+                              kernel_size=3, stride=1, padding=1)
+        self.conv1_2 = Conv2d(
+            in_channels=64, out_channels=64, kernel_size=3, stride=1,
+            padding=1)
+        self.conv2_1 = Conv2d(
+            in_channels=64, out_channels=128, kernel_size=3, stride=1,
+            padding=1)
+        self.conv2_2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=3, stride=1,
+            padding=1)
+        self.conv3_1 = Conv2d(
+            in_channels=128, out_channels=256, kernel_size=3, stride=1,
+            padding=1)
+        self.conv3_2 = Conv2d(
+            in_channels=256, out_channels=256, kernel_size=3, stride=1,
+            padding=1)
+        self.conv3_3 = Conv2d(
+            in_channels=256, out_channels=256, kernel_size=3, stride=1,
+            padding=1)
+        self.conv3_4 = Conv2d(
+            in_channels=256, out_channels=256, kernel_size=3, stride=1,
+            padding=1)
+        self.conv4_1 = Conv2d(
+            in_channels=256, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv4_2 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv4_3 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv4_4 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv5_1 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv5_2 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv5_3_CPM = Conv2d(
+            in_channels=512, out_channels=128, kernel_size=3, stride=1,
+            padding=1)
+
+        # stage1
+        self.conv6_1_CPM = Conv2d(
+            in_channels=128, out_channels=512, kernel_size=1, stride=1,
+            padding=0)
+        self.conv6_2_CPM = Conv2d(
+            in_channels=512, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage2
+        self.Mconv1_stage2 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage2 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage3
+        self.Mconv1_stage3 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage3 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage4
+        self.Mconv1_stage4 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage4 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage5
+        self.Mconv1_stage5 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage5 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage6
+        self.Mconv1_stage6 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage6 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        for m in self.modules():
+            if isinstance(m, Conv2d):
+                init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        """Return a list of heatmaps."""
+        heatmaps = []
+
+        h = self.relu(self.conv1_1(x))
+        h = self.relu(self.conv1_2(h))
+        h = self.max_pooling_2d(h)
+        h = self.relu(self.conv2_1(h))
+        h = self.relu(self.conv2_2(h))
+        h = self.max_pooling_2d(h)
+        h = self.relu(self.conv3_1(h))
+        h = self.relu(self.conv3_2(h))
+        h = self.relu(self.conv3_3(h))
+        h = self.relu(self.conv3_4(h))
+        h = self.max_pooling_2d(h)
+        h = self.relu(self.conv4_1(h))
+        h = self.relu(self.conv4_2(h))
+        h = self.relu(self.conv4_3(h))
+        h = self.relu(self.conv4_4(h))
+        h = self.relu(self.conv5_1(h))
+        h = self.relu(self.conv5_2(h))
+        h = self.relu(self.conv5_3_CPM(h))
+        feature_map = h
+
+        # stage1
+        h = self.relu(self.conv6_1_CPM(h))
+        h = self.conv6_2_CPM(h)
+        heatmaps.append(h)
+
+        # stage2
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage2(h))
+        h = self.relu(self.Mconv2_stage2(h))
+        h = self.relu(self.Mconv3_stage2(h))
+        h = self.relu(self.Mconv4_stage2(h))
+        h = self.relu(self.Mconv5_stage2(h))
+        h = self.relu(self.Mconv6_stage2(h))
+        h = self.Mconv7_stage2(h)
+        heatmaps.append(h)
+
+        # stage3
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage3(h))
+        h = self.relu(self.Mconv2_stage3(h))
+        h = self.relu(self.Mconv3_stage3(h))
+        h = self.relu(self.Mconv4_stage3(h))
+        h = self.relu(self.Mconv5_stage3(h))
+        h = self.relu(self.Mconv6_stage3(h))
+        h = self.Mconv7_stage3(h)
+        heatmaps.append(h)
+
+        # stage4
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage4(h))
+        h = self.relu(self.Mconv2_stage4(h))
+        h = self.relu(self.Mconv3_stage4(h))
+        h = self.relu(self.Mconv4_stage4(h))
+        h = self.relu(self.Mconv5_stage4(h))
+        h = self.relu(self.Mconv6_stage4(h))
+        h = self.Mconv7_stage4(h)
+        heatmaps.append(h)
+
+        # stage5
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage5(h))
+        h = self.relu(self.Mconv2_stage5(h))
+        h = self.relu(self.Mconv3_stage5(h))
+        h = self.relu(self.Mconv4_stage5(h))
+        h = self.relu(self.Mconv5_stage5(h))
+        h = self.relu(self.Mconv6_stage5(h))
+        h = self.Mconv7_stage5(h)
+        heatmaps.append(h)
+
+        # stage6
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage6(h))
+        h = self.relu(self.Mconv2_stage6(h))
+        h = self.relu(self.Mconv3_stage6(h))
+        h = self.relu(self.Mconv4_stage6(h))
+        h = self.relu(self.Mconv5_stage6(h))
+        h = self.relu(self.Mconv6_stage6(h))
+        h = self.Mconv7_stage6(h)
+        heatmaps.append(h)
+
+        return heatmaps
+
+
+LOG = logging.getLogger(__name__)
+TOTEN = ToTensor()
+TOPIL = ToPILImage()
+
+
+params = {
+    'gaussian_sigma': 2.5,
+    'inference_img_size': 736,  # 368, 736, 1312
+    'heatmap_peak_thresh': 0.1,
+    'crop_scale': 1.5,
+    'line_indices': [
+        [0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6],
+        [6, 7], [7, 8], [8, 9], [9, 10], [10, 11], [11, 12], [12, 13],
+        [13, 14], [14, 15], [15, 16],
+        [17, 18], [18, 19], [19, 20], [20, 21],
+        [22, 23], [23, 24], [24, 25], [25, 26],
+        [27, 28], [28, 29], [29, 30],
+        [31, 32], [32, 33], [33, 34], [34, 35],
+        [36, 37], [37, 38], [38, 39], [39, 40], [40, 41], [41, 36],
+        [42, 43], [43, 44], [44, 45], [45, 46], [46, 47], [47, 42],
+        [48, 49], [49, 50], [50, 51], [51, 52], [52, 53], [53, 54],
+        [54, 55], [55, 56], [56, 57], [57, 58], [58, 59], [59, 48],
+        [60, 61], [61, 62], [62, 63], [63, 64], [64, 65], [65, 66],
+        [66, 67], [67, 60]
+    ],
+}
+
+
+class Face(object):
+    """
+    The OpenPose face landmark detector model.
+
+    Args:
+        inference_size: set the size of the inference image size, suggested:
+            368, 736, 1312, default 736
+        gaussian_sigma: blur the heatmaps, default 2.5
+        heatmap_peak_thresh: return landmark if over threshold, default 0.1
+
+    """
+    def __init__(self, face_model_path,
+                 inference_size=None,
+                 gaussian_sigma=None,
+                 heatmap_peak_thresh=None):
+        self.inference_size = inference_size or params["inference_img_size"]
+        self.sigma = gaussian_sigma or params['gaussian_sigma']
+        self.threshold = heatmap_peak_thresh or params["heatmap_peak_thresh"]
+        self.model = FaceNet()
+        self.model.load_state_dict(torch.load(face_model_path))
+        # if torch.cuda.is_available():
+        #     self.model = self.model.cuda()
+            # print('cuda')
+        self.model.eval()
+
+    def __call__(self, face_img):
+        H, W, C = face_img.shape
+
+        w_size = 384
+        x_data = torch.from_numpy(util.smart_resize(face_img, (w_size, w_size))).permute([2, 0, 1]) / 256.0 - 0.5
+
+        x_data = x_data.to(self.cn_device)
+
+        with torch.no_grad():
+            hs = self.model(x_data[None, ...])
+            heatmaps = F.interpolate(
+                hs[-1],
+                (H, W),
+                mode='bilinear', align_corners=True).cpu().numpy()[0]
+        return heatmaps
+
+    def compute_peaks_from_heatmaps(self, heatmaps):
+        all_peaks = []
+        for part in range(heatmaps.shape[0]):
+            map_ori = heatmaps[part].copy()
+            binary = np.ascontiguousarray(map_ori > 0.05, dtype=np.uint8)
+
+            if np.sum(binary) == 0:
+                continue
+
+            positions = np.where(binary > 0.5)
+            intensities = map_ori[positions]
+            mi = np.argmax(intensities)
+            y, x = positions[0][mi], positions[1][mi]
+            all_peaks.append([x, y])
+
+        return np.array(all_peaks)

dwpose/hand.py

@@ -0,0 +1,94 @@
+import cv2
+import json
+import numpy as np
+import math
+import time
+from scipy.ndimage.filters import gaussian_filter
+import matplotlib.pyplot as plt
+import matplotlib
+import torch
+from skimage.measure import label
+
+from .model import handpose_model
+from . import util
+
+class Hand(object):
+    def __init__(self, model_path):
+        self.model = handpose_model()
+        # if torch.cuda.is_available():
+        #     self.model = self.model.cuda()
+            # print('cuda')
+        model_dict = util.transfer(self.model, torch.load(model_path))
+        self.model.load_state_dict(model_dict)
+        self.model.eval()
+
+    def __call__(self, oriImgRaw):
+        scale_search = [0.5, 1.0, 1.5, 2.0]
+        # scale_search = [0.5]
+        boxsize = 368
+        stride = 8
+        padValue = 128
+        thre = 0.05
+        multiplier = [x * boxsize for x in scale_search]
+
+        wsize = 128
+        heatmap_avg = np.zeros((wsize, wsize, 22))
+
+        Hr, Wr, Cr = oriImgRaw.shape
+
+        oriImg = cv2.GaussianBlur(oriImgRaw, (0, 0), 0.8)
+
+        for m in range(len(multiplier)):
+            scale = multiplier[m]
+            imageToTest = util.smart_resize(oriImg, (scale, scale))
+
+            imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
+            im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
+            im = np.ascontiguousarray(im)
+
+            data = torch.from_numpy(im).float()
+            if torch.cuda.is_available():
+                data = data.cuda()
+
+            with torch.no_grad():
+                data = data.to(self.cn_device)
+                output = self.model(data).cpu().numpy()
+
+            # extract outputs, resize, and remove padding
+            heatmap = np.transpose(np.squeeze(output), (1, 2, 0))  # output 1 is heatmaps
+            heatmap = util.smart_resize_k(heatmap, fx=stride, fy=stride)
+            heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
+            heatmap = util.smart_resize(heatmap, (wsize, wsize))
+
+            heatmap_avg += heatmap / len(multiplier)
+
+        all_peaks = []
+        for part in range(21):
+            map_ori = heatmap_avg[:, :, part]
+            one_heatmap = gaussian_filter(map_ori, sigma=3)
+            binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8)
+
+            if np.sum(binary) == 0:
+                all_peaks.append([0, 0])
+                continue
+            label_img, label_numbers = label(binary, return_num=True, connectivity=binary.ndim)
+            max_index = np.argmax([np.sum(map_ori[label_img == i]) for i in range(1, label_numbers + 1)]) + 1
+            label_img[label_img != max_index] = 0
+            map_ori[label_img == 0] = 0
+
+            y, x = util.npmax(map_ori)
+            y = int(float(y) * float(Hr) / float(wsize))
+            x = int(float(x) * float(Wr) / float(wsize))
+            all_peaks.append([x, y])
+        return np.array(all_peaks)
+
+if __name__ == "__main__":
+    hand_estimation = Hand('../model/hand_pose_model.pth')
+
+    # test_image = '../images/hand.jpg'
+    test_image = '../images/hand.jpg'
+    oriImg = cv2.imread(test_image)  # B,G,R order
+    peaks = hand_estimation(oriImg)
+    canvas = util.draw_handpose(oriImg, peaks, True)
+    cv2.imshow('', canvas)
+    cv2.waitKey(0)

dwpose/model.py

@@ -0,0 +1,218 @@
+import torch
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+def make_layers(block, no_relu_layers):
+    layers = []
+    for layer_name, v in block.items():
+        if 'pool' in layer_name:
+            layer = nn.MaxPool2d(kernel_size=v[0], stride=v[1],
+                                    padding=v[2])
+            layers.append((layer_name, layer))
+        else:
+            conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1],
+                               kernel_size=v[2], stride=v[3],
+                               padding=v[4])
+            layers.append((layer_name, conv2d))
+            if layer_name not in no_relu_layers:
+                layers.append(('relu_'+layer_name, nn.ReLU(inplace=True)))
+
+    return nn.Sequential(OrderedDict(layers))
+
+class bodypose_model(nn.Module):
+    def __init__(self):
+        super(bodypose_model, self).__init__()
+
+        # these layers have no relu layer
+        no_relu_layers = ['conv5_5_CPM_L1', 'conv5_5_CPM_L2', 'Mconv7_stage2_L1',\
+                          'Mconv7_stage2_L2', 'Mconv7_stage3_L1', 'Mconv7_stage3_L2',\
+                          'Mconv7_stage4_L1', 'Mconv7_stage4_L2', 'Mconv7_stage5_L1',\
+                          'Mconv7_stage5_L2', 'Mconv7_stage6_L1', 'Mconv7_stage6_L1']
+        blocks = {}
+        block0 = OrderedDict([
+                      ('conv1_1', [3, 64, 3, 1, 1]),
+                      ('conv1_2', [64, 64, 3, 1, 1]),
+                      ('pool1_stage1', [2, 2, 0]),
+                      ('conv2_1', [64, 128, 3, 1, 1]),
+                      ('conv2_2', [128, 128, 3, 1, 1]),
+                      ('pool2_stage1', [2, 2, 0]),
+                      ('conv3_1', [128, 256, 3, 1, 1]),
+                      ('conv3_2', [256, 256, 3, 1, 1]),
+                      ('conv3_3', [256, 256, 3, 1, 1]),
+                      ('conv3_4', [256, 256, 3, 1, 1]),
+                      ('pool3_stage1', [2, 2, 0]),
+                      ('conv4_1', [256, 512, 3, 1, 1]),
+                      ('conv4_2', [512, 512, 3, 1, 1]),
+                      ('conv4_3_CPM', [512, 256, 3, 1, 1]),
+                      ('conv4_4_CPM', [256, 128, 3, 1, 1])
+                  ])
+
+
+        # Stage 1
+        block1_1 = OrderedDict([
+                        ('conv5_1_CPM_L1', [128, 128, 3, 1, 1]),
+                        ('conv5_2_CPM_L1', [128, 128, 3, 1, 1]),
+                        ('conv5_3_CPM_L1', [128, 128, 3, 1, 1]),
+                        ('conv5_4_CPM_L1', [128, 512, 1, 1, 0]),
+                        ('conv5_5_CPM_L1', [512, 38, 1, 1, 0])
+                    ])
+
+        block1_2 = OrderedDict([
+                        ('conv5_1_CPM_L2', [128, 128, 3, 1, 1]),
+                        ('conv5_2_CPM_L2', [128, 128, 3, 1, 1]),
+                        ('conv5_3_CPM_L2', [128, 128, 3, 1, 1]),
+                        ('conv5_4_CPM_L2', [128, 512, 1, 1, 0]),
+                        ('conv5_5_CPM_L2', [512, 19, 1, 1, 0])
+                    ])
+        blocks['block1_1'] = block1_1
+        blocks['block1_2'] = block1_2
+
+        self.model0 = make_layers(block0, no_relu_layers)
+
+        # Stages 2 - 6
+        for i in range(2, 7):
+            blocks['block%d_1' % i] = OrderedDict([
+                    ('Mconv1_stage%d_L1' % i, [185, 128, 7, 1, 3]),
+                    ('Mconv2_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv3_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv4_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv5_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv6_stage%d_L1' % i, [128, 128, 1, 1, 0]),
+                    ('Mconv7_stage%d_L1' % i, [128, 38, 1, 1, 0])
+                ])
+
+            blocks['block%d_2' % i] = OrderedDict([
+                    ('Mconv1_stage%d_L2' % i, [185, 128, 7, 1, 3]),
+                    ('Mconv2_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv3_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv4_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv5_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv6_stage%d_L2' % i, [128, 128, 1, 1, 0]),
+                    ('Mconv7_stage%d_L2' % i, [128, 19, 1, 1, 0])
+                ])
+
+        for k in blocks.keys():
+            blocks[k] = make_layers(blocks[k], no_relu_layers)
+
+        self.model1_1 = blocks['block1_1']
+        self.model2_1 = blocks['block2_1']
+        self.model3_1 = blocks['block3_1']
+        self.model4_1 = blocks['block4_1']
+        self.model5_1 = blocks['block5_1']
+        self.model6_1 = blocks['block6_1']
+
+        self.model1_2 = blocks['block1_2']
+        self.model2_2 = blocks['block2_2']
+        self.model3_2 = blocks['block3_2']
+        self.model4_2 = blocks['block4_2']
+        self.model5_2 = blocks['block5_2']
+        self.model6_2 = blocks['block6_2']
+
+
+    def forward(self, x):
+
+        out1 = self.model0(x)
+
+        out1_1 = self.model1_1(out1)
+        out1_2 = self.model1_2(out1)
+        out2 = torch.cat([out1_1, out1_2, out1], 1)
+
+        out2_1 = self.model2_1(out2)
+        out2_2 = self.model2_2(out2)
+        out3 = torch.cat([out2_1, out2_2, out1], 1)
+
+        out3_1 = self.model3_1(out3)
+        out3_2 = self.model3_2(out3)
+        out4 = torch.cat([out3_1, out3_2, out1], 1)
+
+        out4_1 = self.model4_1(out4)
+        out4_2 = self.model4_2(out4)
+        out5 = torch.cat([out4_1, out4_2, out1], 1)
+
+        out5_1 = self.model5_1(out5)
+        out5_2 = self.model5_2(out5)
+        out6 = torch.cat([out5_1, out5_2, out1], 1)
+
+        out6_1 = self.model6_1(out6)
+        out6_2 = self.model6_2(out6)
+
+        return out6_1, out6_2
+
+class handpose_model(nn.Module):
+    def __init__(self):
+        super(handpose_model, self).__init__()
+
+        # these layers have no relu layer
+        no_relu_layers = ['conv6_2_CPM', 'Mconv7_stage2', 'Mconv7_stage3',\
+                          'Mconv7_stage4', 'Mconv7_stage5', 'Mconv7_stage6']
+        # stage 1
+        block1_0 = OrderedDict([
+                ('conv1_1', [3, 64, 3, 1, 1]),
+                ('conv1_2', [64, 64, 3, 1, 1]),
+                ('pool1_stage1', [2, 2, 0]),
+                ('conv2_1', [64, 128, 3, 1, 1]),
+                ('conv2_2', [128, 128, 3, 1, 1]),
+                ('pool2_stage1', [2, 2, 0]),
+                ('conv3_1', [128, 256, 3, 1, 1]),
+                ('conv3_2', [256, 256, 3, 1, 1]),
+                ('conv3_3', [256, 256, 3, 1, 1]),
+                ('conv3_4', [256, 256, 3, 1, 1]),
+                ('pool3_stage1', [2, 2, 0]),
+                ('conv4_1', [256, 512, 3, 1, 1]),
+                ('conv4_2', [512, 512, 3, 1, 1]),
+                ('conv4_3', [512, 512, 3, 1, 1]),
+                ('conv4_4', [512, 512, 3, 1, 1]),
+                ('conv5_1', [512, 512, 3, 1, 1]),
+                ('conv5_2', [512, 512, 3, 1, 1]),
+                ('conv5_3_CPM', [512, 128, 3, 1, 1])
+            ])
+
+        block1_1 = OrderedDict([
+            ('conv6_1_CPM', [128, 512, 1, 1, 0]),
+            ('conv6_2_CPM', [512, 22, 1, 1, 0])
+        ])
+
+        blocks = {}
+        blocks['block1_0'] = block1_0
+        blocks['block1_1'] = block1_1
+
+        # stage 2-6
+        for i in range(2, 7):
+            blocks['block%d' % i] = OrderedDict([
+                    ('Mconv1_stage%d' % i, [150, 128, 7, 1, 3]),
+                    ('Mconv2_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv3_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv4_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv5_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv6_stage%d' % i, [128, 128, 1, 1, 0]),
+                    ('Mconv7_stage%d' % i, [128, 22, 1, 1, 0])
+                ])
+
+        for k in blocks.keys():
+            blocks[k] = make_layers(blocks[k], no_relu_layers)
+
+        self.model1_0 = blocks['block1_0']
+        self.model1_1 = blocks['block1_1']
+        self.model2 = blocks['block2']
+        self.model3 = blocks['block3']
+        self.model4 = blocks['block4']
+        self.model5 = blocks['block5']
+        self.model6 = blocks['block6']
+
+    def forward(self, x):
+        out1_0 = self.model1_0(x)
+        out1_1 = self.model1_1(out1_0)
+        concat_stage2 = torch.cat([out1_1, out1_0], 1)
+        out_stage2 = self.model2(concat_stage2)
+        concat_stage3 = torch.cat([out_stage2, out1_0], 1)
+        out_stage3 = self.model3(concat_stage3)
+        concat_stage4 = torch.cat([out_stage3, out1_0], 1)
+        out_stage4 = self.model4(concat_stage4)
+        concat_stage5 = torch.cat([out_stage4, out1_0], 1)
+        out_stage5 = self.model5(concat_stage5)
+        concat_stage6 = torch.cat([out_stage5, out1_0], 1)
+        out_stage6 = self.model6(concat_stage6)
+        return out_stage6
+

dwpose/types.py

@@ -0,0 +1,30 @@
+from typing import NamedTuple, List, Optional
+
+class Keypoint(NamedTuple):
+    x: float
+    y: float
+    score: float = 1.0
+    id: int = -1
+
+
+class BodyResult(NamedTuple):
+    # Note: Using `Optional` instead of `|` operator as the ladder is a Python
+    # 3.10 feature.
+    # Annotator code should be Python 3.8 Compatible, as controlnet repo uses
+    # Python 3.8 environment.
+    # https://github.com/lllyasviel/ControlNet/blob/d3284fcd0972c510635a4f5abe2eeb71dc0de524/environment.yaml#L6
+    keypoints: List[Optional[Keypoint]]
+    total_score: float = 0.0
+    total_parts: int = 0
+
+
+HandResult = List[Keypoint]
+FaceResult = List[Keypoint]
+AnimalPoseResult = List[Keypoint]
+
+
+class PoseResult(NamedTuple):
+    body: BodyResult
+    left_hand: Optional[HandResult]
+    right_hand: Optional[HandResult]
+    face: Optional[FaceResult]

dwpose/util.py

@@ -0,0 +1,466 @@
+import math
+import numpy as np
+import matplotlib
+import cv2
+import os
+from typing import List, Tuple, Union, Optional
+
+from .body import BodyResult, Keypoint
+
+eps = 0.01
+
+
+def smart_resize(x, s):
+    Ht, Wt = s
+    if x.ndim == 2:
+        Ho, Wo = x.shape
+        Co = 1
+    else:
+        Ho, Wo, Co = x.shape
+    if Co == 3 or Co == 1:
+        k = float(Ht + Wt) / float(Ho + Wo)
+        return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
+    else:
+        return np.stack([smart_resize(x[:, :, i], s) for i in range(Co)], axis=2)
+
+
+def smart_resize_k(x, fx, fy):
+    if x.ndim == 2:
+        Ho, Wo = x.shape
+        Co = 1
+    else:
+        Ho, Wo, Co = x.shape
+    Ht, Wt = Ho * fy, Wo * fx
+    if Co == 3 or Co == 1:
+        k = float(Ht + Wt) / float(Ho + Wo)
+        return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
+    else:
+        return np.stack([smart_resize_k(x[:, :, i], fx, fy) for i in range(Co)], axis=2)
+
+
+def padRightDownCorner(img, stride, padValue):
+    h = img.shape[0]
+    w = img.shape[1]
+
+    pad = 4 * [None]
+    pad[0] = 0 # up
+    pad[1] = 0 # left
+    pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
+    pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
+
+    img_padded = img
+    pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1))
+    img_padded = np.concatenate((pad_up, img_padded), axis=0)
+    pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1))
+    img_padded = np.concatenate((pad_left, img_padded), axis=1)
+    pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1))
+    img_padded = np.concatenate((img_padded, pad_down), axis=0)
+    pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1))
+    img_padded = np.concatenate((img_padded, pad_right), axis=1)
+
+    return img_padded, pad
+
+
+def transfer(model, model_weights):
+    transfered_model_weights = {}
+    for weights_name in model.state_dict().keys():
+        transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])]
+    return transfered_model_weights
+
+
+def is_normalized(keypoints: List[Optional[Keypoint]]) -> bool:
+    point_normalized = [
+        0 <= abs(k.x) <= 1 and 0 <= abs(k.y) <= 1 
+        for k in keypoints 
+        if k is not None
+    ]
+    if not point_normalized:
+        return False
+    return all(point_normalized)
+
+    
+def draw_bodypose(canvas: np.ndarray, keypoints: List[Keypoint], xinsr_stick_scaling: bool = False) -> np.ndarray:
+    """
+    Draw keypoints and limbs representing body pose on a given canvas.
+
+    Args:
+        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the body pose.
+        keypoints (List[Keypoint]): A list of Keypoint objects representing the body keypoints to be drawn.
+        xinsr_stick_scaling (bool): Whether or not scaling stick width for xinsr ControlNet
+
+    Returns:
+        np.ndarray: A 3D numpy array representing the modified canvas with the drawn body pose.
+
+    Note:
+        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
+    """
+    if not is_normalized(keypoints):
+        H, W = 1.0, 1.0
+    else:
+        H, W, _ = canvas.shape
+
+    CH, CW, _ = canvas.shape
+    stickwidth = 4
+
+    # Ref: https://huggingface.co/xinsir/controlnet-openpose-sdxl-1.0
+    max_side = max(CW, CH)
+    if xinsr_stick_scaling:
+        stick_scale = 1 if max_side < 500 else min(2 + (max_side // 1000), 7)
+    else:
+        stick_scale = 1
+
+    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],
+    ]
+
+    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 (k1_index, k2_index), color in zip(limbSeq, colors):
+        keypoint1 = keypoints[k1_index - 1]
+        keypoint2 = keypoints[k2_index - 1]
+
+        if keypoint1 is None or keypoint2 is None:
+            continue
+
+        Y = np.array([keypoint1.x, keypoint2.x]) * float(W)
+        X = np.array([keypoint1.y, keypoint2.y]) * 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*stick_scale), int(angle), 0, 360, 1)
+        cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])
+
+    for keypoint, color in zip(keypoints, colors):
+        if keypoint is None:
+            continue
+
+        x, y = keypoint.x, keypoint.y
+        x = int(x * W)
+        y = int(y * H)
+        cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
+
+    return canvas
+
+
+def draw_handpose(canvas: np.ndarray, keypoints: Union[List[Keypoint], None]) -> np.ndarray:
+    """
+    Draw keypoints and connections representing hand pose on a given canvas.
+
+    Args:
+        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
+        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
+                                          or None if no keypoints are present.
+
+    Returns:
+        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.
+
+    Note:
+        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
+    """
+    if not keypoints:
+        return canvas
+
+    if not is_normalized(keypoints):
+        H, W = 1.0, 1.0
+    else:
+        H, W, _ = 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 ie, (e1, e2) in enumerate(edges):
+        k1 = keypoints[e1]
+        k2 = keypoints[e2]
+        if k1 is None or k2 is None:
+            continue
+
+        x1 = int(k1.x * W)
+        y1 = int(k1.y * H)
+        x2 = int(k2.x * W)
+        y2 = int(k2.y * H)
+        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]) * 255, thickness=2)
+
+    for keypoint in keypoints:
+        if keypoint is None:
+            continue
+
+        x, y = keypoint.x, keypoint.y
+        x = int(x * W)
+        y = int(y * H)
+        if x > eps and y > eps:
+            cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
+    return canvas
+
+
+def draw_facepose(canvas: np.ndarray, keypoints: Union[List[Keypoint], None]) -> np.ndarray:
+    """
+    Draw keypoints representing face pose on a given canvas.
+
+    Args:
+        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the face pose.
+        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the face keypoints to be drawn
+                                          or None if no keypoints are present.
+
+    Returns:
+        np.ndarray: A 3D numpy array representing the modified canvas with the drawn face pose.
+
+    Note:
+        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
+    """
+    if not keypoints:
+        return canvas
+
+    if not is_normalized(keypoints):
+        H, W = 1.0, 1.0
+    else:
+        H, W, _ = canvas.shape
+
+    for keypoint in keypoints:
+        if keypoint is None:
+            continue
+
+        x, y = keypoint.x, keypoint.y
+        x = int(x * W)
+        y = int(y * H)
+        if x > eps and y > eps:
+            cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
+    return canvas
+
+
+# detect hand according to body pose keypoints
+# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
+def handDetect(body: BodyResult, oriImg) -> List[Tuple[int, int, int, bool]]:
+    """
+    Detect hands in the input body pose keypoints and calculate the bounding box for each hand.
+
+    Args:
+        body (BodyResult): A BodyResult object containing the detected body pose keypoints.
+        oriImg (numpy.ndarray): A 3D numpy array representing the original input image.
+
+    Returns:
+        List[Tuple[int, int, int, bool]]: A list of tuples, each containing the coordinates (x, y) of the top-left
+                                          corner of the bounding box, the width (height) of the bounding box, and
+                                          a boolean flag indicating whether the hand is a left hand (True) or a
+                                          right hand (False).
+
+    Notes:
+        - The width and height of the bounding boxes are equal since the network requires squared input.
+        - The minimum bounding box size is 20 pixels.
+    """
+    ratioWristElbow = 0.33
+    detect_result = []
+    image_height, image_width = oriImg.shape[0:2]
+    
+    keypoints = body.keypoints
+    # right hand: wrist 4, elbow 3, shoulder 2
+    # left hand: wrist 7, elbow 6, shoulder 5
+    left_shoulder = keypoints[5]
+    left_elbow = keypoints[6]
+    left_wrist = keypoints[7]
+    right_shoulder = keypoints[2]
+    right_elbow = keypoints[3]
+    right_wrist = keypoints[4]
+
+    # if any of three not detected
+    has_left = all(keypoint is not None for keypoint in (left_shoulder, left_elbow, left_wrist))
+    has_right = all(keypoint is not None for keypoint in (right_shoulder, right_elbow, right_wrist))
+    if not (has_left or has_right):
+        return []
+    
+    hands = []
+    #left hand
+    if has_left:
+        hands.append([
+            left_shoulder.x, left_shoulder.y,
+            left_elbow.x, left_elbow.y,
+            left_wrist.x, left_wrist.y,
+            True
+        ])
+    # right hand
+    if has_right:
+        hands.append([
+            right_shoulder.x, right_shoulder.y,
+            right_elbow.x, right_elbow.y,
+            right_wrist.x, right_wrist.y,
+            False
+        ])
+
+    for x1, y1, x2, y2, x3, y3, is_left in hands:
+        # pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox
+        # handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]);
+        # handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]);
+        # const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow);
+        # const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder);
+        # handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder);
+        x = x3 + ratioWristElbow * (x3 - x2)
+        y = y3 + ratioWristElbow * (y3 - y2)
+        distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
+        distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+        width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
+        # x-y refers to the center --> offset to topLeft point
+        # handRectangle.x -= handRectangle.width / 2.f;
+        # handRectangle.y -= handRectangle.height / 2.f;
+        x -= width / 2
+        y -= width / 2  # width = height
+        # overflow the image
+        if x < 0: x = 0
+        if y < 0: y = 0
+        width1 = width
+        width2 = width
+        if x + width > image_width: width1 = image_width - x
+        if y + width > image_height: width2 = image_height - y
+        width = min(width1, width2)
+        # the max hand box value is 20 pixels
+        if width >= 20:
+            detect_result.append((int(x), int(y), int(width), is_left))
+
+    '''
+    return value: [[x, y, w, True if left hand else False]].
+    width=height since the network require squared input.
+    x, y is the coordinate of top left 
+    '''
+    return detect_result
+
+
+# Written by Lvmin
+def faceDetect(body: BodyResult, oriImg) -> Union[Tuple[int, int, int], None]:
+    """
+    Detect the face in the input body pose keypoints and calculate the bounding box for the face.
+
+    Args:
+        body (BodyResult): A BodyResult object containing the detected body pose keypoints.
+        oriImg (numpy.ndarray): A 3D numpy array representing the original input image.
+
+    Returns:
+        Tuple[int, int, int] | None: A tuple containing the coordinates (x, y) of the top-left corner of the
+                                   bounding box and the width (height) of the bounding box, or None if the
+                                   face is not detected or the bounding box width is less than 20 pixels.
+
+    Notes:
+        - The width and height of the bounding box are equal.
+        - The minimum bounding box size is 20 pixels.
+    """
+    # left right eye ear 14 15 16 17
+    image_height, image_width = oriImg.shape[0:2]
+    
+    keypoints = body.keypoints
+    head = keypoints[0]
+    left_eye = keypoints[14]
+    right_eye = keypoints[15]
+    left_ear = keypoints[16]
+    right_ear = keypoints[17]
+    
+    if head is None or all(keypoint is None for keypoint in (left_eye, right_eye, left_ear, right_ear)):
+        return None
+
+    width = 0.0
+    x0, y0 = head.x, head.y
+
+    if left_eye is not None:
+        x1, y1 = left_eye.x, left_eye.y
+        d = max(abs(x0 - x1), abs(y0 - y1))
+        width = max(width, d * 3.0)
+
+    if right_eye is not None:
+        x1, y1 = right_eye.x, right_eye.y
+        d = max(abs(x0 - x1), abs(y0 - y1))
+        width = max(width, d * 3.0)
+
+    if left_ear is not None:
+        x1, y1 = left_ear.x, left_ear.y
+        d = max(abs(x0 - x1), abs(y0 - y1))
+        width = max(width, d * 1.5)
+
+    if right_ear is not None:
+        x1, y1 = right_ear.x, right_ear.y
+        d = max(abs(x0 - x1), abs(y0 - y1))
+        width = max(width, d * 1.5)
+
+    x, y = x0, y0
+
+    x -= width
+    y -= width
+
+    if x < 0:
+        x = 0
+
+    if y < 0:
+        y = 0
+
+    width1 = width * 2
+    width2 = width * 2
+
+    if x + width > image_width:
+        width1 = image_width - x
+
+    if y + width > image_height:
+        width2 = image_height - y
+
+    width = min(width1, width2)
+
+    if width >= 20:
+        return int(x), int(y), int(width)
+    else:
+        return None
+
+
+# get max index of 2d array
+def npmax(array):
+    arrayindex = array.argmax(1)
+    arrayvalue = array.max(1)
+    i = arrayvalue.argmax()
+    j = arrayindex[i]
+    return i, j
+
+def guess_onnx_input_shape_dtype(filename):
+    dtype = np.float32
+    if "fp16" in filename:
+        dtype = np.float16
+    elif "int8" in filename:
+        dtype = np.uint8
+    input_size = (640, 640) if "yolo" in filename else (192, 256)
+    if "384" in filename:
+        input_size = (288, 384)
+    elif "256" in filename:
+        input_size = (256, 256)
+    return input_size, dtype
+
+if os.getenv('AUX_ORT_PROVIDERS'):
+    ONNX_PROVIDERS = os.getenv('AUX_ORT_PROVIDERS').split(',')
+else:
+    ONNX_PROVIDERS = ["CUDAExecutionProvider", "DirectMLExecutionProvider", "OpenVINOExecutionProvider", "ROCMExecutionProvider", "CPUExecutionProvider"]
+def get_ort_providers() -> List[str]:
+    providers = []
+    try:
+        import onnxruntime as ort
+        for provider in ONNX_PROVIDERS:
+            if provider in ort.get_available_providers():
+                providers.append(provider)
+        return providers
+    except:
+        return []
+
+def is_model_torchscript(model) -> bool:
+    return bool(type(model).__name__ == "RecursiveScriptModule")
+
+def get_model_type(Nodesname, filename) -> str:
+    ort_providers = list(filter(lambda x : x != "CPUExecutionProvider", get_ort_providers()))
+    if filename is None:
+        return None
+    elif ("onnx" in filename) and ort_providers:
+        print(f"{Nodesname}: Caching ONNXRuntime session {filename}...")
+        return "ort"
+    elif ("onnx" in filename):
+        print(f"{Nodesname}: Caching OpenCV DNN module {filename} on cv2.DNN...")
+        return "cv2"
+    else:
+        print(f"{Nodesname}: Caching TorchScript module {filename} on ...")
+        return  "torchscript"

dwpose/wholebody.py

@@ -0,0 +1,181 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import cv2
+import numpy as np
+
+from .dw_onnx.cv_ox_det import inference_detector as inference_onnx_yolox
+from .dw_onnx.cv_ox_yolo_nas import inference_detector as inference_onnx_yolo_nas
+from .dw_onnx.cv_ox_pose import inference_pose as inference_onnx_pose
+
+from .dw_torchscript.jit_det import inference_detector as inference_jit_yolox
+from .dw_torchscript.jit_pose import inference_pose as inference_jit_pose
+
+from typing import List, Optional
+from .types import PoseResult, BodyResult, Keypoint
+from timeit import default_timer
+import os
+from custom_controlnet_aux.dwpose.util import guess_onnx_input_shape_dtype, get_model_type, get_ort_providers, is_model_torchscript
+import torch
+
+class Wholebody:
+    def __init__(self, det_model_path: Optional[str] = None, pose_model_path: Optional[str] = None, torchscript_device="cuda"):
+        self.det_filename = det_model_path and os.path.basename(det_model_path)
+        self.pose_filename = pose_model_path and os.path.basename(pose_model_path)
+        self.det, self.pose = None, None
+        # return type: None ort cv2 torchscript
+        self.det_model_type = get_model_type("DWPose",self.det_filename)
+        self.pose_model_type = get_model_type("DWPose",self.pose_filename)
+        # Always loads to CPU to avoid building OpenCV.
+        cv2_device = 'cpu'
+        cv2_backend = cv2.dnn.DNN_BACKEND_OPENCV if cv2_device == 'cpu' else cv2.dnn.DNN_BACKEND_CUDA
+        # You need to manually build OpenCV through cmake to work with your GPU.
+        cv2_providers = cv2.dnn.DNN_TARGET_CPU if cv2_device == 'cpu' else cv2.dnn.DNN_TARGET_CUDA
+        ort_providers = get_ort_providers()
+
+        if self.det_model_type is None:
+            pass
+        elif self.det_model_type == "ort":
+            try:
+                import onnxruntime as ort
+                self.det = ort.InferenceSession(det_model_path, providers=ort_providers)
+            except:
+                print(f"Failed to load onnxruntime with {self.det.get_providers()}.\nPlease change EP_list in the config.yaml and restart ComfyUI")
+                self.det = ort.InferenceSession(det_model_path, providers=["CPUExecutionProvider"])
+        elif self.det_model_type == "cv2":
+            try:
+                self.det = cv2.dnn.readNetFromONNX(det_model_path)
+                self.det.setPreferableBackend(cv2_backend)
+                self.det.setPreferableTarget(cv2_providers)
+            except:
+                print("TopK operators may not work on your OpenCV, try use onnxruntime with CPUExecutionProvider")
+                try:
+                    import onnxruntime as ort
+                    self.det = ort.InferenceSession(det_model_path, providers=["CPUExecutionProvider"])
+                except:
+                    print(f"Failed to load {det_model_path}, you can use other models instead")
+        else:
+            self.det = torch.jit.load(det_model_path)
+            self.det.to(torchscript_device)
+
+        if self.pose_model_type is None:
+            pass
+        elif self.pose_model_type == "ort":
+            try:
+                import onnxruntime as ort
+                self.pose = ort.InferenceSession(pose_model_path, providers=ort_providers)
+            except:
+                print(f"Failed to load onnxruntime with {self.pose.get_providers()}.\nPlease change EP_list in the config.yaml and restart ComfyUI")
+                self.pose = ort.InferenceSession(pose_model_path, providers=["CPUExecutionProvider"])
+        elif self.pose_model_type == "cv2":
+            self.pose = cv2.dnn.readNetFromONNX(pose_model_path)
+            self.pose.setPreferableBackend(cv2_backend)
+            self.pose.setPreferableTarget(cv2_providers)
+        else:
+            self.pose = torch.jit.load(pose_model_path)
+            self.pose.to(torchscript_device)
+        
+        if self.pose_filename is not None:
+            self.pose_input_size, _ = guess_onnx_input_shape_dtype(self.pose_filename)
+
+    def __call__(self, oriImg) -> Optional[np.ndarray]:
+        #Sacrifice accurate time measurement for compatibility 
+        
+        det_result = None
+        
+        if self.det is None:
+            print("DWPose: No detector specified, using full image for pose estimation.") # pragma: no cover
+            det_result = []
+        else:
+            det_start = default_timer()
+            if is_model_torchscript(self.det):
+                det_result = inference_jit_yolox(self.det, oriImg, detect_classes=[0])
+            else:
+                if "yolox" in self.det_filename:
+                    det_result = inference_onnx_yolox(self.det, oriImg, detect_classes=[0], dtype=np.float32)
+                else:
+                    #FP16 and INT8 YOLO NAS accept uint8 input
+                    det_result = inference_onnx_yolo_nas(self.det, oriImg, detect_classes=[0], dtype=np.uint8)
+            print(f"DWPose: Bbox {((default_timer() - det_start) * 1000):.2f}ms")
+            if (det_result is None) or (det_result.shape[0] == 0):
+                return None
+
+        pose_start = default_timer()
+        if is_model_torchscript(self.pose):
+            keypoints, scores = inference_jit_pose(self.pose, det_result, oriImg, self.pose_input_size)
+        else:
+            _, pose_onnx_dtype = guess_onnx_input_shape_dtype(self.pose_filename)
+            keypoints, scores = inference_onnx_pose(self.pose, det_result, oriImg, self.pose_input_size, dtype=pose_onnx_dtype)
+            
+        num_subjects_log = 'full image'
+        if hasattr(det_result, 'shape') and det_result.shape[0] > 0:
+            num_subjects_log = f"{det_result.shape[0]} people"
+        elif isinstance(det_result, list) and len(det_result) > 0 and isinstance(det_result[0], (list, np.ndarray)):
+             num_subjects_log = f"{len(det_result)} people"
+        
+        print(f"DWPose: Pose {((default_timer() - pose_start) * 1000):.2f}ms on {num_subjects_log}\n")
+
+        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
+
+        return keypoints_info
+
+    @staticmethod
+    def format_result(keypoints_info: Optional[np.ndarray]) -> List[PoseResult]:
+        def format_keypoint_part(
+            part: np.ndarray,
+        ) -> Optional[List[Optional[Keypoint]]]:
+            keypoints = [
+                Keypoint(x, y, score, i) if score >= 0.3 else None
+                for i, (x, y, score) in enumerate(part)
+            ]
+            return (
+                None if all(keypoint is None for keypoint in keypoints) else keypoints
+            )
+
+        def total_score(keypoints: Optional[List[Optional[Keypoint]]]) -> float:
+            return (
+                sum(keypoint.score for keypoint in keypoints if keypoint is not None)
+                if keypoints is not None
+                else 0.0
+            )
+
+        pose_results = []
+        if keypoints_info is None:
+            return pose_results
+
+        for instance in keypoints_info:
+            body_keypoints = format_keypoint_part(instance[:18]) or ([None] * 18)
+            left_hand = format_keypoint_part(instance[92:113])
+            right_hand = format_keypoint_part(instance[113:134])
+            face = format_keypoint_part(instance[24:92])
+
+            # Openpose face consists of 70 points in total, while DWPose only
+            # provides 68 points. Padding the last 2 points.
+            if face is not None:
+                # left eye
+                face.append(body_keypoints[14])
+                # right eye
+                face.append(body_keypoints[15])
+
+            body = BodyResult(
+                body_keypoints, total_score(body_keypoints), len(body_keypoints)
+            )
+            pose_results.append(PoseResult(body, left_hand, right_hand, face))
+
+        return pose_results

frame_cau.py

@@ -0,0 +1,107 @@
+import json
+import numpy as np
+
+class ApplyPoseScalesToFrames:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "frame_list": ("POSE_KEYPOINT",),  
+                "scales":     ("SCALES",),         
+                "ref_pose":   ("POSE_KEYPOINT",),  
+            }
+        }
+
+    RETURN_TYPES   = ("POSE_KEYPOINT",)
+    RETURN_NAMES   = ("new_frame_list",)
+    FUNCTION       = "apply_scales"
+    CATEGORY       = "Snap Processing"
+
+
+    NECK = 1
+    L_SHOULDER, L_ELBOW, L_WRIST = 2, 3, 4
+    R_SHOULDER, R_ELBOW, R_WRIST = 5, 6, 7
+    L_HIP, L_KNEE, L_ANKLE       = 8, 9, 10
+    R_HIP, R_KNEE, R_ANKLE       = 11, 12, 13
+
+
+    BONE_STEPS = [
+        ("neck-lshoulder", NECK, L_SHOULDER, [L_SHOULDER, L_ELBOW, L_WRIST]),
+        ("lshoulder-lelbow", L_SHOULDER, L_ELBOW, [L_ELBOW, L_WRIST]),
+        ("lelbow-lwrist", L_ELBOW, L_WRIST, [L_WRIST]),
+        ("neck-rshoulder", NECK, R_SHOULDER, [R_SHOULDER, R_ELBOW, R_WRIST]),
+        ("rshoulder-relbow", R_SHOULDER, R_ELBOW, [R_ELBOW, R_WRIST]),
+        ("relbow-rwrist", R_ELBOW, R_WRIST, [R_WRIST]),
+        ("neck-lhip", NECK, L_HIP, [L_HIP, L_KNEE, L_ANKLE]),
+        ("lhip-lknee", L_HIP, L_KNEE, [L_KNEE, L_ANKLE]),
+        ("lknee-lankle", L_KNEE, L_ANKLE, [L_ANKLE]),
+        ("neck-rhip", NECK, R_HIP, [R_HIP, R_KNEE, R_ANKLE]),
+        ("rhip-rknee", R_HIP, R_KNEE, [R_KNEE, R_ANKLE]),
+        ("rknee-rankle", R_KNEE, R_ANKLE, [R_ANKLE]),
+    ]
+
+
+    @staticmethod
+    def move_branch(kp, parent, child, scale, subtree):
+        p = kp[parent, :2].copy()
+        c = kp[child,  :2].copy()
+        new_c = p + (c - p) * scale
+        offset = new_c - c
+        kp[child, :2] = new_c
+        for idx in subtree[1:]:
+            kp[idx, :2] += offset
+
+
+    def apply_scales(self, frame_list, scales, ref_pose):
+        smooth_alpha = 0.2               
+
+
+        ref_json = ref_pose[0] if isinstance(ref_pose, list) else ref_pose
+        if isinstance(ref_json, str):
+            ref_json = json.loads(ref_json)
+        ref_kp = np.array(
+            ref_json["people"][0]["pose_keypoints_2d"], dtype=np.float32
+        ).reshape(-1, 3)
+
+        delta_map = None          
+        prev_smoothed = None      
+
+        new_frames = []
+        for frame_idx, frame in enumerate(frame_list):
+            person = frame["people"][0]
+            kp = np.array(person["pose_keypoints_2d"], dtype=np.float32).reshape(-1, 3)
+
+
+            for name, parent, child, subtree in self.BONE_STEPS:
+                self.move_branch(
+                    kp, parent, child,
+                    scales.get(name, 1.0), subtree
+                )
+
+
+            if delta_map is None:
+                delta_map = ref_kp[:, :2] - kp[:, :2]   
+
+
+            kp[:, :2] += delta_map
+
+
+            if prev_smoothed is not None:
+                kp[:, :2] = prev_smoothed[:, :2] + smooth_alpha * (
+                    kp[:, :2] - prev_smoothed[:, :2]
+                )
+            prev_smoothed = kp.copy()
+
+
+            person["pose_keypoints_2d"]      = kp.reshape(-1).tolist()
+            person["face_keypoints_2d"]      = []
+            person["hand_left_keypoints_2d"] = []
+            person["hand_right_keypoints_2d"] = []
+            new_frames.append(frame)
+
+        return (new_frames,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "ApplyPoseScalesToFrames": ApplyPoseScalesToFrames,
+}

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "Comfyui_SSsnap_pose-Remapping"
+description = "A pose remapping node with support for joint locking, motion-aware scaling, and optional easing - perfect for animation refinement and mech rig control."
+
+version = "1.0.0"
+
+[project.urls]
+Repository = "https://github.com/SS-snap/Comfyui_SSsnap_pose-Remapping"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "snap-demo"
+DisplayName = "Comfyui_SSsnap_pose-Remapping"
+Icon = ""

ratio_cau.py

@@ -0,0 +1,105 @@
+import json
+import numpy as np
+
+class CalculatePoseScales:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "ref_json": ("POSE_KEYPOINT",), 
+                "raw_json": ("POSE_KEYPOINT",), 
+            }
+        }
+
+    RETURN_TYPES = ("SCALES",)
+    RETURN_NAMES = ("scales",)
+    FUNCTION = "calculate_scales"
+    CATEGORY = "Snap Processing"
+
+    BONES = [
+        ("neck-lshoulder", 1, 2),
+        ("lshoulder-lelbow", 2, 3),
+        ("lelbow-lwrist", 3, 4),
+        ("neck-rshoulder", 1, 5),
+        ("rshoulder-relbow", 5, 6),
+        ("relbow-rwrist", 6, 7),
+        ("neck-lhip", 1, 8),
+        ("lhip-lknee", 8, 9),
+        ("lknee-lankle", 9, 10),
+        ("neck-rhip", 1, 11),
+        ("rhip-rknee", 11, 12),
+        ("rknee-rankle", 12, 13),
+    ]
+
+    def _to_kps_array(self, data):
+        """
+        将各种可能的输入:
+          - JSON 字符串 (dict 或 list of dict)
+          - Python dict
+          - list of frame-dicts
+          - list of floats 或 list of [x,y,c]
+        统一转成 (N,3) 的 NumPy array。
+        """
+
+        if isinstance(data, str):
+            data = json.loads(data)
+
+
+        if isinstance(data, dict):
+            data = [data]
+
+
+        if isinstance(data, list) and data and isinstance(data[0], dict):
+            frame = data[0]
+            people = frame.get("people")
+            if not isinstance(people, list) or not people:
+                raise ValueError("frame 中找不到 people 列表")
+            kp_list = people[0].get("pose_keypoints_2d")
+            if not isinstance(kp_list, list):
+                raise ValueError("people[0] 中找不到 pose_keypoints_2d")
+            arr = np.array(kp_list, dtype=np.float32)
+
+
+        elif isinstance(data, list) and data and isinstance(data[0], (int, float)):
+            arr = np.array(data, dtype=np.float32)
+
+
+        elif isinstance(data, list) and data and isinstance(data[0], list):
+            arr = np.array(data, dtype=np.float32).reshape(-1, 3)
+
+        else:
+            raise ValueError(f"无法识别的 keypoints 输入格式: {type(data)}")
+
+
+        if arr.ndim == 1:
+            arr = arr.reshape(-1, 3)
+        if arr.shape[1] != 3:
+            raise ValueError(f"关键点格式错误，应为 N×3，但得到 shape={arr.shape}")
+        return arr
+
+    def _segment_length(self, kps, i, j):
+        a, b = kps[i, :2], kps[j, :2]
+
+        if np.any(a == 0) or np.any(b == 0):
+            return np.nan
+        return float(np.linalg.norm(a - b))
+
+    def calculate_scales(self, ref_json, raw_json):
+
+        kps_ref = self._to_kps_array(ref_json)
+        kps_raw = self._to_kps_array(raw_json)
+
+        scales = {}
+        for name, i, j in self.BONES:
+            len_ref = self._segment_length(kps_ref, i, j)
+            len_raw = self._segment_length(kps_raw, i, j)
+            if np.isnan(len_ref) or np.isnan(len_raw) or len_raw == 0:
+                scales[name] = 1.0
+            else:
+                scales[name] = len_ref / len_raw
+        return (scales,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "CalculatePoseScales": CalculatePoseScales,
+}

render_kps.py

@@ -0,0 +1,58 @@
+import json
+import math
+import cv2
+import numpy as np
+import torch
+from collections import namedtuple
+from typing import List, Union, NamedTuple
+from .dwpose import *
+
+def numpy2torch(np_image: np.ndarray) -> torch.Tensor:
+    """ [H, W, C] => [B=1, H, W, C]"""
+    return torch.from_numpy(np_image.astype(np.float32) / 255).unsqueeze(0)
+
+class RenderKps:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "kps": ("POSE_KEYPOINT",),
+                "render_body": ("BOOLEAN", {"default": True}),
+            }
+        }
+
+    RETURN_TYPES   = ("IMAGE",)
+    FUNCTION = "render"
+    CATEGORY = "Snap Processing"
+
+    def render(self, kps, render_body):
+        frames = kps if isinstance(kps, list) else [kps]
+
+        imgs = []
+        for data in frames:
+            if isinstance(data, list):
+                data = data[0]
+            if isinstance(data, str):
+                try:
+                    data = json.loads(data)
+                except json.JSONDecodeError:
+                    continue
+            if not isinstance(data, dict):
+                continue
+
+            poses, _, H, W = decode_json_as_poses(data)
+            np_img = draw_poses(poses, H, W, render_body)
+            imgs.append(torch.from_numpy(np_img.astype(np.float32) / 255.0))  # [H,W,3]
+
+        if not imgs:
+            imgs = [torch.zeros((1, 1, 3), dtype=torch.float32)]
+
+
+        batch = torch.stack(imgs, dim=0)  
+        print("batch shape:", batch.shape)  
+
+        return (batch,)
+
+NODE_CLASS_MAPPINGS = {
+    "RenderKps": RenderKps,
+}

test.py

@@ -0,0 +1,51 @@
+# show.py
+import json
+import numpy as np
+import torch
+from PIL import Image, ImageDraw
+
+# 连接关节的骨架列表，参考 BODY_25
+POSE_PAIRS = [
+    (0,1),(0,15),(0,16),(15,17),(16,18),
+    (1,2),(1,5),(1,8),(2,3),(3,4),(5,6),(6,7),
+    (8,9),(8,12),(9,10),(10,11),(11,22),(11,24),
+    (22,23),(12,13),(13,14),(14,21),(14,19),(19,20)
+]
+
+def decode_json_to_batch(json_batch, canvas_size=(512,768)):
+    imgs = []
+    for rec in json_batch:
+        width, height = rec.get('canvas_width',canvas_size[0]), rec.get('canvas_height',canvas_size[1])
+        img = Image.new("RGB", (width, height), "white")
+        draw = ImageDraw.Draw(img)
+        for person in rec.get('people', []):
+            pts = person.get('pose_keypoints_2d', [])
+            coords = [(pts[i], pts[i+1]) for i in range(0, len(pts), 3)]
+            # 绘制点
+            for x,y in coords:
+                draw.ellipse((x-3, y-3, x+3, y+3), fill="red")
+            # 绘制骨架
+            for a,b in POSE_PAIRS:
+                if a < len(coords) and b < len(coords):
+                    xa, ya = coords[a]
+                    xb, yb = coords[b]
+                    draw.line((xa, ya, xb, yb), fill="blue", width=2)
+        # 转为 Tensor (H,W,C)->(C,H,W) 并归一化
+        arr = np.array(img).astype(np.float32) / 255.0
+        t = torch.from_numpy(arr).permute(2,0,1)
+        imgs.append(t)
+    if not imgs:
+        return torch.zeros(1,3,1,1)
+    return torch.stack(imgs, dim=0)
+
+if __name__ == "__main__":
+    # 示例输入，可替换为文件读取
+    sample = [
+        # （此处插入你提供的 JSON 对象列表）
+    ]
+    batch = decode_json_to_batch(sample)
+    print("Batch shape:", batch.shape)  # 比如 (5,3,768,512)
+    # 可将每张图保存为 PNG
+    for i, im in enumerate(batch):
+        np_img = (im.permute(1,2,0).numpy()*255).astype(np.uint8)
+        Image.fromarray(np_img).save(f"out_{i:03}.png")

tools.py

@@ -0,0 +1,350 @@
+import os
+import random
+import tempfile
+import warnings
+from contextlib import suppress
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+from huggingface_hub import constants, hf_hub_download
+from torch.utils.model_zoo import load_url
+from ast import literal_eval
+
+
+HF_MODEL_NAME = "lllyasviel/Annotators"
+DWPOSE_MODEL_NAME = "yzd-v/DWPose"
+BDS_MODEL_NAME = "bdsqlsz/qinglong_controlnet-lllite"
+DENSEPOSE_MODEL_NAME = "LayerNorm/DensePose-TorchScript-with-hint-image"
+MESH_GRAPHORMER_MODEL_NAME = "hr16/ControlNet-HandRefiner-pruned"
+SAM_MODEL_NAME = "dhkim2810/MobileSAM"
+UNIMATCH_MODEL_NAME = "hr16/Unimatch"
+DEPTH_ANYTHING_MODEL_NAME = "LiheYoung/Depth-Anything" #HF Space
+DIFFUSION_EDGE_MODEL_NAME = "hr16/Diffusion-Edge"
+METRIC3D_MODEL_NAME = "JUGGHM/Metric3D"
+
+DEPTH_ANYTHING_V2_MODEL_NAME_DICT = {
+    "depth_anything_v2_vits.pth": "depth-anything/Depth-Anything-V2-Small",
+    "depth_anything_v2_vitb.pth": "depth-anything/Depth-Anything-V2-Base",
+    "depth_anything_v2_vitl.pth": "depth-anything/Depth-Anything-V2-Large",
+    "depth_anything_v2_vitg.pth": "depth-anything/Depth-Anything-V2-Giant",
+    "depth_anything_v2_metric_vkitti_vitl.pth": "depth-anything/Depth-Anything-V2-Metric-VKITTI-Large",
+    "depth_anything_v2_metric_hypersim_vitl.pth": "depth-anything/Depth-Anything-V2-Metric-Hypersim-Large"
+}
+
+temp_dir = tempfile.gettempdir()
+annotator_ckpts_path = os.path.join(Path(__file__).parents[2], 'ckpts')
+USE_SYMLINKS = False
+
+try:
+    annotator_ckpts_path = os.environ['AUX_ANNOTATOR_CKPTS_PATH']
+except:
+    warnings.warn("Custom pressesor model path not set successfully.")
+    pass
+
+try:
+    USE_SYMLINKS = literal_eval(os.environ['AUX_USE_SYMLINKS'])
+except:
+    warnings.warn("USE_SYMLINKS not set successfully. Using default value: False to download models.")
+    pass
+
+try:
+    temp_dir = os.environ['AUX_TEMP_DIR']
+    if len(temp_dir) >= 60:
+        warnings.warn(f"custom temp dir is too long. Using default")
+        temp_dir = tempfile.gettempdir()
+except:
+    warnings.warn(f"custom temp dir not set successfully")
+    pass
+
+here = Path(__file__).parent.resolve()
+
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
+
+
+def make_noise_disk(H, W, C, F, rng=None):
+    if rng:
+        noise = rng.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
+    else:
+        noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
+    noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)
+    noise = noise[F: F + H, F: F + W]
+    noise -= np.min(noise)
+    noise /= np.max(noise)
+    if C == 1:
+        noise = noise[:, :, None]
+    return noise
+
+
+def nms(x, t, s):
+    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
+
+    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
+    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
+    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
+    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+
+    y = np.zeros_like(x)
+
+    for f in [f1, f2, f3, f4]:
+        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+
+    z = np.zeros_like(y, dtype=np.uint8)
+    z[y > t] = 255
+    return z
+
+def min_max_norm(x):
+    x -= np.min(x)
+    x /= np.maximum(np.max(x), 1e-5)
+    return x
+
+
+def safe_step(x, step=2):
+    y = x.astype(np.float32) * float(step + 1)
+    y = y.astype(np.int32).astype(np.float32) / float(step)
+    return y
+
+
+def img2mask(img, H, W, low=10, high=90):
+    assert img.ndim == 3 or img.ndim == 2
+    assert img.dtype == np.uint8
+
+    if img.ndim == 3:
+        y = img[:, :, random.randrange(0, img.shape[2])]
+    else:
+        y = img
+
+    y = cv2.resize(y, (W, H), interpolation=cv2.INTER_CUBIC)
+
+    if random.uniform(0, 1) < 0.5:
+        y = 255 - y
+
+    return y < np.percentile(y, random.randrange(low, high))
+
+def safer_memory(x):
+    # Fix many MAC/AMD problems
+    return np.ascontiguousarray(x.copy()).copy()
+
+UPSCALE_METHODS = ["INTER_NEAREST", "INTER_LINEAR", "INTER_AREA", "INTER_CUBIC", "INTER_LANCZOS4"]
+def get_upscale_method(method_str):
+    assert method_str in UPSCALE_METHODS, f"Method {method_str} not found in {UPSCALE_METHODS}"
+    return getattr(cv2, method_str)
+
+def pad64(x):
+    return int(np.ceil(float(x) / 64.0) * 64 - x)
+
+#https://github.com/Mikubill/sd-webui-controlnet/blob/main/scripts/processor.py#L17
+#Added upscale_method, mode params
+def resize_image_with_pad(input_image, resolution, upscale_method = "", skip_hwc3=False, mode='edge'):
+    if skip_hwc3:
+        img = input_image
+    else:
+        img = HWC3(input_image)
+    H_raw, W_raw, _ = img.shape
+    if resolution == 0:
+        return img, lambda x: x
+    k = float(resolution) / float(min(H_raw, W_raw))
+    H_target = int(np.round(float(H_raw) * k))
+    W_target = int(np.round(float(W_raw) * k))
+    img = cv2.resize(img, (W_target, H_target), interpolation=get_upscale_method(upscale_method) if k > 1 else cv2.INTER_AREA)
+    H_pad, W_pad = pad64(H_target), pad64(W_target)
+    img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode=mode)
+
+    def remove_pad(x):
+        return safer_memory(x[:H_target, :W_target, ...])
+
+    return safer_memory(img_padded), remove_pad
+    
+def common_input_validate(input_image, output_type, **kwargs):
+    if "img" in kwargs:
+            warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)
+            input_image = kwargs.pop("img")
+    
+    if "return_pil" in kwargs:
+            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+            output_type = "pil" if kwargs["return_pil"] else "np"
+    
+    if type(output_type) is bool:
+        warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+        if output_type:
+            output_type = "pil"
+
+    if input_image is None:
+        raise ValueError("input_image must be defined.")
+
+    if not isinstance(input_image, np.ndarray):
+        input_image = np.array(input_image, dtype=np.uint8)
+        output_type = output_type or "pil"
+    else:
+        output_type = output_type or "np"
+    
+    return (input_image, output_type)
+
+def torch_gc():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+
+def ade_palette():
+    """ADE20K palette that maps each class to RGB values."""
+    return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
+            [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
+            [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
+            [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
+            [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
+            [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
+            [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
+            [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
+            [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
+            [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
+            [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
+            [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
+            [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
+            [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
+            [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
+            [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
+            [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
+            [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
+            [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
+            [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
+            [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
+            [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
+            [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
+            [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
+            [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
+            [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
+            [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
+            [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
+            [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
+            [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
+            [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
+            [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
+            [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
+            [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
+            [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
+            [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
+            [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
+            [102, 255, 0], [92, 0, 255]]
+
+#https://stackoverflow.com/a/44873382
+#Assume that the minimum version of Python ppl use is 3.9
+def sha256sum(file_path):
+    import hashlib
+    h  = hashlib.sha256()
+    b  = bytearray(128*1024)
+    mv = memoryview(b)
+    with open(file_path, 'rb', buffering=0) as f:
+        while n := f.readinto(mv):
+            h.update(mv[:n])
+    return h.hexdigest()
+
+def check_hash_from_torch_hub(file_path, filename):
+    basename, _ = filename.split('.')
+    _, ref_hash = basename.split('-')
+    curr_hash = sha256sum(file_path)
+    return curr_hash[:len(ref_hash)] == ref_hash
+
+def custom_torch_download(filename, ckpts_dir=annotator_ckpts_path):
+    """Download PyTorch models using PyTorch 2.7's built-in download mechanism."""
+    model_url = "https://download.pytorch.org/models/" + filename
+    
+    # Use PyTorch's built-in model downloading with custom cache directory
+    local_dir = os.path.join(ckpts_dir, "torch")
+    if not os.path.exists(local_dir):
+        os.makedirs(local_dir, exist_ok=True)
+    
+    model_path = os.path.join(local_dir, filename)
+    
+    if not os.path.exists(model_path):
+        print(f"Downloading {filename} from pytorch.org...")
+        try:
+            # Use PyTorch 2.7's load_url which handles caching, progress, and hash checking
+            state_dict = load_url(model_url, model_dir=local_dir, file_name=filename, progress=True, check_hash=True)
+            # The file is already saved by load_url, we just need the path
+        except Exception as e:
+            warnings.warn(f"Download failed with error: {e}")
+            raise
+    
+    print(f"model_path is {model_path}")
+    return model_path
+
+def custom_hf_download(pretrained_model_or_path, filename, cache_dir=temp_dir, ckpts_dir=annotator_ckpts_path, subfolder='', use_symlinks=USE_SYMLINKS, repo_type="model"):
+
+    local_dir = os.path.join(ckpts_dir, pretrained_model_or_path)
+    model_path = Path(local_dir).joinpath(*subfolder.split('/'), filename).__str__()
+
+    if len(str(model_path)) >= 255:
+        warnings.warn(f"Path {model_path} is too long, \n please change annotator_ckpts_path in config.yaml")
+
+    if not os.path.exists(model_path):
+        print(f"Failed to find {model_path}.\n Downloading from huggingface.co")
+        print(f"cacher folder is {cache_dir}, you can change it by custom_tmp_path in config.yaml")
+        if use_symlinks:
+            cache_dir_d = constants.HF_HUB_CACHE    # use huggingface newer env variables `HF_HUB_CACHE`
+            if cache_dir_d is None:
+                import platform
+                if platform.system() == "Windows":
+                    cache_dir_d = Path(os.getenv("USERPROFILE")).joinpath(".cache", "huggingface", "hub").__str__()
+                else:
+                    cache_dir_d = os.path.join(os.getenv("HOME"), ".cache", "huggingface", "hub")
+            try:
+                # test_link
+                Path(cache_dir_d).mkdir(parents=True, exist_ok=True)
+                Path(ckpts_dir).mkdir(parents=True, exist_ok=True)
+                (Path(cache_dir_d) / f"linktest_{filename}.txt").touch()
+                # symlink instead of link avoid `invalid cross-device link` error.
+                os.symlink(os.path.join(cache_dir_d, f"linktest_{filename}.txt"), os.path.join(ckpts_dir, f"linktest_{filename}.txt"))
+                print("Using symlinks to download models. \n",\
+                      "Make sure you have enough space on your cache folder. \n",\
+                      "And do not purge the cache folder after downloading.\n",\
+                      "Otherwise, you will have to re-download the models every time you run the script.\n",\
+                      "You can use USE_SYMLINKS: False in config.yaml to avoid this behavior.")
+            except:
+                print("Maybe not able to create symlink. Disable using symlinks.")
+                use_symlinks = False
+                cache_dir_d = Path(cache_dir).joinpath("ckpts", pretrained_model_or_path).__str__()
+            finally:    # always remove test link files
+                with suppress(FileNotFoundError):
+                    os.remove(os.path.join(ckpts_dir, f"linktest_{filename}.txt"))
+                    os.remove(os.path.join(cache_dir_d, f"linktest_{filename}.txt"))
+        else:
+            cache_dir_d = os.path.join(cache_dir, "ckpts", pretrained_model_or_path)
+
+        model_path = hf_hub_download(repo_id=pretrained_model_or_path,
+            cache_dir=cache_dir_d,
+            local_dir=local_dir,
+            subfolder=subfolder,
+            filename=filename,
+            local_dir_use_symlinks=use_symlinks,
+            resume_download=True,
+            etag_timeout=100,
+            repo_type=repo_type
+        )
+        if not use_symlinks:
+            try:
+                import shutil
+                shutil.rmtree(os.path.join(cache_dir, "ckpts"))
+            except Exception as e :
+                print(e)
+
+    print(f"model_path is {model_path}")
+
+    return model_path

types.py

@@ -0,0 +1,30 @@
+from typing import NamedTuple, List, Optional
+
+class Keypoint(NamedTuple):
+    x: float
+    y: float
+    score: float = 1.0
+    id: int = -1
+
+
+class BodyResult(NamedTuple):
+    # Note: Using `Optional` instead of `|` operator as the ladder is a Python
+    # 3.10 feature.
+    # Annotator code should be Python 3.8 Compatible, as controlnet repo uses
+    # Python 3.8 environment.
+    # https://github.com/lllyasviel/ControlNet/blob/d3284fcd0972c510635a4f5abe2eeb71dc0de524/environment.yaml#L6
+    keypoints: List[Optional[Keypoint]]
+    total_score: float = 0.0
+    total_parts: int = 0
+
+
+HandResult = List[Keypoint]
+FaceResult = List[Keypoint]
+AnimalPoseResult = List[Keypoint]
+
+
+class PoseResult(NamedTuple):
+    body: BodyResult
+    left_hand: Optional[HandResult]
+    right_hand: Optional[HandResult]
+    face: Optional[FaceResult]

wholebody.py

@@ -0,0 +1,181 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import cv2
+import numpy as np
+
+from .dw_onnx.cv_ox_det import inference_detector as inference_onnx_yolox
+from .dw_onnx.cv_ox_yolo_nas import inference_detector as inference_onnx_yolo_nas
+from .dw_onnx.cv_ox_pose import inference_pose as inference_onnx_pose
+
+from .dw_torchscript.jit_det import inference_detector as inference_jit_yolox
+from .dw_torchscript.jit_pose import inference_pose as inference_jit_pose
+
+from typing import List, Optional
+from .types import PoseResult, BodyResult, Keypoint
+from timeit import default_timer
+import os
+from custom_controlnet_aux.dwpose.util import guess_onnx_input_shape_dtype, get_model_type, get_ort_providers, is_model_torchscript
+import torch
+
+class Wholebody:
+    def __init__(self, det_model_path: Optional[str] = None, pose_model_path: Optional[str] = None, torchscript_device="cuda"):
+        self.det_filename = det_model_path and os.path.basename(det_model_path)
+        self.pose_filename = pose_model_path and os.path.basename(pose_model_path)
+        self.det, self.pose = None, None
+        # return type: None ort cv2 torchscript
+        self.det_model_type = get_model_type("DWPose",self.det_filename)
+        self.pose_model_type = get_model_type("DWPose",self.pose_filename)
+        # Always loads to CPU to avoid building OpenCV.
+        cv2_device = 'cpu'
+        cv2_backend = cv2.dnn.DNN_BACKEND_OPENCV if cv2_device == 'cpu' else cv2.dnn.DNN_BACKEND_CUDA
+        # You need to manually build OpenCV through cmake to work with your GPU.
+        cv2_providers = cv2.dnn.DNN_TARGET_CPU if cv2_device == 'cpu' else cv2.dnn.DNN_TARGET_CUDA
+        ort_providers = get_ort_providers()
+
+        if self.det_model_type is None:
+            pass
+        elif self.det_model_type == "ort":
+            try:
+                import onnxruntime as ort
+                self.det = ort.InferenceSession(det_model_path, providers=ort_providers)
+            except:
+                print(f"Failed to load onnxruntime with {self.det.get_providers()}.\nPlease change EP_list in the config.yaml and restart ComfyUI")
+                self.det = ort.InferenceSession(det_model_path, providers=["CPUExecutionProvider"])
+        elif self.det_model_type == "cv2":
+            try:
+                self.det = cv2.dnn.readNetFromONNX(det_model_path)
+                self.det.setPreferableBackend(cv2_backend)
+                self.det.setPreferableTarget(cv2_providers)
+            except:
+                print("TopK operators may not work on your OpenCV, try use onnxruntime with CPUExecutionProvider")
+                try:
+                    import onnxruntime as ort
+                    self.det = ort.InferenceSession(det_model_path, providers=["CPUExecutionProvider"])
+                except:
+                    print(f"Failed to load {det_model_path}, you can use other models instead")
+        else:
+            self.det = torch.jit.load(det_model_path)
+            self.det.to(torchscript_device)
+
+        if self.pose_model_type is None:
+            pass
+        elif self.pose_model_type == "ort":
+            try:
+                import onnxruntime as ort
+                self.pose = ort.InferenceSession(pose_model_path, providers=ort_providers)
+            except:
+                print(f"Failed to load onnxruntime with {self.pose.get_providers()}.\nPlease change EP_list in the config.yaml and restart ComfyUI")
+                self.pose = ort.InferenceSession(pose_model_path, providers=["CPUExecutionProvider"])
+        elif self.pose_model_type == "cv2":
+            self.pose = cv2.dnn.readNetFromONNX(pose_model_path)
+            self.pose.setPreferableBackend(cv2_backend)
+            self.pose.setPreferableTarget(cv2_providers)
+        else:
+            self.pose = torch.jit.load(pose_model_path)
+            self.pose.to(torchscript_device)
+        
+        if self.pose_filename is not None:
+            self.pose_input_size, _ = guess_onnx_input_shape_dtype(self.pose_filename)
+
+    def __call__(self, oriImg) -> Optional[np.ndarray]:
+        #Sacrifice accurate time measurement for compatibility 
+        
+        det_result = None
+        
+        if self.det is None:
+            print("DWPose: No detector specified, using full image for pose estimation.") # pragma: no cover
+            det_result = []
+        else:
+            det_start = default_timer()
+            if is_model_torchscript(self.det):
+                det_result = inference_jit_yolox(self.det, oriImg, detect_classes=[0])
+            else:
+                if "yolox" in self.det_filename:
+                    det_result = inference_onnx_yolox(self.det, oriImg, detect_classes=[0], dtype=np.float32)
+                else:
+                    #FP16 and INT8 YOLO NAS accept uint8 input
+                    det_result = inference_onnx_yolo_nas(self.det, oriImg, detect_classes=[0], dtype=np.uint8)
+            print(f"DWPose: Bbox {((default_timer() - det_start) * 1000):.2f}ms")
+            if (det_result is None) or (det_result.shape[0] == 0):
+                return None
+
+        pose_start = default_timer()
+        if is_model_torchscript(self.pose):
+            keypoints, scores = inference_jit_pose(self.pose, det_result, oriImg, self.pose_input_size)
+        else:
+            _, pose_onnx_dtype = guess_onnx_input_shape_dtype(self.pose_filename)
+            keypoints, scores = inference_onnx_pose(self.pose, det_result, oriImg, self.pose_input_size, dtype=pose_onnx_dtype)
+            
+        num_subjects_log = 'full image'
+        if hasattr(det_result, 'shape') and det_result.shape[0] > 0:
+            num_subjects_log = f"{det_result.shape[0]} people"
+        elif isinstance(det_result, list) and len(det_result) > 0 and isinstance(det_result[0], (list, np.ndarray)):
+             num_subjects_log = f"{len(det_result)} people"
+        
+        print(f"DWPose: Pose {((default_timer() - pose_start) * 1000):.2f}ms on {num_subjects_log}\n")
+
+        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
+
+        return keypoints_info
+
+    @staticmethod
+    def format_result(keypoints_info: Optional[np.ndarray]) -> List[PoseResult]:
+        def format_keypoint_part(
+            part: np.ndarray,
+        ) -> Optional[List[Optional[Keypoint]]]:
+            keypoints = [
+                Keypoint(x, y, score, i) if score >= 0.3 else None
+                for i, (x, y, score) in enumerate(part)
+            ]
+            return (
+                None if all(keypoint is None for keypoint in keypoints) else keypoints
+            )
+
+        def total_score(keypoints: Optional[List[Optional[Keypoint]]]) -> float:
+            return (
+                sum(keypoint.score for keypoint in keypoints if keypoint is not None)
+                if keypoints is not None
+                else 0.0
+            )
+
+        pose_results = []
+        if keypoints_info is None:
+            return pose_results
+
+        for instance in keypoints_info:
+            body_keypoints = format_keypoint_part(instance[:18]) or ([None] * 18)
+            left_hand = format_keypoint_part(instance[92:113])
+            right_hand = format_keypoint_part(instance[113:134])
+            face = format_keypoint_part(instance[24:92])
+
+            # Openpose face consists of 70 points in total, while DWPose only
+            # provides 68 points. Padding the last 2 points.
+            if face is not None:
+                # left eye
+                face.append(body_keypoints[14])
+                # right eye
+                face.append(body_keypoints[15])
+
+            body = BodyResult(
+                body_keypoints, total_score(body_keypoints), len(body_keypoints)
+            )
+            pose_results.append(PoseResult(body, left_hand, right_hand, face))
+
+        return pose_results

文件.py

@@ -0,0 +1,13 @@
+import os
+
+def print_tree(startpath, indent=''):
+    for entry in os.scandir(startpath):
+        if entry.is_dir():
+            print(f"{indent}{entry.name}/")
+            print_tree(entry.path, indent + '    ')
+        else:
+            print(f"{indent}{entry.name}")
+
+if __name__ == "__main__":
+    base = r"F:\AI\comfyui\ComfyUI_windows_portable\ComfyUI\custom_nodes\comfyui_poseremapping"  # 改成你的路径
+    print_tree(base)

缩放前.json

@@ -0,0 +1,69 @@
+[
+  {
+    "canvas_width": 512,
+    "canvas_height": 768,
+    "people": [
+      {
+        "pose_keypoints_2d": [
+          235.92669018109643,
+          134.29254957040143,
+          1.0,
+          241.18313191334406,
+          195.7929178376993,
+          1.0,
+          181.25969616572064,
+          197.3698503573736,
+          1.0,
+          148.66975742578506,
+          287.78064815203345,
+          1.0,
+          125.54141380389535,
+          366.6272741357485,
+          1.0,
+          301.10656766096747,
+          194.215985318025,
+          1.0,
+          333.696506400903,
+          287.78064815203345,
+          1.0,
+          352.6196966369946,
+          368.7298508286476,
+          1.0,
+          201.23417474826175,
+          389.7556177576382,
+          1.0,
+          187.5674262444178,
+          522.2179494102795,
+          1.0,
+          174.95196608702338,
+          647.3212626377742,
+          1.0,
+          282.18337742487586,
+          388.70432941118884,
+          1.0,
+          296.9014142751694,
+          524.3205261031787,
+          1.0,
+          307.4142977396648,
+          648.3725509842236,
+          1.0,
+          222.25994167725253,
+          121.67708941300702,
+          1.0,
+          247.49086199204135,
+          120.62580106655764,
+          1.0,
+          204.38803978761035,
+          122.72837775945663,
+          1.0,
+          266.41405222813296,
+          119.57451272010803,
+          1.0
+        ],
+        "face_keypoints_2d": [],
+        "hand_left_keypoints_2d": [],
+        "hand_right_keypoints_2d": []
+      }
+    ]
+  }
+]

缩放后.json

@@ -0,0 +1,69 @@
+[
+  {
+    "canvas_width": 512,
+    "canvas_height": 768,
+    "people": [
+      {
+        "pose_keypoints_2d": [
+          240.60175927025875,
+          327.8728573232966,
+          1,
+          245.85820100250635,
+          389.3732255905945,
+          1,
+          157.3458791993362,
+          417.4138496894658,
+          1,
+          124.75594045940045,
+          410.55882169240374,
+          1,
+          84.09079772882873,
+          393.7341436186873,
+          1,
+          331.9620236494394,
+          420.52726298082996,
+          1,
+          374.83711438553235,
+          408.4176515342724,
+          1,
+          407.78551188911126,
+          405.20150301397564,
+          1,
+          189.86329858678482,
+          561.2221837712721,
+          1,
+          189.47190290660046,
+          576.2334005651164,
+          1,
+          189.82157337510864,
+          593.9566681131229,
+          1,
+          291.80488802304865,
+          561.9205542616965,
+          1,
+          294.2773659302132,
+          575.0606323725194,
+          1,
+          297.792787141492,
+          592.3834682042616,
+          1,
+          226.93501076641485,
+          315.2573971659022,
+          1,
+          252.16593108120367,
+          314.2061088194528,
+          1,
+          209.06310887677267,
+          316.3086855123518,
+          1,
+          271.0891213172953,
+          313.1548204730032,
+          1
+        ],
+        "face_keypoints_2d": [],
+        "hand_left_keypoints_2d": [],
+        "hand_right_keypoints_2d": []
+      }
+    ]
+  }
+]