diff --git a/comfy_extras/nodes_sdpose.py b/comfy_extras/nodes_sdpose.py index 46b5fb226..2fa39872e 100644 --- a/comfy_extras/nodes_sdpose.py +++ b/comfy_extras/nodes_sdpose.py @@ -251,6 +251,10 @@ class KeypointDraw: Y = [keypoints[idx1][0], keypoints[idx2][0]] X = [keypoints[idx1][1], keypoints[idx2][1]] + + if math.isnan(X[0]) or math.isnan(X[1]) or math.isnan(Y[0]) or math.isnan(Y[1]): + continue + 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) @@ -268,6 +272,8 @@ class KeypointDraw: for i in range(18): if scores is not None and scores[i] < threshold: continue + if math.isnan(keypoints[i][0]) or math.isnan(keypoints[i][1]): + 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) @@ -277,6 +283,8 @@ class KeypointDraw: for i in range(18, 24): if scores is not None and scores[i] < threshold: continue + if math.isnan(keypoints[i][0]) or math.isnan(keypoints[i][1]): + 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) @@ -290,6 +298,9 @@ class KeypointDraw: if scores[idx1] < threshold or scores[idx2] < threshold: continue + if math.isnan(keypoints[idx1][0]) or math.isnan(keypoints[idx1][1]) or math.isnan(keypoints[idx2][0]) or math.isnan(keypoints[idx2][1]): + continue + x1, y1 = int(keypoints[idx1][0]), int(keypoints[idx1][1]) x2, y2 = int(keypoints[idx2][0]), int(keypoints[idx2][1]) @@ -304,6 +315,8 @@ class KeypointDraw: for i in range(92, 113): if scores is not None and scores[i] < threshold: continue + if math.isnan(keypoints[i][0]) or math.isnan(keypoints[i][1]): + 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) @@ -317,6 +330,9 @@ class KeypointDraw: if scores[idx1] < threshold or scores[idx2] < threshold: continue + if math.isnan(keypoints[idx1][0]) or math.isnan(keypoints[idx1][1]) or math.isnan(keypoints[idx2][0]) or math.isnan(keypoints[idx2][1]): + continue + x1, y1 = int(keypoints[idx1][0]), int(keypoints[idx1][1]) x2, y2 = int(keypoints[idx2][0]), int(keypoints[idx2][1]) @@ -331,6 +347,8 @@ class KeypointDraw: for i in range(113, 134): if scores is not None and i < len(scores) and scores[i] < threshold: continue + if math.isnan(keypoints[i][0]) or math.isnan(keypoints[i][1]): + 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) @@ -341,6 +359,8 @@ class KeypointDraw: for i in range(24, 92): if scores is not None and scores[i] < threshold: continue + if math.isnan(keypoints[i][0]) or math.isnan(keypoints[i][1]): + 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)