commit 70d1dc2591447e2d180639d907a9cb6e0e76819e Author: Craig R. Hughes Date: Sat Feb 24 13:35:21 2024 -0800 Initial commit of working version diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bee8a64 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ diff --git a/README.md b/README.md new file mode 100644 index 0000000..d830351 --- /dev/null +++ b/README.md @@ -0,0 +1,11 @@ +# Openpose Keypoint Extractor +- Take the keypoint output from [OpenPose estimator node](https://github.com/Fannovel16/comfyui_controlnet_aux?tab=readme-ov-file#faces-and-poses-estimators) and calculate bounding boxes around those keypoints. +- You need to give it the width and height of the original image and it will output `(x,y,width,height)` bounding box within that image +- Note that the points on the OpenPose skeleton are *inside* the particular limb (eg center of wrist, middle of shoulder), so you probably will want to apply some padding around the bounding box to get the whole arm, leg or whatever you're looking for. +- Reference for which control point is which - provide comma-separated list in the Node: [OpenPose Docs](https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/02_output.md) + +# Example workflow + +Drag this to your ComfyUI: + +![Workflow image](workflow.png) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..605d787 --- /dev/null +++ b/__init__.py @@ -0,0 +1,2 @@ +from .openpose_keypoint_extractor import NODE_CLASS_MAPPINGS +__all__ = ["NODE_CLASS_MAPPINGS"] diff --git a/openpose_keypoint_extractor.py b/openpose_keypoint_extractor.py new file mode 100644 index 0000000..5b445e5 --- /dev/null +++ b/openpose_keypoint_extractor.py @@ -0,0 +1,51 @@ +import json +from nodes import MAX_RESOLUTION + +class OpenPoseKeyPointExtractor: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "pose_keypoint": ("POSE_KEYPOINT",), + "image_width": ("INT", { "min": 0, "max": MAX_RESOLUTION }), + "image_height": ("INT", { "min": 0, "max": MAX_RESOLUTION }), + "points_list": ("STRING", {"multiline": True, "default": ""}), + }, + "optional": { + "person_number": ("INT", { "default": 0 }), + } + } + + RETURN_TYPES = ("INT", "INT", "INT", "INT") + RETURN_NAMES = ("x", "y", "width", "height") + FUNCTION = "box_keypoints" + CATEGORY = "utils" + + def get_keypoint_from_list(self, list, item): + idx_x = item*3 + idx_y = idx_x + 1 + idx_conf = idx_y + 1 + return (list[idx_x], list[idx_y], list[idx_conf]) + + def box_keypoints(self, pose_keypoint, image_width, image_height, points_list, person_number=0): + points_we_want = [int(element) for element in points_list.split(",")] + + min_x = MAX_RESOLUTION + min_y = MAX_RESOLUTION + max_x = 0 + max_y = 0 + for element in points_we_want: + (x,y,z) = self.get_keypoint_from_list(pose_keypoint[0]["people"][person_number]["pose_keypoints_2d"], element) + if x < min_x: + min_x = x + if y < min_y: + min_y = y + if x > max_x: + max_x = x + if y > max_y: + max_y = y + return (int(min_x*image_width), int(min_y*image_height), int((max_x-min_x)*image_width), int((max_y-min_y)*image_height)) + +NODE_CLASS_MAPPINGS = { + "Openpose Keypoint Extractor": OpenPoseKeyPointExtractor, +} diff --git a/workflow.png b/workflow.png new file mode 100644 index 0000000..a2b5944 Binary files /dev/null and b/workflow.png differ