diff --git a/README.md b/README.md index c78787f..f6e5829 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,13 @@ # ComfyUI-ultimate-openpose-editor -ultimate openpose editor with render + +This is an improved version of [ComfyUI-openpose-editor](https://github.com/huchenlei/ComfyUI-openpose-editor) +in ComfyUI, enable input and output with flexible choices. Much more convenient and easier to use. It integrates the render function which you also can intall it separately from my [ultimate-openpose-render](https://github.com/westNeighbor/ComfyUI-ultimate-openpose-render) repo or search in the Custom Nodes Manager + +## Install & Features & Usage + +- Just install this repo through the Manager or the Git, or you can just simply copy to your *ComfyUI/custom\_nodes/* directory +- Insert node by `Right Click -> ultimate-openpose -> Openpose Edito Node` +- The node enable multiple functions: + - It can accept POSE\_KEYPOINTS or poses in json format as input. + - To edit poses, right click the node, and open the editor through *Open in Openpose Editor* menu. The send back poses will be shown in the pose json input area. Be ware that the edit priority is POSE\_KEYPOINT > POSE\_JSON + - It can output pose images, or POSE\_KEYPOINTS or the json poses. Be ware that the output priority is POSE\_JSON > POSE\_KEYPOINT diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..bda963a --- /dev/null +++ b/__init__.py @@ -0,0 +1,14 @@ +from .openpose_editor_nodes import OpenposeEditorNode + + +WEB_DIRECTORY = "js" + +NODE_CLASS_MAPPINGS = { + "OpenposeEditorNode": OpenposeEditorNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "OpenposeEditorNode": "Openpose Editor Node", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/js/editor_textDisplay.js b/js/editor_textDisplay.js new file mode 100644 index 0000000..7f31596 --- /dev/null +++ b/js/editor_textDisplay.js @@ -0,0 +1,69 @@ +import { app } from "../../../scripts/app.js"; +import { ComfyWidgets } from "../../../scripts/widgets.js"; + +// Displays input text on a node +app.registerExtension({ + name: "EditOpenposeText", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "OpenposeEditorNode") { + function populate(text) { + let removed = 0; + if (this.widgets) { + for (let i = 8; i < this.widgets.length; i++) { + if (this.widgets[i].name === "display_only") { + this.widgets[i].onRemove?.(); + removed++; + } + } + this.widgets.length = this.widgets.length - removed; + } + + if (this.inputs?.[0].link || this.inputs?.[this.inputs.length-1].widget) { + const v = [...text]; + if (!v[0]) { + v.shift(); + } + for (const list of v) { + // const w = this.widgets[0]; + // w.value = list; + // if (this.inputs?.[this.inputs.length-1].widget) { + + const w = ComfyWidgets["STRING"](this, "display_only", ["STRING", { multiline: true }], app).widget; + w.inputEl.readOnly = true; + w.inputEl.style.opacity = 0.6; + w.value = list; + + //} + } + } + + requestAnimationFrame(() => { + const sz = this.computeSize(); + if (sz[0] < this.size[0]) { + sz[0] = this.size[0]; + } + if (sz[1] < this.size[1]) { + sz[1] = this.size[1]; + } + this.onResize?.(sz); + app.graph.setDirtyCanvas(true, false); + }); + } + + // When the node is executed we will be sent the input text, display this in the widget + const onExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function (message) { + onExecuted?.apply(this, arguments); + populate.call(this, message.POSE_JSON); + }; + + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function () { + onConfigure?.apply(this, arguments); + if (this.widgets_values?.length) { + populate.call(this, this.widgets_values.slice(+this.widgets_values.length > 1)); + } + }; + } + }, +}); diff --git a/js/openpose_editor.js b/js/openpose_editor.js new file mode 100644 index 0000000..eefcfa6 --- /dev/null +++ b/js/openpose_editor.js @@ -0,0 +1,161 @@ +import { app } from "../../scripts/app.js"; +import { ComfyDialog, $el } from "../../scripts/ui.js"; +import { ComfyApp } from "../../scripts/app.js"; + + +function addMenuHandler(nodeType, cb) { + const getOpts = nodeType.prototype.getExtraMenuOptions; + nodeType.prototype.getExtraMenuOptions = function () { + const r = getOpts.apply(this, arguments); + cb.apply(this, arguments); + return r; + }; +} + +class OpenposeEditorDialog extends ComfyDialog { + static timeout = 5000; + static instance = null; + + static getInstance() { + if (!OpenposeEditorDialog.instance) { + OpenposeEditorDialog.instance = new OpenposeEditorDialog(); + } + + return OpenposeEditorDialog.instance; + } + + constructor() { + super(); + this.element = $el("div.comfy-modal", { + parent: document.body, + style: { + width: "80vw", + height: "80vh", + }, + }, [ + $el("div.comfy-modal-content", { + style: { + width: "100%", + height: "100%", + }, + }, this.createButtons()), + ]); + this.is_layout_created = false; + + window.addEventListener("message", (event) => { + if (event.source !== this.iframeElement.contentWindow) { + return; + } + + const message = event.data; + if (message.modalId === 0) { + const targetNode = ComfyApp.clipspace_return_node; + const textAreaElement = targetNode.widgets[7].element; + textAreaElement.value = JSON.stringify(event.data.poses); + this.close(); + } + }); + } + + createButtons() { + const closeBtn = $el("button", { + type: "button", + textContent: "Close", + onclick: () => this.close(), + }); + return [ + closeBtn, + ]; + } + + close() { + super.close(); + } + + async show() { + if (!this.is_layout_created) { + this.createLayout(); + this.is_layout_created = true; + await this.waitIframeReady(); + } + + const targetNode = ComfyApp.clipspace_return_node; + if (targetNode.inputs?.[0].link || targetNode.inputs?.[targetNode.inputs.length-1].widget){ + const textAreaElement = targetNode.widgets[8].element; + this.element.style.display = "flex"; + this.setCanvasJSONString(textAreaElement.value.replace(/'/g, '"')); + } else { + const textAreaElement = targetNode.widgets[7].element; + this.element.style.display = "flex"; + this.setCanvasJSONString(textAreaElement.value.replace(/'/g, '"')); + } + + } + + createLayout() { + this.iframeElement = $el("iframe", { + // Change to for local dev + // src: "http://localhost:5173", + src: "https://huchenlei.github.io/sd-webui-openpose-editor?theme=dark", + style: { + width: "100%", + height: "100%", + border: "none", + }, + }); + const modalContent = this.element.querySelector(".comfy-modal-content"); + while (modalContent.firstChild) { + modalContent.removeChild(modalContent.firstChild); + } + modalContent.appendChild(this.iframeElement); + } + + waitIframeReady() { + return new Promise((resolve, reject) => { + const receiveMessage = (event) => { + if (event.source !== this.iframeElement.contentWindow) { + return; + } + if (event.data.ready) { + window.removeEventListener("message", receiveMessage); + clearTimeout(timeoutHandle); + resolve(); + } + }; + const timeoutHandle = setTimeout(() => { + reject(new Error("Timeout")); + }, OpenposeEditorDialog.timeout); + + window.addEventListener("message", receiveMessage); + }); + } + + setCanvasJSONString(jsonString) { + this.iframeElement.contentWindow.postMessage({ + modalId: 0, + poses: JSON.parse(jsonString) + }, "*"); + } +} + +app.registerExtension({ + name: "OpenposeEditor", + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name === "OpenposeEditorNode") { + addMenuHandler(nodeType, function (_, options) { + options.unshift({ + content: "Open in Openpose Editor", + callback: () => { + // `this` is the node instance + ComfyApp.copyToClipspace(this); + ComfyApp.clipspace_return_node = this; + + const dlg = OpenposeEditorDialog.getInstance(); + dlg.show(); + }, + }); + }); + } + } +}); diff --git a/openpose_editor_nodes.py b/openpose_editor_nodes.py new file mode 100644 index 0000000..6c1a782 --- /dev/null +++ b/openpose_editor_nodes.py @@ -0,0 +1,70 @@ +import json +import torch +import numpy as np +from .util import draw_pose_json + +OpenposeJSON = dict + +class OpenposeEditorNode: + @classmethod + def INPUT_TYPES(s): + return { + "optional": { + "show_body": ("BOOLEAN", {"default": True}), + "show_face": ("BOOLEAN", {"default": True}), + "show_hands": ("BOOLEAN", {"default": True}), + "resolution_x": ("INT", { + "default": -1, + "min": -1, + "max": 12800 + }), + "pose_marker_size": ("INT", { + "default": 4, + "min": 0, + "max": 100 + }), + "face_marker_size": ("INT", { + "default": 3, + "min": 0, + "max": 100 + }), + "hand_marker_size": ("INT", { + "default": 2, + "min": 0, + "max": 100 + }), + "POSE_JSON": ("STRING", {"multiline": True}), + "POSE_KEYPOINT": ("POSE_KEYPOINT",{"default": None}), + }, + } + + RETURN_NAMES = ("POSE_IMAGE", "POSE_KEYPOINT", "POSE_JSON") + RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT", "STRING") + OUTPUT_NODE = True + FUNCTION = "load_pose" + CATEGORY = "ultimate-openpose" + + def load_pose(self, show_body, show_face, show_hands, resolution_x, pose_marker_size, face_marker_size, hand_marker_size, POSE_JSON: str, POSE_KEYPOINT=None) -> tuple[OpenposeJSON]: + ''' + priority output is: POSE_JSON > POSE_KEYPOINT + priority edit is: POSE_KEYPOINT > POSE_JSON + ''' + if POSE_JSON: + POSE_JSON = POSE_JSON.replace("'",'"').replace('None','[]') + POSE_PASS = POSE_JSON + if POSE_KEYPOINT is not None: + POSE_PASS = json.dumps(POSE_KEYPOINT,indent=4).replace("'",'"').replace('None','[]') + pose_imgs = draw_pose_json(POSE_JSON, resolution_x, show_body, show_face, show_hands, pose_marker_size, face_marker_size, hand_marker_size) + pose_imgs_np = np.array(pose_imgs).astype(np.float32) / 255 + return { + "ui": {"POSE_JSON": [POSE_PASS]}, + "result": (torch.from_numpy(pose_imgs_np), json.loads(POSE_JSON), POSE_JSON) + } + elif POSE_KEYPOINT is not None: + POSE_JSON = json.dumps(POSE_KEYPOINT,indent=4).replace("'",'"').replace('None','[]') + pose_imgs = draw_pose_json(POSE_JSON, resolution_x, show_body, show_face, show_hands, pose_marker_size, face_marker_size, hand_marker_size) + pose_imgs_np = np.array(pose_imgs).astype(np.float32) / 255 + return { + "ui": {"POSE_JSON": [POSE_JSON]}, + "result": (torch.from_numpy(pose_imgs_np), json.loads(POSE_JSON), POSE_JSON) + } diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..ec0bace --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "ComfyUI-ultimate-openpose-editor" +description = "Improved version of [ComfyUI-openpose-editor](https://github.com/huchenlei/ComfyUI-openpose-editor) in ComfyUI, enable input and output with flexible choices. Much more convenient and easier to use." +version = "1.0.0" +license = "LICENSE" + +[project.urls] +Repository = "https://github.com/westNeighbor/ComfyUI-ultimate-openpose-editor" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "westNeighbor" +DisplayName = "ComfyUI-ultimate-openpose-editor" +Icon = "" diff --git a/util.py b/util.py new file mode 100644 index 0000000..87ce06e --- /dev/null +++ b/util.py @@ -0,0 +1,170 @@ +import math +import json +import numpy as np +import matplotlib +import cv2 + +eps = 0.01 + +def draw_pose_json(pose_json, resolution_x, show_body, show_face, show_hands, pose_marker_size, face_marker_size, hand_marker_size): + pose_imgs = [] + if pose_json: + if pose_json.startswith('{'): + pose_json = '[{}]'.format(pose_json) + images = json.loads(pose_json) + for image in images: + figures = image['people'] + H = image['canvas_height'] + W = image['canvas_width'] + bodies = [] + candidate = [] + subset = [[]] + faces = [] + hands = [] + for figure in figures: + body = figure['pose_keypoints_2d'] + face = figure['face_keypoints_2d'] + lhand = figure['hand_left_keypoints_2d'] + rhand = figure['hand_right_keypoints_2d'] + if body: + for i in range(0,len(body),3): + candidate.append(body[i:i+2]) + if not subset[0]: + subset[0].extend([len(subset[0])+(i//3) if body[i+2]>0 else -1 for i in range(0,len(body),3)]) + else: + subset.append([len(subset[0])*len(subset)+(i//3) if body[i+2]>0 else -1 for i in range(0,len(body),3)]) + if face: + faces.append([face[i:i+2] for i in range(0,len(face),3)]) + if lhand: + hands.append([lhand[i:i+2] for i in range(0,len(lhand),3)]) + if rhand: + hands.append([rhand[i:i+2] for i in range(0,len(rhand),3)]) + + candidate = np.array(candidate).astype(float) + subset = np.array(subset) + faces = np.array(faces).astype(float) + hands = np.array(hands).astype(float) + normalized = max(np.max(candidate[...,0]),np.max(candidate[...,1]),np.max(faces[...,0]),np.max(faces[...,1]),np.max(hands[...,0]),np.max(hands[...,1])) + if normalized>2.0: + candidate[...,0] /= float(W) + candidate[...,1] /= float(H) + faces[...,0] /= float(W) + faces[...,1] /= float(H) + hands[...,0] /= float(W) + hands[...,1] /= float(H) + bodies = dict(candidate=candidate, subset=subset) + pose = dict(bodies=bodies, faces=faces, hands=hands) + pose = dict(bodies=bodies if show_body else {'candidate':[], 'subset':[]}, faces=faces if show_face else [], hands=hands if show_hands else []) + W_scaled = resolution_x + if resolution_x < 64: + W_scaled = W + H_scaled = int(H*(W_scaled*1.0/W)) + pose_img = draw_pose(pose, H_scaled, W_scaled, pose_marker_size, face_marker_size, hand_marker_size) + pose_imgs.append(pose_img) + + return pose_imgs + +def draw_pose(pose, H, W, pose_marker_size, face_marker_size, hand_marker_size): + bodies = pose['bodies'] + faces = pose['faces'] + hands = pose['hands'] + candidate = bodies['candidate'] + subset = bodies['subset'] + canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8) + + canvas = draw_bodypose(canvas, candidate, subset, pose_marker_size) + + canvas = draw_handpose(canvas, hands, hand_marker_size) + + canvas = draw_facepose(canvas, faces, face_marker_size) + + return canvas + +def draw_bodypose(canvas, candidate, subset, pose_marker_size): + H, W, C = canvas.shape + candidate = np.array(candidate) + subset = np.array(subset) + + # stickwidth = 4 + + limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \ + [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \ + [1, 16], [16, 18], [3, 17], [6, 18]] + + colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \ + [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \ + [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]] + + for i in range(17): + for n in range(len(subset)): + index = subset[n][np.array(limbSeq[i]) - 1] + if -1 in index: + continue + Y = candidate[index.astype(int), 0] * float(W) + X = candidate[index.astype(int), 1] * float(H) + mX = np.mean(X) + mY = np.mean(Y) + length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 + angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) + polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), pose_marker_size), int(angle), 0, 360, 1) + cv2.fillConvexPoly(canvas, polygon, colors[i]) + + canvas = (canvas * 0.6).astype(np.uint8) + + for i in range(18): + for n in range(len(subset)): + index = int(subset[n][i]) + if index == -1: + continue + x, y = candidate[index][0:2] + x = int(x * W) + y = int(y * H) + cv2.circle(canvas, (int(x), int(y)), pose_marker_size, colors[i], thickness=-1) + + return canvas + + +def draw_handpose(canvas, all_hand_peaks, hand_marker_size): + H, W, C = canvas.shape + + edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \ + [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] + + for peaks in all_hand_peaks: + peaks = np.array(peaks) + + for ie, e in enumerate(edges): + x1, y1 = peaks[e[0]] + x2, y2 = peaks[e[1]] + x1 = int(x1 * W) + y1 = int(y1 * H) + x2 = int(x2 * W) + y2 = int(y2 * H) + 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=1 if hand_marker_size == 0 else hand_marker_size) + + joint_size=0 + if hand_marker_size < 2: + joint_size = hand_marker_size + 1 + else: + joint_size = hand_marker_size + 2 + for i, keyponit in enumerate(peaks): + x, y = keyponit + x = int(x * W) + y = int(y * H) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), joint_size, (0, 0, 255), thickness=-1) + return canvas + + +def draw_facepose(canvas, all_lmks, face_marker_size): + H, W, C = canvas.shape + for lmks in all_lmks: + lmks = np.array(lmks) + for lmk in lmks: + x, y = lmk + x = int(x * W) + y = int(y * H) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), face_marker_size, (255, 255, 255), thickness=-1) + return canvas