xiaden
7d82aa22e6
Sub-packages: - edge/ — Canny, Lineart, LineartAnime (3 nodes) - restore/ — SCUNet, NAFNet, Adversarial (3 nodes) - compare/ — CCIP, LPIPS (2 nodes) - validate/ — Classify (11 ops), Check (5 ops), Metric (2 ops) - tagging/ — WD14, DeepDanbooru, DeepGelbooru, MLDanbooru, Camie, PixAI (6 nodes) Root nodes: - ImgUtilsOCR — PaddleOCR with heat/box threshold - ImgUtilsDetect — 10 detection types + confidence knob - ImgUtilsPose — DWPose + auto_detect toggle - ImgUtilsSegment — isnetis + scale knob - ImgUtilsTags — tag string processing - ImgUtilsCDC — upscaling + tile_size - ImgUtilsCensor — NSFW censoring with method/confidence/body-part toggles - ImgUtilsAlign — resize to max dimension - ImgUtilsSqueeze — transparency-based auto-crop Every node exposes its native function parameters as knobs. Removed mega-nodes: transform (11 ops), compare (3 ops), validate (18 ops).
107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
"""
|
|
Node 5: ImgUtilsPose — "Detect pose"
|
|
|
|
Detects human keypoints in anime images using DWPose and visualizes them.
|
|
Outputs the visualization image and keypoints as JSON.
|
|
|
|
Visualization presets (mapped to op18_visualize flags):
|
|
body_only -> body=True, hands=False, feet=False, face=False
|
|
body_with_face -> body=True, hands=False, feet=False, face=True
|
|
body_with_hands -> body=True, hands=True, feet=False, face=False
|
|
full -> body=True, hands=True, feet=True, face=True
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from comfy_api.latest import io, ui
|
|
|
|
from .utils import comfy_to_pil, pil_to_comfy, keypoints_to_json
|
|
|
|
|
|
class ImgUtilsPose(io.ComfyNode):
|
|
"""Detect human pose keypoints and visualize them."""
|
|
|
|
VISUALIZATION_PRESETS = [
|
|
"body_only",
|
|
"body_with_face",
|
|
"body_with_hands",
|
|
"full",
|
|
]
|
|
|
|
# Mapping from preset name to op18_visualize boolean flags
|
|
PRESET_FLAGS = {
|
|
"body_only": {"draw_body": True, "draw_hands": False, "draw_feet": False, "draw_face": False},
|
|
"body_with_face": {"draw_body": True, "draw_hands": False, "draw_feet": False, "draw_face": True},
|
|
"body_with_hands": {"draw_body": True, "draw_hands": True, "draw_feet": False, "draw_face": False},
|
|
"full": {"draw_body": True, "draw_hands": True, "draw_feet": True, "draw_face": True},
|
|
}
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="ImgUtilsPose",
|
|
display_name="Imgutils Pose",
|
|
category="imgutils/pose",
|
|
description=(
|
|
"Detect OpenPose 18-point human keypoints in anime images "
|
|
"using DWPose. Outputs a visualization image and keypoints JSON."
|
|
),
|
|
search_aliases=[
|
|
"pose", "keypoints", "dwpose", "skeleton", "openpose",
|
|
"body", "face", "hands", "feet",
|
|
],
|
|
inputs=[
|
|
io.Image.Input("image", tooltip="Input image for pose detection"),
|
|
io.Combo.Input(
|
|
"visualization",
|
|
options=cls.VISUALIZATION_PRESETS,
|
|
default="full",
|
|
tooltip="Visualization detail preset: body only, body+face, body+hands, or full (all keypoints).",
|
|
),
|
|
io.Boolean.Input(
|
|
"auto_detect", default=True,
|
|
tooltip="Auto-detect person bounding boxes. Disable to use full image.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output(display_name="visualization"),
|
|
io.String.Output(display_name="keypoints_json"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, image, visualization, auto_detect=True) -> io.NodeOutput:
|
|
"""
|
|
Detect pose keypoints and render visualization.
|
|
|
|
Args:
|
|
image: ComfyUI IMAGE tensor (B,H,W,C), float32, [0,1]
|
|
visualization: preset name for op18_visualize flags
|
|
|
|
Returns:
|
|
NodeOutput with (visualization_IMAGE, keypoints_json_STRING)
|
|
"""
|
|
from imgutils.pose import dwpose_estimate, op18_visualize
|
|
|
|
pil_image = comfy_to_pil(image.numpy() if hasattr(image, "numpy") else image)
|
|
|
|
# Estimate keypoints
|
|
keypoints_list = dwpose_estimate(pil_image, auto_detect=bool(auto_detect))
|
|
|
|
# Build JSON from keypoints
|
|
json_str = keypoints_to_json(keypoints_list)
|
|
|
|
# Visualize with selected preset
|
|
flags = cls.PRESET_FLAGS.get(visualization, cls.PRESET_FLAGS["full"])
|
|
vis_image = op18_visualize(
|
|
pil_image,
|
|
keypoints_list,
|
|
draw_body=flags["draw_body"],
|
|
draw_hands=flags["draw_hands"],
|
|
draw_feet=flags["draw_feet"],
|
|
draw_face=flags["draw_face"],
|
|
)
|
|
|
|
vis_tensor = pil_to_comfy(vis_image)
|
|
return io.NodeOutput(vis_tensor, json_str, ui=ui.PreviewImage(vis_tensor, cls=cls))
|