File size: 3,203 Bytes
6ae64bb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """AnimaPoseControl ComfyUI node: detect -> edit (JS) -> styled skeleton -> control IMAGE."""
import numpy as np
try: # keep the helper importable without the package context (tests)
from .pose_render import render, STYLES
from .keypoints import from_json, to_json
from .detect import detect
except ImportError: # flat import when running tests inside the package dir
from pose_render import render, STYLES
from keypoints import from_json, to_json
from detect import detect
def compose(image, pose_json, style, hands, face, feet, redetect, resolution, estimator=None):
"""Pure branch logic. Returns (skeleton uint8 [H,W,3] RGB, pose_json_out str)."""
if redetect and image is not None:
kp, sc = detect(image, int(resolution), estimator=estimator)
elif pose_json and pose_json.strip():
kp, sc, _ = from_json(pose_json)
elif image is not None:
kp, sc = detect(image, int(resolution), estimator=estimator) # no pose yet -> detect
else:
raise ValueError("AnimaPoseControl: no pose to render — connect an image, or Queue once with redetect on first")
img = render(kp, sc, style, int(resolution), hands=hands, face=face, feet=feet)
return img, to_json(kp, sc, int(resolution))
class AnimaPoseControl:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"style": (list(STYLES),),
"hands": ("BOOLEAN", {"default": True}),
"face": ("BOOLEAN", {"default": True}),
"feet": ("BOOLEAN", {"default": True}),
"redetect": ("BOOLEAN", {"default": True}),
"resolution": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 64}),
"pose_json": ("STRING", {"multiline": True, "default": ""}),
},
"optional": {"image": ("IMAGE",)},
}
RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
RETURN_NAMES = ("skeleton", "keypoints")
FUNCTION = "run"
OUTPUT_NODE = True
CATEGORY = "Anima/pose"
def run(self, style, hands, face, feet, redetect, resolution, pose_json, image=None):
import torch
img_in = None
if image is not None:
img_in = (image[0].cpu().numpy() * 255).astype(np.uint8) # ComfyUI IMAGE -> HWC uint8
skel, pj = compose(img_in, pose_json, style, hands, face, feet, redetect, resolution)
t = torch.from_numpy(skel.astype(np.float32) / 255.0)[None, ...]
import json
kpts = json.loads(pj)
# Show the styled render inline via ComfyUI's standard image preview.
return {"ui": {"images": _save_preview(skel), "pose_json": [pj]}, "result": (t, kpts)}
def _save_preview(skel_uint8):
"""Write the styled skeleton to ComfyUI's temp dir; return a /view descriptor for the JS to load."""
import folder_paths, os, uuid
from PIL import Image
d = folder_paths.get_temp_directory(); os.makedirs(d, exist_ok=True)
fn = "anima_pose_%s.png" % uuid.uuid4().hex[:8]
Image.fromarray(skel_uint8).save(os.path.join(d, fn))
return [{"filename": fn, "subfolder": "", "type": "temp"}]
|