From 6cc0d36dbee7009a72ffaee02d527cfb491a49cd Mon Sep 17 00:00:00 2001 From: grmchn Date: Fri, 1 May 2026 18:39:34 +0900 Subject: [PATCH] Add mascot auto-detect nodes --- README.md | 9 + README_ja.md | 9 + nodes.py | 29 + proportion_changer/__init__.py | 16 + proportion_changer/mascot_nodes.py | 1004 ++++++++++++++++++++++++++++ requirements.txt | 6 +- 6 files changed, 1071 insertions(+), 2 deletions(-) create mode 100644 proportion_changer/mascot_nodes.py diff --git a/README.md b/README.md index 558c056..6e34e4b 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,11 @@ Additionally, the Openpose Editor node from [toyxyz/ComfyUI-ultimate-openpose-ed - **pose_keypoint resize**: Resizes `POSE_KEYPOINT` to a target size (pads then scales to avoid stretching when aspect differs) - **pose_keypoint input**: Converts JSON text to KeyPoints - **pose_keypoint preview**: Converts KeyPoints to JSON +- **(Down)Load Mascot DWPose Model**: Downloads and loads a mascot DWPose ONNX model from HuggingFace +- **(Down)Load Mascot BBox Model**: Downloads and loads a mascot bbox ONNX model from HuggingFace +- **Mascot DWPose Detector**: Detects mascot body pose as `POSE_KEYPOINT` +- **Mascot BBox Detector**: Detects mascot part bounding boxes as `BOUNDING_BOX` +- **pose_keypoint to dw_poses**: Converts 25-point `POSE_KEYPOINT` to SCAIL-Pose `DWPOSES` ## Installation ### Install via ComfyUI Manager @@ -64,6 +69,10 @@ Please refer to example_workflows. 2. **Incorrect body proportions after transformation with reference image**: The `pose_keypoint` and `reference_pose_keypoint` aspect ratios (canvas width/height) may not match. **ProportionChanger Reference** has `auto_resize_reference` (default ON) to automatically align the reference to the pose canvas size. If needed, use **pose_keypoint resize** to explicitly align `width`/`height`. Fine-tune individual body parts using the "ProportionChanger Params" node. 3. **Nothing displays after transformation with reference image**: The reference image's pose estimation by DWPose has failed. Use OpenposeEditor or similar tools to input parameter values manually. +### Mascot OpenPose Models + +Mascot OpenPose nodes download model artifacts from `grmchn/mascot-openpose-detect-test` on HuggingFace into `ComfyUI/models/mascot_body_detect`. The current test model package is CC-BY-NC 4.0 because it includes the ViTPose-L variant derived from a non-commercial pretrained backbone. The node code remains part of this repository, but use of the downloaded test models is subject to that model license. + ## Attribution and Credits ### Special Thanks - **[kijai/ComfyUI-WanVideoWrapper](https://github.com/kijai/ComfyUI-WanVideoWrapper)** diff --git a/README_ja.md b/README_ja.md index dcd139a..dd45737 100644 --- a/README_ja.md +++ b/README_ja.md @@ -20,6 +20,11 @@ - **pose_keypoint resize**: `POSE_KEYPOINT`を指定サイズにリサイズ（アスペクト比が違う場合はパディング→スケールで歪みを防ぐ） - **pose_keypoint input**: JSONテキストからKeyPointに変換する - **pose_keypoint preview**: KeyPointからJSONに変換する +- **(Down)Load Mascot DWPose Model**: HuggingFaceからマスコット用DWPose ONNXモデルをダウンロードして読み込む +- **(Down)Load Mascot BBox Model**: HuggingFaceからマスコット用bbox ONNXモデルをダウンロードして読み込む +- **Mascot DWPose Detector**: マスコットの姿勢を`POSE_KEYPOINT`として検出する +- **Mascot BBox Detector**: マスコットのパーツbboxを`BOUNDING_BOX`として検出する +- **pose_keypoint to dw_poses**: 25点`POSE_KEYPOINT`をSCAIL-Pose用`DWPOSES`に変換する ## インストール ### ComfyUI Managerでインストール @@ -65,6 +70,10 @@ example_workflowsを参照してください。 2. **リファレンス画像ありでプロポーション変換した後の体型がおかしい**: `pose_keypoint`と`reference_pose_keypoint`の縦横比（canvas幅/高さ）が一致していない可能性があります。**ProportionChanger Reference**の`auto_resize_reference`（デフォルトON）でリファレンス側を自動で揃えます。必要に応じて**pose_keypoint resize**ノードで明示的に`width`/`height`を揃えてください。体の各パーツの微調整は「ProportionChanger Params」ノードで行ってください。 3. **リファレンス画像ありでプロポーション変換した後に何も表示されない**:リファレンス画像のDWPoseによる姿勢推定が失敗しています。OpenposeEditorなどを利用してパラメータの数値を入力してください。 +### Mascot OpenPoseモデル + +Mascot OpenPoseノードは、HuggingFaceの`grmchn/mascot-openpose-detect-test`から`ComfyUI/models/mascot_body_detect`へモデルをダウンロードします。現在のtestモデルパッケージは、非商用のpretrained backboneに由来するViTPose-L variantを含むため、CC-BY-NC 4.0です。ノードコードはこのリポジトリの一部ですが、ダウンロードされたtestモデルの利用はモデル側のライセンスに従います。 + ## 帰属とクレジット ### Special Thanks - **[kijai/ComfyUI-WanVideoWrapper](https://github.com/kijai/ComfyUI-WanVideoWrapper)** diff --git a/nodes.py b/nodes.py index 92dede9..cb6455b 100644 --- a/nodes.py +++ b/nodes.py @@ -5,6 +5,16 @@ Backward compatibility maintained for existing workflows from .utils import log import json +import os + +try: + import folder_paths + folder_paths.add_model_folder_path( + "mascot_body_detect", + os.path.join(folder_paths.models_dir, "mascot_body_detect"), + ) +except ImportError: + folder_paths = None # Import all nodes from our organized structure from .proportion_changer import ( @@ -15,6 +25,11 @@ from .proportion_changer import ( ProportionChangerInterpolator, PoseDataToPoseKeypoint, PoseKeypointResize, + DownloadAndLoadMascotDWPoseModel, + DownloadAndLoadMascotBBoxModel, + MascotDWPoseDetector, + MascotBBoxDetector, + ConvertToSCAILPose, # TODO: Re-enable after fixing Kalman filter issues # ProportionChangerKeypointDenoiser, # ProportionChangerKeypointDenoiserAdvanced @@ -208,6 +223,13 @@ NODE_CLASS_MAPPINGS = { "PoseKeypointResize": PoseKeypointResize, "PoseJSONToPoseKeypoint": PoseJSONToPoseKeypoint, "PoseKeypointPreview": PoseKeypointPreview, + + # Mascot OpenPose nodes + "DownloadAndLoadMascotDWPoseModel": DownloadAndLoadMascotDWPoseModel, + "DownloadAndLoadMascotBBoxModel": DownloadAndLoadMascotBBoxModel, + "MascotDWPoseDetector": MascotDWPoseDetector, + "MascotBBoxDetector": MascotBBoxDetector, + "ConvertToSCAILPose": ConvertToSCAILPose, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -227,4 +249,11 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PoseKeypointResize": "pose_keypoint resize", "PoseJSONToPoseKeypoint": "pose_keypoint input", "PoseKeypointPreview": "pose_keypoint preview", + + # Mascot OpenPose nodes + "DownloadAndLoadMascotDWPoseModel": "(Down)Load Mascot DWPose Model", + "DownloadAndLoadMascotBBoxModel": "(Down)Load Mascot BBox Model", + "MascotDWPoseDetector": "Mascot DWPose Detector", + "MascotBBoxDetector": "Mascot BBox Detector", + "ConvertToSCAILPose": "pose_keypoint to dw_poses", } diff --git a/proportion_changer/__init__.py b/proportion_changer/__init__.py index ce09f1f..ead3d51 100644 --- a/proportion_changer/__init__.py +++ b/proportion_changer/__init__.py @@ -34,6 +34,15 @@ from .converter_nodes import ( PoseKeypointResize, ) +# Mascot OpenPose nodes +from .mascot_nodes import ( + DownloadAndLoadMascotDWPoseModel, + DownloadAndLoadMascotBBoxModel, + MascotDWPoseDetector, + MascotBBoxDetector, + ConvertToSCAILPose, +) + # Import keypoint denoiser node classes from .keypoint_denoiser_node import ( ProportionChangerKeypointDenoiser, @@ -59,6 +68,13 @@ __all__ = [ # Converter nodes "PoseDataToPoseKeypoint", "PoseKeypointResize", + + # Mascot OpenPose nodes + "DownloadAndLoadMascotDWPoseModel", + "DownloadAndLoadMascotBBoxModel", + "MascotDWPoseDetector", + "MascotBBoxDetector", + "ConvertToSCAILPose", # KeyPoint Denoiser nodes "ProportionChangerKeypointDenoiser", diff --git a/proportion_changer/mascot_nodes.py b/proportion_changer/mascot_nodes.py new file mode 100644 index 0000000..7929e6e --- /dev/null +++ b/proportion_changer/mascot_nodes.py @@ -0,0 +1,1004 @@ +""" +Mascot OpenPose nodes for ProportionChanger. + +The model artifacts are downloaded from HuggingFace and executed with +ONNXRuntime. This module intentionally has no runtime dependency on the +mascot_body_detect source repository. +""" + +from __future__ import annotations + +import json +import logging +import math +import os +from typing import Any + +import cv2 +import numpy as np + +try: + import folder_paths +except ImportError: # pragma: no cover - ComfyUI provides this at runtime. + folder_paths = None + +try: + from comfy.utils import ProgressBar +except ImportError: # pragma: no cover - standalone import smoke tests. + ProgressBar = None + +try: + from ..utils import log +except ImportError: # pragma: no cover + log = logging.getLogger(__name__) + + +DEFAULT_REPO_ID = "grmchn/mascot-openpose-detect-test" +HF_BASE_URL = f"https://huggingface.co/{DEFAULT_REPO_ID}/resolve/main" +MASCOT_BBOX_MODEL_URLS = [ + f"{HF_BASE_URL}/bbox/model.onnx", +] +MASCOT_DWPOSE_MODEL_URLS = [ + f"{HF_BASE_URL}/keypoint/vitpose_l/model.onnx", + f"{HF_BASE_URL}/keypoint/rtmpose_s/model.onnx", +] +POSE_URL_TO_VARIANT = { + f"{HF_BASE_URL}/keypoint/vitpose_l/model.onnx": "vitpose_l", + f"{HF_BASE_URL}/keypoint/rtmpose_s/model.onnx": "rtmpose_s", +} +BBOX_KEYS = ( + "full", + "head", + "body", + "hand_left", + "hand_right", + "foot_left", + "foot_right", +) + +IMAGENET_MEAN = np.array([123.675, 116.28, 103.53], dtype=np.float32) +IMAGENET_STD = np.array([58.395, 57.12, 57.375], dtype=np.float32) + +COCO17_TO_DWPOSE25 = (0, 15, 14, 17, 16, 5, 2, 6, 3, 7, 4, 8, 11, 9, 12, 10, 13) + +NEUTRAL_HAND_TEMPLATE_RIGHT = np.array( + [ + (0.00, 0.00), + (0.18, 0.12), + (0.32, 0.26), + (0.42, 0.40), + (0.50, 0.52), + (0.20, 0.55), + (0.22, 0.72), + (0.23, 0.83), + (0.24, 0.92), + (0.06, 0.60), + (0.07, 0.78), + (0.07, 0.90), + (0.08, 1.00), + (-0.08, 0.58), + (-0.09, 0.74), + (-0.10, 0.86), + (-0.11, 0.95), + (-0.20, 0.52), + (-0.22, 0.64), + (-0.24, 0.74), + (-0.26, 0.82), + ], + dtype=np.float32, +) +NEUTRAL_HAND_TEMPLATE_LEFT = NEUTRAL_HAND_TEMPLATE_RIGHT * np.array([-1.0, 1.0], dtype=np.float32) + + +def _models_root() -> str: + if folder_paths is None: + return os.path.abspath(os.path.join("models", "mascot_body_detect")) + return os.path.join(folder_paths.models_dir, "mascot_body_detect") + + +def _read_json(path: str) -> Any: + with open(path, "r", encoding="utf-8") as handle: + return json.load(handle) + + +def _download_artifacts(url: str, *, include_bbox: bool, keypoint_variant: str | None = None) -> str: + if include_bbox and url not in MASCOT_BBOX_MODEL_URLS and keypoint_variant is None: + raise ValueError(f"URL {url} is not in the list of allowed Mascot BBox models.") + if keypoint_variant is not None and url not in MASCOT_DWPOSE_MODEL_URLS: + raise ValueError(f"URL {url} is not in the list of allowed Mascot DWPose models.") + + from huggingface_hub import snapshot_download + + base = _models_root() + os.makedirs(base, exist_ok=True) + allow_patterns = [] + required_paths = [] + if include_bbox: + allow_patterns.append("bbox/*") + required_paths.extend( + [ + os.path.join(base, "bbox", "model.onnx"), + os.path.join(base, "bbox", "classes.json"), + os.path.join(base, "bbox", "decode_params.json"), + ] + ) + if keypoint_variant is not None: + allow_patterns.append(f"keypoint/{keypoint_variant}/*") + required_paths.extend( + [ + os.path.join(base, "keypoint", keypoint_variant, "model.onnx"), + os.path.join(base, "keypoint", keypoint_variant, "meta.json"), + ] + ) + + if not all(os.path.exists(path) for path in required_paths): + snapshot_download( + repo_id=DEFAULT_REPO_ID, + repo_type="model", + allow_patterns=allow_patterns, + local_dir=base, + ) + return base + + +def _providers(cuda: bool) -> list[str]: + return ["CUDAExecutionProvider", "CPUExecutionProvider"] if cuda else ["CPUExecutionProvider"] + + +def _progress(total: int): + if ProgressBar is None: + return None + return ProgressBar(total) + + +def sigmoid(x: np.ndarray) -> np.ndarray: + return 1.0 / (1.0 + np.exp(-x)) + + +def letterbox(img_rgb: np.ndarray, size_wh: tuple[int, int]) -> tuple[np.ndarray, float, tuple[int, int]]: + img_h, img_w = img_rgb.shape[:2] + target_w, target_h = size_wh + scale = min(target_w / img_w, target_h / img_h) + resized_w = int(round(img_w * scale)) + resized_h = int(round(img_h * scale)) + resized = cv2.resize(img_rgb, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR) + canvas = np.full((target_h, target_w, 3), 114, dtype=np.uint8) + pad_x = (target_w - resized_w) // 2 + pad_y = (target_h - resized_h) // 2 + canvas[pad_y : pad_y + resized_h, pad_x : pad_x + resized_w] = resized + return canvas, scale, (pad_x, pad_y) + + +def decode_rtmdet_scale(cls_logits: np.ndarray, bbox_dist: np.ndarray, stride: int, score_thr: float) -> np.ndarray: + height, width = cls_logits.shape[-2:] + yy, xx = np.meshgrid(np.arange(height), np.arange(width), indexing="ij") + center_x = (xx + 0.5) * stride + center_y = (yy + 0.5) * stride + scores = sigmoid(cls_logits[0]) + cls_idx = scores.argmax(axis=0) + cls_score = scores.max(axis=0) + keep = cls_score >= score_thr + if not keep.any(): + return np.zeros((0, 6), dtype=np.float32) + + # Exported RTMDet ONNX bbox distances are already in input-image pixels. + # Do not multiply them by stride; doing so expands boxes to the whole image. + left = bbox_dist[0, 0][keep] + top = bbox_dist[0, 1][keep] + right = bbox_dist[0, 2][keep] + bottom = bbox_dist[0, 3][keep] + kept_x = center_x[keep] + kept_y = center_y[keep] + return np.stack( + [ + kept_x - left, + kept_y - top, + kept_x + right, + kept_y + bottom, + cls_score[keep], + cls_idx[keep], + ], + axis=1, + ).astype(np.float32) + + +def iou_xyxy(a: np.ndarray, b: np.ndarray) -> float: + x1 = max(float(a[0]), float(b[0])) + y1 = max(float(a[1]), float(b[1])) + x2 = min(float(a[2]), float(b[2])) + y2 = min(float(a[3]), float(b[3])) + inter = max(0.0, x2 - x1) * max(0.0, y2 - y1) + area_a = max(0.0, float(a[2] - a[0])) * max(0.0, float(a[3] - a[1])) + area_b = max(0.0, float(b[2] - b[0])) * max(0.0, float(b[3] - b[1])) + union = area_a + area_b - inter + return inter / union if union > 0.0 else 0.0 + + +def nms_per_class(boxes: np.ndarray, iou_thr: float) -> np.ndarray: + if len(boxes) == 0: + return boxes + + keep_global: list[int] = [] + class_ids = boxes[:, 5].astype(np.int32) + for class_id in np.unique(class_ids): + indices = np.where(class_ids == class_id)[0] + class_boxes = boxes[indices] + order = np.argsort(-class_boxes[:, 4]) + suppressed = np.zeros(len(order), dtype=bool) + for pos, src_order_idx in enumerate(order): + if suppressed[pos]: + continue + keep_global.append(int(indices[src_order_idx])) + for later_pos in range(pos + 1, len(order)): + if suppressed[later_pos]: + continue + if iou_xyxy(class_boxes[src_order_idx, :4], class_boxes[order[later_pos], :4]) >= iou_thr: + suppressed[later_pos] = True + return boxes[sorted(keep_global)] + + +def _run_bbox(session, img_rgb: np.ndarray, params: dict[str, Any], score_threshold: float) -> np.ndarray: + input_size = tuple(params["input_size"]) + padded, scale, (pad_x, pad_y) = letterbox(img_rgb, input_size) + inp = (padded.astype(np.float32) - np.asarray(params["mean"], dtype=np.float32)) / np.asarray( + params["std"], dtype=np.float32 + ) + inp = inp.transpose(2, 0, 1)[None] + + input_name = session.get_inputs()[0].name + outputs = session.run(None, {input_name: inp}) + cls_outputs = outputs[:3] + bbox_outputs = outputs[3:] + + boxes_by_scale = [ + decode_rtmdet_scale(cls_out, bbox_out, int(stride), score_threshold) + for cls_out, bbox_out, stride in zip(cls_outputs, bbox_outputs, params["strides"]) + ] + boxes = np.concatenate(boxes_by_scale, axis=0) if boxes_by_scale else np.zeros((0, 6), dtype=np.float32) + boxes = nms_per_class(boxes, float(params["nms_iou_threshold"])) + if len(boxes) == 0: + return boxes + + img_h, img_w = img_rgb.shape[:2] + boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_x) / scale + boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_y) / scale + boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, img_w) + boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, img_h) + return boxes + + +def infer_bbox_raw( + session, + img_rgb: np.ndarray, + score_threshold: float, + params: dict[str, Any], + classes_map: dict[str, str], +) -> list[dict[str, float | str]]: + boxes = _run_bbox(session, img_rgb, params, score_threshold) + out: list[dict[str, float | str]] = [] + for x1, y1, x2, y2, score, class_id_raw in boxes: + class_name = classes_map.get(str(int(class_id_raw))) + if class_name is None: + continue + out.append( + { + "x": float(x1), + "y": float(y1), + "width": float(max(0.0, x2 - x1)), + "height": float(max(0.0, y2 - y1)), + "label": class_name, + "score": float(score), + } + ) + return out + + +def infer_bbox_dict( + session, + img_rgb: np.ndarray, + score_threshold: float, + params: dict[str, Any], + classes_map: dict[str, str], +) -> dict[str, list[float] | None]: + boxes = _run_bbox(session, img_rgb, params, score_threshold) + img_h, img_w = img_rgb.shape[:2] + out: dict[str, list[float] | None] = {name: None for name in BBOX_KEYS} + if len(boxes) == 0: + return out + + labelled_boxes = [] + for x1, y1, x2, y2, score, class_id_raw in boxes: + class_name = classes_map.get(str(int(class_id_raw))) + if class_name in BBOX_KEYS: + labelled_boxes.append( + { + "xyxy": (float(x1), float(y1), float(x2), float(y2)), + "score": float(score), + "label": class_name, + } + ) + if not labelled_boxes: + return out + + roi_candidates = [box for box in labelled_boxes if box["label"] == "full"] + if not roi_candidates: + roi_candidates = [box for box in labelled_boxes if box["label"] == "body"] + if not roi_candidates: + return out + + roi = max(roi_candidates, key=lambda box: box["score"]) + roi_x1, roi_y1, roi_x2, roi_y2 = roi["xyxy"] + roi_w = max(1.0, roi_x2 - roi_x1) + roi_h = max(1.0, roi_y2 - roi_y1) + # Body boxes do not always enclose hands and feet; expand fallback ROIs. + margin_x = roi_w * (0.35 if roi["label"] == "body" else 0.05) + margin_y = roi_h * (0.35 if roi["label"] == "body" else 0.05) + gate = ( + max(0.0, roi_x1 - margin_x), + max(0.0, roi_y1 - margin_y), + min(float(img_w), roi_x2 + margin_x), + min(float(img_h), roi_y2 + margin_y), + ) + + def inside_gate(box: dict[str, Any]) -> bool: + x1, y1, x2, y2 = box["xyxy"] + cx = (x1 + x2) / 2.0 + cy = (y1 + y2) / 2.0 + gx1, gy1, gx2, gy2 = gate + return gx1 <= cx <= gx2 and gy1 <= cy <= gy2 + + def to_yxyx_1000(box: dict[str, Any]) -> list[float]: + x1, y1, x2, y2 = box["xyxy"] + return [ + float(y1) / img_h * 1000.0, + float(x1) / img_w * 1000.0, + float(y2) / img_h * 1000.0, + float(x2) / img_w * 1000.0, + ] + + out[roi["label"]] = to_yxyx_1000(roi) + for class_name in BBOX_KEYS: + if class_name == roi["label"]: + continue + candidates = [box for box in labelled_boxes if box["label"] == class_name and inside_gate(box)] + if not candidates: + continue + out[class_name] = to_yxyx_1000(max(candidates, key=lambda box: box["score"])) + return out + + +def bbox_yxyx_1000_to_xyxy_pixel( + bbox: list[float] | None, + img_w: int, + img_h: int, +) -> tuple[float, float, float, float] | None: + if bbox is None: + return None + ymin, xmin, ymax, xmax = bbox + return ( + xmin / 1000.0 * img_w, + ymin / 1000.0 * img_h, + xmax / 1000.0 * img_w, + ymax / 1000.0 * img_h, + ) + + +def topdown_affine( + bbox_xyxy: tuple[float, float, float, float], + input_size: tuple[int, int], + pad_ratio: float, +) -> np.ndarray: + x1, y1, x2, y2 = bbox_xyxy + target_w, target_h = input_size + center_x = (x1 + x2) / 2.0 + center_y = (y1 + y2) / 2.0 + bbox_w = max(1.0, (x2 - x1) * pad_ratio) + bbox_h = max(1.0, (y2 - y1) * pad_ratio) + aspect = target_w / target_h + if bbox_w / bbox_h > aspect: + bbox_h = bbox_w / aspect + else: + bbox_w = bbox_h * aspect + crop_x1 = center_x - bbox_w / 2.0 + crop_y1 = center_y - bbox_h / 2.0 + return np.array( + [ + [target_w / bbox_w, 0.0, -crop_x1 * target_w / bbox_w], + [0.0, target_h / bbox_h, -crop_y1 * target_h / bbox_h], + ], + dtype=np.float32, + ) + + +def run_topdown_crop( + session, + img_rgb: np.ndarray, + bbox_yxyx_0_1000: list[float] | None, + input_size: tuple[int, int], + pad_ratio: float, +) -> tuple[list[np.ndarray] | None, np.ndarray | None]: + img_h, img_w = img_rgb.shape[:2] + bbox_xyxy = bbox_yxyx_1000_to_xyxy_pixel(bbox_yxyx_0_1000, img_w, img_h) + if bbox_xyxy is None: + return None, None + matrix = topdown_affine(bbox_xyxy, input_size, pad_ratio) + crop = cv2.warpAffine( + img_rgb, + matrix, + input_size, + flags=cv2.INTER_LINEAR, + borderValue=(114, 114, 114), + ) + inp = (crop.astype(np.float32) - IMAGENET_MEAN) / IMAGENET_STD + inp = inp.transpose(2, 0, 1)[None] + input_name = session.get_inputs()[0].name + return session.run(None, {input_name: inp}), matrix + + +def decode_heatmap(heatmap: np.ndarray, heatmap_to_input_scale: float) -> np.ndarray: + keypoints, heat_h, heat_w = heatmap.shape + out = np.zeros((keypoints, 3), dtype=np.float32) + for idx in range(keypoints): + flat_idx = int(heatmap[idx].argmax()) + y, x = divmod(flat_idx, heat_w) + score = float(heatmap[idx, y, x]) + dx = 0.0 + dy = 0.0 + if 1 <= x < heat_w - 1: + dx = 0.25 * float(np.sign(heatmap[idx, y, x + 1] - heatmap[idx, y, x - 1])) + if 1 <= y < heat_h - 1: + dy = 0.25 * float(np.sign(heatmap[idx, y + 1, x] - heatmap[idx, y - 1, x])) + out[idx] = [ + (x + dx) * heatmap_to_input_scale, + (y + dy) * heatmap_to_input_scale, + score, + ] + return out + + +def decode_simcc(simcc_x: np.ndarray, simcc_y: np.ndarray, split_ratio: float) -> np.ndarray: + x = simcc_x[0] + y = simcc_y[0] + keypoints = x.shape[0] + out = np.zeros((keypoints, 3), dtype=np.float32) + for idx in range(keypoints): + x_idx = int(x[idx].argmax()) + y_idx = int(y[idx].argmax()) + score = min(float(x[idx, x_idx]), float(y[idx, y_idx])) + out[idx] = [x_idx / split_ratio, y_idx / split_ratio, score] + return out + + +def crop_to_image_points(kp_crop_xy_score: np.ndarray, affine_matrix: np.ndarray) -> np.ndarray: + inv_matrix = cv2.invertAffineTransform(affine_matrix) + out = np.zeros_like(kp_crop_xy_score, dtype=np.float32) + for idx, (x, y, score) in enumerate(kp_crop_xy_score): + mapped = inv_matrix @ np.array([float(x), float(y), 1.0], dtype=np.float32) + out[idx] = [float(mapped[0]), float(mapped[1]), float(score)] + return out + + +def coco17_to_dwpose25(kp_17: list[list[float]], derive_neck: bool = True) -> list[list[float]]: + arr = np.asarray(kp_17, dtype=np.float32) + out = np.zeros((25, 3), dtype=np.float32) + for coco_idx, dw_idx in enumerate(COCO17_TO_DWPOSE25): + out[dw_idx] = arr[coco_idx] + + if derive_neck: + nose = arr[0] + left_shoulder = arr[5] + right_shoulder = arr[6] + if left_shoulder[2] > 0 and right_shoulder[2] > 0: + midpoint = (left_shoulder[:2] + right_shoulder[:2]) / 2.0 + if nose[2] > 0: + centroid = (nose[:2] + left_shoulder[:2] + right_shoulder[:2]) / 3.0 + out[1, :2] = (midpoint + centroid) / 2.0 + out[1, 2] = float(min(nose[2], left_shoulder[2], right_shoulder[2])) + else: + out[1, :2] = midpoint + out[1, 2] = float(min(left_shoulder[2], right_shoulder[2])) + return out.tolist() + + +def derive_toes_from_bbox( + kp_25_pixel: list[list[float]], + foot_left: list[float] | None, + foot_right: list[float] | None, + img_w: int, + img_h: int, +) -> list[list[float]]: + arr = np.asarray(kp_25_pixel, dtype=np.float32) + for toe_idx, ankle_idx, bbox in ( + (18, 13, foot_right), + (19, 10, foot_left), + ): + arr[toe_idx] = [0.0, 0.0, 0.0] + if bbox is None: + continue + ymin, xmin, ymax, xmax = bbox + arr[toe_idx, 0] = (xmin + xmax) / 2.0 / 1000.0 * img_w + arr[toe_idx, 1] = (ymin + ymax) / 2.0 / 1000.0 * img_h + ankle_score = float(arr[ankle_idx, 2]) + arr[toe_idx, 2] = ankle_score if ankle_score > 0.0 else 0.5 + return arr.tolist() + + +def infer_keypoints_pixel( + session, + img_rgb: np.ndarray, + bbox_dict: dict[str, list[float] | None], + kp_meta: dict[str, Any], + variant: str, +) -> list[list[float]]: + img_h, img_w = img_rgb.shape[:2] + roi = bbox_dict.get("full") or bbox_dict.get("body") + input_size = tuple(kp_meta["input_size"]) + pad_ratio = float(kp_meta.get("topdown_bbox_pad_ratio", 1.25)) + outputs, matrix = run_topdown_crop(session, img_rgb, roi, input_size, pad_ratio) + if outputs is None or matrix is None: + return [[0.0, 0.0, 0.0] for _ in range(25)] + + if variant == "vitpose_l": + scale = float(kp_meta.get("heatmap_to_input_scale", 4.0)) + kp_crop = decode_heatmap(outputs[0][0], scale) + kp_image_17 = crop_to_image_points(kp_crop, matrix) + kp_25 = coco17_to_dwpose25(kp_image_17.tolist(), derive_neck=True) + elif variant == "rtmpose_s": + kp_crop = decode_simcc(outputs[0], outputs[1], float(kp_meta["split_ratio"])) + kp_image_25 = crop_to_image_points(kp_crop, matrix) + kp_25 = kp_image_25[:25].tolist() + if len(kp_25) < 25: + kp_25.extend([[0.0, 0.0, 0.0] for _ in range(25 - len(kp_25))]) + else: + raise ValueError(f"unknown keypoint variant: {variant}") + + kp_25 = derive_toes_from_bbox( + kp_25, + bbox_dict.get("foot_left"), + bbox_dict.get("foot_right"), + img_w=img_w, + img_h=img_h, + ) + for idx in range(20, 25): + kp_25[idx] = [0.0, 0.0, 0.0] + return kp_25 + + +def _bbox_short_side_pixel(bbox: list[float], img_w: int, img_h: int) -> float: + ymin, xmin, ymax, xmax = bbox + width = max(1.0, (xmax - xmin) / 1000.0 * img_w) + height = max(1.0, (ymax - ymin) / 1000.0 * img_h) + return min(width, height) + + +def _build_hand_2d( + hand_bbox: list[float] | None, + kp_25_pixel: list[list[float]], + *, + elbow_idx: int, + wrist_idx: int, + side: str, + img_w: int, + img_h: int, + canvas_w: int, + canvas_h: int, +) -> list[float]: + if hand_bbox is None: + return [0.0] * (21 * 3) + + arr = np.asarray(kp_25_pixel, dtype=np.float32) + wrist = arr[wrist_idx] + elbow = arr[elbow_idx] + if wrist[2] <= 0: + ymin, xmin, ymax, xmax = hand_bbox + wrist_xy = np.array( + [ + (xmin + xmax) / 2.0 / 1000.0 * img_w, + (ymin + ymax) / 2.0 / 1000.0 * img_h, + ], + dtype=np.float32, + ) + else: + wrist_xy = wrist[:2].astype(np.float32) + + forearm_len = 0.0 + if elbow[2] > 0 and wrist[2] > 0: + direction = wrist[:2] - elbow[:2] + forearm_len = float(np.linalg.norm(direction)) + else: + direction = np.array([0.0, 1.0], dtype=np.float32) + norm = float(np.linalg.norm(direction)) + if norm < 1e-6: + direction = np.array([0.0, 1.0], dtype=np.float32) + else: + direction = direction / norm + + dx, dy = float(direction[0]), float(direction[1]) + theta = math.atan2(-dx, dy) + cos_t = math.cos(theta) + sin_t = math.sin(theta) + rot = np.array([[cos_t, -sin_t], [sin_t, cos_t]], dtype=np.float32) + template = NEUTRAL_HAND_TEMPLATE_LEFT if side == "left" else NEUTRAL_HAND_TEMPLATE_RIGHT + scale = _bbox_short_side_pixel(hand_bbox, img_w, img_h) + caps = [max(8.0, min(float(img_w), float(img_h)) * 0.18)] + if forearm_len > 1e-6: + caps.append(max(8.0, forearm_len * 1.25)) + scale = min(scale, *caps) + points = template @ rot.T * scale + wrist_xy + + out: list[float] = [] + for x, y in points: + out.extend([float(x) / canvas_w, float(y) / canvas_h, 0.5]) + return out + + +def build_person_dict( + stage1_bbox: dict[str, list[float] | None], + kp_25_pixel: list[list[float]], + canvas_width: int, + canvas_height: int, + img_width: int, + img_height: int, +) -> dict[str, list[float]]: + pose_kp: list[float] = [] + for idx in range(25): + if 20 <= idx < 25: + pose_kp.extend([0.0, 0.0, 0.0]) + continue + x, y, score = kp_25_pixel[idx] + pose_kp.extend([float(x) / img_width, float(y) / img_height, float(score)]) + + right_hand = _build_hand_2d( + stage1_bbox.get("hand_right"), + kp_25_pixel, + elbow_idx=3, + wrist_idx=4, + side="right", + img_w=img_width, + img_h=img_height, + canvas_w=img_width, + canvas_h=img_height, + ) + left_hand = _build_hand_2d( + stage1_bbox.get("hand_left"), + kp_25_pixel, + elbow_idx=6, + wrist_idx=7, + side="left", + img_w=img_width, + img_h=img_height, + canvas_w=img_width, + canvas_h=img_height, + ) + + return { + "pose_keypoints_2d": pose_kp, + "face_keypoints_2d": [0.0] * (70 * 3), + "hand_left_keypoints_2d": left_hand, + "hand_right_keypoints_2d": right_hand, + } + + +class DownloadAndLoadMascotDWPoseModel: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ( + MASCOT_DWPOSE_MODEL_URLS, + {"default": f"{HF_BASE_URL}/keypoint/vitpose_l/model.onnx"}, + ), + "cuda": ( + "BOOLEAN", + {"default": True, "tooltip": "Use CUDAExecutionProvider when available."}, + ), + }, + "optional": { + "warmup": ( + "BOOLEAN", + {"default": True, "tooltip": "Run a small dummy inference after loading."}, + ), + }, + } + + RETURN_TYPES = ("MASCOT_POSE_MODEL",) + RETURN_NAMES = ("mascot_pose_model",) + FUNCTION = "loadmodel" + CATEGORY = "ProportionChanger" + DESCRIPTION = "Download and load a mascot DWPose ONNX model from HuggingFace." + + def loadmodel(self, url: str, cuda: bool, warmup: bool = True): + if folder_paths is None: + raise RuntimeError("folder_paths is required inside ComfyUI to locate the model directory") + + import onnxruntime as ort + + keypoint_variant = POSE_URL_TO_VARIANT.get(url) + if keypoint_variant is None: + raise ValueError(f"URL {url} is not in the list of allowed Mascot DWPose models.") + + # The current top-down DWPose ONNX still needs a body ROI. Keep the bbox + # sidecar inside this handle so the DWPose detector remains a single-input + # node while exposing a separate bbox loader for bbox-only workflows. + base = _download_artifacts(url, include_bbox=True, keypoint_variant=keypoint_variant) + bbox_onnx = os.path.join(base, "bbox", "model.onnx") + kp_onnx = os.path.join(base, "keypoint", keypoint_variant, "model.onnx") + providers = _providers(bool(cuda)) + bbox_session = ort.InferenceSession(bbox_onnx, providers=providers) + kp_session = ort.InferenceSession(kp_onnx, providers=providers) + + if warmup: + log.info("Warming up Mascot DWPose model...") + bbox_input = bbox_session.get_inputs()[0] + kp_input = kp_session.get_inputs()[0] + bbox_shape = [1 if not isinstance(dim, int) else dim for dim in bbox_input.shape] + kp_shape = [1 if not isinstance(dim, int) else dim for dim in kp_input.shape] + bbox_session.run(None, {bbox_input.name: np.zeros(bbox_shape, dtype=np.float32)}) + kp_session.run(None, {kp_input.name: np.zeros(kp_shape, dtype=np.float32)}) + log.info("Mascot DWPose model warmed up") + + return ( + { + "variant": keypoint_variant, + "url": url, + "cuda": bool(cuda), + "bbox_session": bbox_session, + "bbox_decode": _read_json(os.path.join(base, "bbox", "decode_params.json")), + "bbox_classes": _read_json(os.path.join(base, "bbox", "classes.json")), + "kp_session": kp_session, + "kp_meta": _read_json(os.path.join(base, "keypoint", keypoint_variant, "meta.json")), + }, + ) + + +class DownloadAndLoadMascotBBoxModel: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": (MASCOT_BBOX_MODEL_URLS, {"default": f"{HF_BASE_URL}/bbox/model.onnx"}), + "cuda": ( + "BOOLEAN", + {"default": True, "tooltip": "Use CUDAExecutionProvider when available."}, + ), + }, + "optional": { + "warmup": ( + "BOOLEAN", + {"default": True, "tooltip": "Run a small dummy inference after loading."}, + ), + }, + } + + RETURN_TYPES = ("MASCOT_BBOX_MODEL",) + RETURN_NAMES = ("mascot_bbox_model",) + FUNCTION = "loadmodel" + CATEGORY = "ProportionChanger" + DESCRIPTION = "Download and load a mascot bbox ONNX model from HuggingFace." + + def loadmodel(self, url: str, cuda: bool, warmup: bool = True): + if folder_paths is None: + raise RuntimeError("folder_paths is required inside ComfyUI to locate the model directory") + + import onnxruntime as ort + + if url not in MASCOT_BBOX_MODEL_URLS: + raise ValueError(f"URL {url} is not in the list of allowed Mascot BBox models.") + + base = _download_artifacts(url, include_bbox=True) + bbox_onnx = os.path.join(base, "bbox", "model.onnx") + session = ort.InferenceSession(bbox_onnx, providers=_providers(bool(cuda))) + + if warmup: + log.info("Warming up Mascot BBox model...") + input_info = session.get_inputs()[0] + input_shape = [1 if not isinstance(dim, int) else dim for dim in input_info.shape] + session.run(None, {input_info.name: np.zeros(input_shape, dtype=np.float32)}) + log.info("Mascot BBox model warmed up") + + return ( + { + "url": url, + "cuda": bool(cuda), + "bbox_session": session, + "bbox_decode": _read_json(os.path.join(base, "bbox", "decode_params.json")), + "bbox_classes": _read_json(os.path.join(base, "bbox", "classes.json")), + }, + ) + + +class MascotDWPoseDetector: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mascot_pose_model": ("MASCOT_POSE_MODEL", {"tooltip": "Output of (Down)Load Mascot DWPose Model"}), + "image": ("IMAGE", {"tooltip": "Input image, RGB float [0,1]."}), + "width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8}), + "height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8}), + "threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ("POSE_KEYPOINT",) + RETURN_NAMES = ("pose_keypoint",) + FUNCTION = "detect" + CATEGORY = "ProportionChanger" + DESCRIPTION = "Detect mascot body pose as ProportionChanger POSE_KEYPOINT." + + def detect(self, mascot_pose_model, image, width: int, height: int, threshold: float): + out_frames = [] + pbar = _progress(len(image)) + + for img_t in image: + img_np = np.clip(img_t.cpu().numpy() * 255.0, 0, 255).astype(np.uint8) + img_h, img_w = img_np.shape[:2] + bbox_dict = infer_bbox_dict( + mascot_pose_model["bbox_session"], + img_np, + threshold, + mascot_pose_model["bbox_decode"], + mascot_pose_model["bbox_classes"], + ) + person = None + if bbox_dict.get("full") is not None or bbox_dict.get("body") is not None: + kp_25_pixel = infer_keypoints_pixel( + mascot_pose_model["kp_session"], + img_np, + bbox_dict, + mascot_pose_model["kp_meta"], + mascot_pose_model["variant"], + ) + person = build_person_dict( + stage1_bbox=bbox_dict, + kp_25_pixel=kp_25_pixel, + canvas_width=int(width), + canvas_height=int(height), + img_width=img_w, + img_height=img_h, + ) + + out_frames.append( + { + "version": "1.0", + "people": [person] if person else [], + "canvas_width": int(width), + "canvas_height": int(height), + } + ) + if pbar is not None: + pbar.update(1) + + return (out_frames,) + + +class MascotBBoxDetector: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mascot_bbox_model": ("MASCOT_BBOX_MODEL", {"tooltip": "Output of (Down)Load Mascot BBox Model"}), + "image": ("IMAGE",), + "threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), + "class_filter": ( + ["all", "full", "head", "body", "hand_left", "hand_right", "foot_left", "foot_right"], + {"default": "all"}, + ), + "max_detections": ("INT", {"default": 7, "min": 1, "max": 100, "step": 1}), + }, + } + + RETURN_TYPES = ("BOUNDING_BOX",) + RETURN_NAMES = ("bboxes",) + FUNCTION = "detect" + CATEGORY = "ProportionChanger" + DESCRIPTION = "Detect mascot part bounding boxes as comfy-core BOUNDING_BOX." + + def detect(self, mascot_bbox_model, image, threshold: float, class_filter: str, max_detections: int): + out_frames = [] + pbar = _progress(len(image)) + for img_t in image: + img_np = np.clip(img_t.cpu().numpy() * 255.0, 0, 255).astype(np.uint8) + boxes = infer_bbox_raw( + mascot_bbox_model["bbox_session"], + img_np, + threshold, + mascot_bbox_model["bbox_decode"], + mascot_bbox_model["bbox_classes"], + ) + if class_filter != "all": + boxes = [box for box in boxes if box["label"] == class_filter] + boxes = sorted(boxes, key=lambda box: -float(box["score"]))[: int(max_detections)] + out_frames.append(boxes) + if pbar is not None: + pbar.update(1) + return (out_frames,) + + +class ConvertToSCAILPose: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "keypoints": ("POSE_KEYPOINT",), + }, + } + + RETURN_TYPES = ("DWPOSES",) + RETURN_NAMES = ("dw_poses",) + FUNCTION = "convert" + CATEGORY = "ProportionChanger" + DESCRIPTION = "Convert 25-point POSE_KEYPOINT to SCAIL-Pose DWPOSES." + + def convert(self, keypoints): + num_body = 18 + num_face = 68 + num_hand = 21 + out_frames = [] + + for frame in keypoints or []: + people = frame.get("people", []) if isinstance(frame, dict) else [] + people = people[:1] if isinstance(people, list) else [] + + bodies = [] + body_scores = [] + hands = [] + hand_scores = [] + faces = [] + face_scores = [] + + for person in people: + if not isinstance(person, dict): + continue + pose_raw = person.get("pose_keypoints_2d") or [] + if len(pose_raw) < num_body * 3: + continue + body_arr = np.asarray(pose_raw[: num_body * 3], dtype=np.float32).reshape(num_body, 3) + bodies.append(body_arr[:, :2]) + body_scores.append(body_arr[:, 2]) + + face_raw = person.get("face_keypoints_2d") or [] + if len(face_raw) >= num_face * 3: + face_arr = np.asarray(face_raw[: num_face * 3], dtype=np.float32).reshape(num_face, 3) + faces.append(face_arr[:, :2]) + face_scores.append(face_arr[:, 2]) + else: + faces.append(np.zeros((num_face, 2), dtype=np.float32)) + face_scores.append(np.zeros(num_face, dtype=np.float32)) + + for hand_key in ("hand_left_keypoints_2d", "hand_right_keypoints_2d"): + hand_raw = person.get(hand_key) or [] + if len(hand_raw) == num_hand * 3: + hand_arr = np.asarray(hand_raw, dtype=np.float32).reshape(num_hand, 3) + hands.append(hand_arr[:, :2]) + hand_scores.append(hand_arr[:, 2]) + + out_frames.append( + { + "bodies": { + "candidate": np.asarray(bodies, dtype=np.float32) + if bodies + else np.zeros((0, num_body, 2), dtype=np.float32), + "subset": np.asarray([np.arange(num_body) for _ in bodies], dtype=np.float32) + if bodies + else np.zeros((0, num_body), dtype=np.float32), + }, + "hands": np.asarray(hands, dtype=np.float32) + if hands + else np.zeros((0, num_hand, 2), dtype=np.float32), + "faces": np.asarray(faces, dtype=np.float32) + if faces + else np.zeros((0, num_face, 2), dtype=np.float32), + "body_score": np.asarray(body_scores, dtype=np.float32) + if body_scores + else np.zeros((0, num_body), dtype=np.float32), + "hand_score": np.asarray(hand_scores, dtype=np.float32) + if hand_scores + else np.zeros((0, num_hand), dtype=np.float32), + "face_score": np.asarray(face_scores, dtype=np.float32) + if face_scores + else np.zeros((0, num_face), dtype=np.float32), + } + ) + + return ({"poses": out_frames, "swap_hands": False},) diff --git a/requirements.txt b/requirements.txt index 1b5e544..383d82d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,6 +7,8 @@ opencv-python # Progress bars and utilities tqdm +huggingface_hub -# Optional: ONNX Runtime (for ONNX-based pose detection - currently not used in active nodes) -# onnxruntime \ No newline at end of file +# ONNX Runtime for Mascot OpenPose nodes. Install onnxruntime-gpu instead if +# your ComfyUI environment manages CUDA providers separately. +onnxruntime