Refactor sdpose again due to upstream change

comfy_extras/nodes_sdpose.py

@@ -2,11 +2,10 @@ import torch
 import comfy.utils
 import comfy.model_management
 import numpy as np
-import math
-import colorsys
 from tqdm import tqdm
 from typing_extensions import override
 from comfy_api.latest import ComfyExtension, io
+from comfy_extras.pose.keypoint_draw import KeypointDraw
 from comfy_extras.nodes_lotus import LotusConditioning
 
 
@@ -73,299 +72,6 @@ def _to_openpose_frames(all_keypoints, all_scores, height, width):
     return frames
 
 
-class KeypointDraw:
-    """
-    Pose keypoint drawing class that supports both numpy and cv2 backends.
-    """
-    def __init__(self):
-        try:
-            import cv2
-            self.draw = cv2
-        except ImportError:
-            self.draw = self
-
-        # Hand connections (same for both hands)
-        self.hand_edges = [
-            [0, 1], [1, 2], [2, 3], [3, 4],      # thumb
-            [0, 5], [5, 6], [6, 7], [7, 8],      # index
-            [0, 9], [9, 10], [10, 11], [11, 12], # middle
-            [0, 13], [13, 14], [14, 15], [15, 16], # ring
-            [0, 17], [17, 18], [18, 19], [19, 20], # pinky
-        ]
-
-        # Body connections - matching DWPose limbSeq (1-indexed, converted to 0-indexed)
-        self.body_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]
-        ]
-
-        # Head connections (1-indexed, converted to 0-indexed)
-        self.head_edges = [
-            [2, 1], [1, 15], [15, 17], [1, 16], [16, 18]
-        ]
-
-        # Colors matching DWPose
-        self.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]
-        ]
-
-    @staticmethod
-    def circle(canvas_np, center, radius, color, **kwargs):
-        """Draw a filled circle using NumPy vectorized operations."""
-        cx, cy = center
-        h, w = canvas_np.shape[:2]
-
-        radius_int = int(np.ceil(radius))
-
-        y_min, y_max = max(0, cy - radius_int), min(h, cy + radius_int + 1)
-        x_min, x_max = max(0, cx - radius_int), min(w, cx + radius_int + 1)
-
-        if y_max <= y_min or x_max <= x_min:
-            return
-
-        y, x = np.ogrid[y_min:y_max, x_min:x_max]
-        mask = (x - cx)**2 + (y - cy)**2 <= radius**2
-        canvas_np[y_min:y_max, x_min:x_max][mask] = color
-
-    @staticmethod
-    def line(canvas_np, pt1, pt2, color, thickness=1, **kwargs):
-        """Draw line using Bresenham's algorithm with NumPy operations."""
-        x0, y0, x1, y1 = *pt1, *pt2
-        h, w = canvas_np.shape[:2]
-        dx, dy = abs(x1 - x0), abs(y1 - y0)
-        sx, sy = (1 if x0 < x1 else -1), (1 if y0 < y1 else -1)
-        err, x, y, line_points = dx - dy, x0, y0, []
-
-        while True:
-            line_points.append((x, y))
-            if x == x1 and y == y1:
-                break
-            e2 = 2 * err
-            if e2 > -dy:
-                err, x = err - dy, x + sx
-            if e2 < dx:
-                err, y = err + dx, y + sy
-
-        if thickness > 1:
-            radius, radius_int = (thickness / 2.0) + 0.5, int(np.ceil((thickness / 2.0) + 0.5))
-            for px, py in line_points:
-                y_min, y_max, x_min, x_max = max(0, py - radius_int), min(h, py + radius_int + 1), max(0, px - radius_int), min(w, px + radius_int + 1)
-                if y_max > y_min and x_max > x_min:
-                    yy, xx = np.ogrid[y_min:y_max, x_min:x_max]
-                    canvas_np[y_min:y_max, x_min:x_max][(xx - px)**2 + (yy - py)**2 <= radius**2] = color
-        else:
-            line_points = np.array(line_points)
-            valid = (line_points[:, 1] >= 0) & (line_points[:, 1] < h) & (line_points[:, 0] >= 0) & (line_points[:, 0] < w)
-            if (valid_points := line_points[valid]).size:
-                canvas_np[valid_points[:, 1], valid_points[:, 0]] = color
-
-    @staticmethod
-    def fillConvexPoly(canvas_np, pts, color, **kwargs):
-        """Fill polygon using vectorized scanline algorithm."""
-        if len(pts) < 3:
-            return
-        pts = np.array(pts, dtype=np.int32)
-        h, w = canvas_np.shape[:2]
-        y_min, y_max, x_min, x_max = max(0, pts[:, 1].min()), min(h, pts[:, 1].max() + 1), max(0, pts[:, 0].min()), min(w, pts[:, 0].max() + 1)
-        if y_max <= y_min or x_max <= x_min:
-            return
-        yy, xx = np.mgrid[y_min:y_max, x_min:x_max]
-        mask = np.zeros((y_max - y_min, x_max - x_min), dtype=bool)
-
-        for i in range(len(pts)):
-            p1, p2 = pts[i], pts[(i + 1) % len(pts)]
-            y1, y2 = p1[1], p2[1]
-            if y1 == y2:
-                continue
-            if y1 > y2:
-                p1, p2, y1, y2 = p2, p1, p2[1], p1[1]
-            if not (edge_mask := (yy >= y1) & (yy < y2)).any():
-                continue
-            mask ^= edge_mask & (xx >= p1[0] + (yy - y1) * (p2[0] - p1[0]) / (y2 - y1))
-
-        canvas_np[y_min:y_max, x_min:x_max][mask] = color
-
-    @staticmethod
-    def ellipse2Poly(center, axes, angle, arc_start, arc_end, delta=1, **kwargs):
-        """Python implementation of cv2.ellipse2Poly."""
-        axes = (axes[0] + 0.5, axes[1] + 0.5) # to better match cv2 output
-        angle = angle % 360
-        if arc_start > arc_end:
-            arc_start, arc_end = arc_end, arc_start
-        while arc_start < 0:
-            arc_start, arc_end = arc_start + 360, arc_end + 360
-        while arc_end > 360:
-            arc_end, arc_start = arc_end - 360, arc_start - 360
-        if arc_end - arc_start > 360:
-            arc_start, arc_end = 0, 360
-
-        angle_rad = math.radians(angle)
-        alpha, beta = math.cos(angle_rad), math.sin(angle_rad)
-        pts = []
-        for i in range(arc_start, arc_end + delta, delta):
-            theta_rad = math.radians(min(i, arc_end))
-            x, y = axes[0] * math.cos(theta_rad), axes[1] * math.sin(theta_rad)
-            pts.append([int(round(center[0] + x * alpha - y * beta)), int(round(center[1] + x * beta + y * alpha))])
-
-        unique_pts, prev_pt = [], (float('inf'), float('inf'))
-        for pt in pts:
-            if (pt_tuple := tuple(pt)) != prev_pt:
-                unique_pts.append(pt)
-                prev_pt = pt_tuple
-
-        return unique_pts if len(unique_pts) > 1 else [[center[0], center[1]], [center[0], center[1]]]
-
-    def draw_wholebody_keypoints(self, canvas, keypoints, scores=None, threshold=0.3,
-                                 draw_body=True, draw_head=True, draw_feet=True, draw_face=True, draw_hands=True, stick_width=4, face_point_size=3):
-        """
-        Draw wholebody keypoints (134 keypoints after processing) in DWPose style.
-
-        Expected keypoint format (after neck insertion and remapping):
-        - Body: 0-17 (18 keypoints in OpenPose format, neck at index 1)
-        - Foot: 18-23 (6 keypoints)
-        - Face: 24-91 (68 landmarks)
-        - Right hand: 92-112 (21 keypoints)
-        - Left hand: 113-133 (21 keypoints)
-
-        Args:
-            canvas: The canvas to draw on (numpy array)
-            keypoints: Array of keypoint coordinates
-            scores: Optional confidence scores for each keypoint
-            threshold: Minimum confidence threshold for drawing keypoints
-
-        Returns:
-            canvas: The canvas with keypoints drawn
-        """
-        H, W, C = canvas.shape
-
-        # Draw body limbs & head connections
-        if (draw_body or draw_head) and len(keypoints) >= 18:
-            colorIndexOffset = 0
-            edges = []
-            if draw_body:
-                edges += self.body_limbSeq
-            else:
-                colorIndexOffset += len(self.body_limbSeq)
-            if draw_head:
-                edges += self.head_edges
-            for i, limb in enumerate(edges):
-                # Convert from 1-indexed to 0-indexed
-                idx1, idx2 = limb[0] - 1, limb[1] - 1
-
-                if idx1 >= 18 or idx2 >= 18:
-                    continue
-
-                if scores is not None:
-                    if scores[idx1] < threshold or scores[idx2] < threshold:
-                        continue
-
-                Y = [keypoints[idx1][0], keypoints[idx2][0]]
-                X = [keypoints[idx1][1], keypoints[idx2][1]]
-                mX, mY = (X[0] + X[1]) / 2, (Y[0] + Y[1]) / 2
-                length = math.sqrt((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2)
-
-                if length < 1:
-                    continue
-
-                angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
-
-                polygon = self.draw.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stick_width), int(angle), 0, 360, 1)
-
-                self.draw.fillConvexPoly(canvas, polygon, self.colors[(i + colorIndexOffset) % len(self.colors)])
-
-        # Draw body & head keypoints
-        if (draw_body or draw_head) and len(keypoints) >= 18:
-            head_keypoints = {0, 14, 15, 16, 17} # nose, eyes, ears
-            neck_point = 1
-            for i in range(18):
-                if not draw_head and i in head_keypoints:
-                    continue
-                if not draw_body and i not in head_keypoints and i != neck_point:
-                    continue
-                if scores is not None and scores[i] < threshold:
-                    continue
-                x, y = int(keypoints[i][0]), int(keypoints[i][1])
-                if 0 <= x < W and 0 <= y < H:
-                    self.draw.circle(canvas, (x, y), 4, self.colors[i % len(self.colors)], thickness=-1)
-
-        # Draw foot keypoints (18-23, 6 keypoints)
-        if draw_feet and len(keypoints) >= 24:
-            for i in range(18, 24):
-                if scores is not None and scores[i] < threshold:
-                    continue
-                x, y = int(keypoints[i][0]), int(keypoints[i][1])
-                if 0 <= x < W and 0 <= y < H:
-                    self.draw.circle(canvas, (x, y), 4, self.colors[i % len(self.colors)], thickness=-1)
-
-        # Draw right hand (92-112)
-        if draw_hands and len(keypoints) >= 113:
-            eps = 0.01
-            for ie, edge in enumerate(self.hand_edges):
-                idx1, idx2 = 92 + edge[0], 92 + edge[1]
-                if scores is not None:
-                    if scores[idx1] < threshold or scores[idx2] < threshold:
-                        continue
-
-                x1, y1 = int(keypoints[idx1][0]), int(keypoints[idx1][1])
-                x2, y2 = int(keypoints[idx2][0]), int(keypoints[idx2][1])
-
-                if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
-                    if 0 <= x1 < W and 0 <= y1 < H and 0 <= x2 < W and 0 <= y2 < H:
-                        # HSV to RGB conversion for rainbow colors
-                        r, g, b = colorsys.hsv_to_rgb(ie / float(len(self.hand_edges)), 1.0, 1.0)
-                        color = (int(r * 255), int(g * 255), int(b * 255))
-                        self.draw.line(canvas, (x1, y1), (x2, y2), color, thickness=2)
-
-            # Draw right hand keypoints
-            for i in range(92, 113):
-                if scores is not None and scores[i] < threshold:
-                    continue
-                x, y = int(keypoints[i][0]), int(keypoints[i][1])
-                if x > eps and y > eps and 0 <= x < W and 0 <= y < H:
-                    self.draw.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
-
-        # Draw left hand (113-133)
-        if draw_hands and len(keypoints) >= 134:
-            eps = 0.01
-            for ie, edge in enumerate(self.hand_edges):
-                idx1, idx2 = 113 + edge[0], 113 + edge[1]
-                if scores is not None:
-                    if scores[idx1] < threshold or scores[idx2] < threshold:
-                        continue
-
-                x1, y1 = int(keypoints[idx1][0]), int(keypoints[idx1][1])
-                x2, y2 = int(keypoints[idx2][0]), int(keypoints[idx2][1])
-
-                if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
-                    if 0 <= x1 < W and 0 <= y1 < H and 0 <= x2 < W and 0 <= y2 < H:
-                        # HSV to RGB conversion for rainbow colors
-                        r, g, b = colorsys.hsv_to_rgb(ie / float(len(self.hand_edges)), 1.0, 1.0)
-                        color = (int(r * 255), int(g * 255), int(b * 255))
-                        self.draw.line(canvas, (x1, y1), (x2, y2), color, thickness=2)
-
-            # Draw left hand keypoints
-            for i in range(113, 134):
-                if scores is not None and i < len(scores) and scores[i] < threshold:
-                    continue
-                x, y = int(keypoints[i][0]), int(keypoints[i][1])
-                if x > eps and y > eps and 0 <= x < W and 0 <= y < H:
-                    self.draw.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
-
-        # Draw face keypoints (24-91) - white dots only, no lines
-        if draw_face and len(keypoints) >= 92:
-            eps = 0.01
-            for i in range(24, 92):
-                if scores is not None and scores[i] < threshold:
-                    continue
-                x, y = int(keypoints[i][0]), int(keypoints[i][1])
-                if x > eps and y > eps and 0 <= x < W and 0 <= y < H:
-                    self.draw.circle(canvas, (x, y), face_point_size, (255, 255, 255), thickness=-1)
-
-        return canvas
-
 class SDPoseDrawKeypoints(io.ComfyNode):
     @classmethod
     def define_schema(cls):

