Merge pull request #6 from andypotato/feat/min-confidence

Feature: Add minimum confidence for keypoints

.github/workflows/publish.yml

@@ -8,15 +8,19 @@ on:
     paths:
       - "pyproject.toml"
 
+permissions:
+  issues: write
+
 jobs:
   publish-node:
     name: Publish Custom Node to registry
     runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'hughescr' }}
     steps:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
+        uses: Comfy-Org/publish-node-action@v1
         with:
           ## Add your own personal access token to your Github Repository secrets and reference it here.
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

openpose_keypoint_extractor.py

@@ -1,4 +1,3 @@
-import json
 from nodes import MAX_RESOLUTION
 
 class OpenPoseKeyPointExtractor:
@@ -10,6 +9,7 @@ class OpenPoseKeyPointExtractor:
                 "image_width": ("INT", { "min": 0, "max": MAX_RESOLUTION }),
                 "image_height": ("INT", { "min": 0, "max": MAX_RESOLUTION }),
                 "points_list": ("STRING", {"multiline": True, "default": ""}),
+                "min_confidence": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
             },
             "optional": {
                 "person_number": ("INT", { "default": 0 }),
@@ -25,17 +25,25 @@ class OpenPoseKeyPointExtractor:
         idx_x = item*3
         idx_y = idx_x + 1
         idx_conf = idx_y + 1
+        # prevent out of bounds errors
+        if idx_conf >= len(list):
+            return (0.0, 0.0, 0.0)
         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):
+    def box_keypoints(self, pose_keypoint, image_width, image_height, points_list,
+                      person_number=0, min_confidence=0.5):
         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)
+
+            # only include points that meet the minimum confidence
+            if z >= min_confidence:
                 if x < min_x:
                     min_x = x
                 if y < min_y:
@@ -44,6 +52,11 @@ class OpenPoseKeyPointExtractor:
                     max_x = x
                 if y > max_y:
                     max_y = y
+
+        # in case no valid keypoints were found return the full image
+        if min_x == MAX_RESOLUTION:
+            return (0, 0, image_width, image_height)
+
         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 = {

pyproject.toml

@@ -2,7 +2,7 @@
 name = "comfyui-openpose-keypoint-extractor"
 description = "This is a single node which can take the POSE_KEYPOINT output from the OpenPose extractor node, parse it, and return x,y,width,height bounding boxes around any elements of the OpenPose skeleton"
 version = "1.0.0"
-license = "LICENSE"
+license = { file = "LICENSE" }
 
 [project.urls]
 Repository = "https://github.com/hughescr/ComfyUI-OpenPose-Keypoint-Extractor"