80 lines
3.2 KiB
Python
80 lines
3.2 KiB
Python
# ComfyUI_DWposeDeluxe/scripts/keypoint_utils.py
|
|
|
|
import json
|
|
import copy
|
|
import numpy as np
|
|
from .logger import logger
|
|
|
|
class UnifiedKeypointHandler:
|
|
@staticmethod
|
|
def detect_format(data, soft_boundary=1.1):
|
|
"""
|
|
Detects if keypoints are 'normalized' or 'absolute'.
|
|
Uses a soft boundary to allow for points slightly outside the [0, 1] range.
|
|
"""
|
|
if not data:
|
|
return "unknown"
|
|
|
|
# Check a sample of points from the first few frames
|
|
for frame in data[:5]:
|
|
for person in frame.get("people", []):
|
|
for key in ["pose_keypoints_2d", "face_keypoints_2d", "hand_left_keypoints_2d", "hand_right_keypoints_2d"]:
|
|
pts = person.get(key, [])
|
|
if pts:
|
|
# Check first 5 triplets (x, y, conf)
|
|
for i in range(0, min(len(pts), 15), 3):
|
|
if abs(pts[i]) > soft_boundary or abs(pts[i+1]) > soft_boundary:
|
|
return "absolute"
|
|
return "normalized"
|
|
|
|
@staticmethod
|
|
def to_normalized(data, width, height):
|
|
"""Converts absolute keypoints to normalized (0-1)."""
|
|
if width <= 0 or height <= 0:
|
|
return data
|
|
|
|
new_data = copy.deepcopy(data)
|
|
for frame in new_data:
|
|
for person in frame.get("people", []):
|
|
for key in ["pose_keypoints_2d", "face_keypoints_2d", "hand_left_keypoints_2d", "hand_right_keypoints_2d"]:
|
|
if key in person:
|
|
pts = person[key]
|
|
for i in range(0, len(pts), 3):
|
|
pts[i] /= width
|
|
pts[i+1] /= height
|
|
frame["canvas_width"] = width
|
|
frame["canvas_height"] = height
|
|
return new_data
|
|
|
|
@staticmethod
|
|
def to_absolute(data, width, height):
|
|
"""Converts normalized keypoints to absolute pixels."""
|
|
if width <= 0 or height <= 0:
|
|
return data
|
|
|
|
new_data = copy.deepcopy(data)
|
|
for frame in new_data:
|
|
for person in frame.get("people", []):
|
|
for key in ["pose_keypoints_2d", "face_keypoints_2d", "hand_left_keypoints_2d", "hand_right_keypoints_2d"]:
|
|
if key in person:
|
|
pts = person[key]
|
|
for i in range(0, len(pts), 3):
|
|
pts[i] *= width
|
|
pts[i+1] *= height
|
|
frame["canvas_width"] = width
|
|
frame["canvas_height"] = height
|
|
return new_data
|
|
|
|
@staticmethod
|
|
def reset_confidence(data, threshold=0.25):
|
|
"""Sets confidence to 1.0 if >= threshold, else 0.0. For external compatibility."""
|
|
new_data = copy.deepcopy(data)
|
|
for frame in new_data:
|
|
for person in frame.get("people", []):
|
|
for key in ["pose_keypoints_2d", "face_keypoints_2d", "hand_left_keypoints_2d", "hand_right_keypoints_2d"]:
|
|
if key in person:
|
|
pts = person[key]
|
|
for i in range(2, len(pts), 3):
|
|
pts[i] = 1.0 if pts[i] >= threshold else 0.0
|
|
return new_data
|