comfy_extras/pose/keypoint_draw.py

@@ -41,12 +41,18 @@ class KeypointDraw:
             [0, 17], [17, 18], [18, 19], [19, 20], # pinky
         ]
 
-        # Body connections - matching DWPose limbSeq (1-indexed, converted to 0-indexed)
+        # Head connections (1-indexed, converted to 0-indexed): nose-neck, eyes, ears
+        self.head_edges = [
+            [2, 1], [1, 15], [15, 17], [1, 16], [16, 18]
+        ]
+
+        # Body connections - matching DWPose limbSeq (1-indexed, converted to 0-indexed).
+        # body_limbSeq is the full 18-point skeleton (body + head_edges last); the head
+        # edges are kept as the trailing entries so callers can toggle them via draw_head.
         self.body_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],
+            [10, 11], [2, 12], [12, 13], [13, 14],
-            [1, 16], [16, 18]
+        ] + self.head_edges
-        ]
 
         # Colors matching DWPose
         self.colors = [
@@ -163,7 +169,7 @@ class KeypointDraw:
         return unique_pts if len(unique_pts) > 1 else [[center[0], center[1]], [center[0], center[1]]]
 
     def draw_wholebody_keypoints(self, canvas, keypoints, scores=None, threshold=0.3,
-                                 draw_body=True, draw_feet=True, draw_face=True, draw_hands=True,
+                                 draw_body=True, draw_head=True, draw_feet=True, draw_face=True, draw_hands=True,
                                  stick_width=4, face_point_size=3,
                                  marker_radius=4, hand_stick_width=2, hand_marker_radius=4,
                                  limb_alpha=1.0, hand_dot_color=(0, 0, 255)):
@@ -182,6 +188,7 @@ class KeypointDraw:
             keypoints: Array of keypoint coordinates
             scores: Optional confidence scores for each keypoint
             threshold: Minimum confidence threshold for drawing keypoints
+            draw_head: Toggle head edges/keypoints (nose, eyes, ears) independently of draw_body.
             stick_width: Body limb half-width (passed to ellipse2Poly).
             face_point_size: Radius of the white face dots.
             marker_radius: Radius of body/foot dots. Defaults to 4 (DWPose).
@@ -207,9 +214,19 @@ class KeypointDraw:
 
         do_alpha = float(limb_alpha) < 1.0
 
-        # Draw body limbs
+        # Draw body limbs & head connections. body_limbSeq holds the full skeleton
-        if draw_body and len(keypoints) >= 18:
+        # with head edges trailing; draw_body / draw_head toggle each group while the
-            for i, limb in enumerate(self.body_limbSeq):
+        # color index stays aligned to the full sequence.
+        if (draw_body or draw_head) and len(keypoints) >= 18:
+            body_core = self.body_limbSeq[:len(self.body_limbSeq) - len(self.head_edges)]
+            edges, color_offset = [], 0
+            if draw_body:
+                edges += body_core
+            else:
+                color_offset += len(body_core)
+            if draw_head:
+                edges += self.head_edges
+            for i, limb in enumerate(edges):
                 # Convert from 1-indexed to 0-indexed
                 idx1, idx2 = limb[0] - 1, limb[1] - 1
 
@@ -232,15 +249,21 @@ class KeypointDraw:
 
                 polygon = self.draw.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stick_width), int(angle), 0, 360, 1)
 
-                color = self.colors[i % len(self.colors)]
+                color = self.colors[(i + color_offset) % len(self.colors)]
                 if do_alpha:
                     _fill_poly_alpha(canvas, polygon, color, limb_alpha, self.draw)
                 else:
                     self.draw.fillConvexPoly(canvas, polygon, color)
 
-        # Draw body keypoints
+        # Draw body & head keypoints
-        if draw_body and len(keypoints) >= 18:
+        if (draw_body or draw_head) and len(keypoints) >= 18:
+            head_keypoints = {0, 14, 15, 16, 17} # nose, eyes, ears
+            neck_point = 1
             for i in range(18):
+                if not draw_head and i in head_keypoints:
+                    continue
+                if not draw_body and i not in head_keypoints and i != neck_point:
+                    continue
                 if scores is not None and scores[i] < threshold:
                     continue
                 x, y = int(keypoints[i][0]), int(keypoints[i][1])