48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
import json
|
|
import numpy as np
|
|
|
|
class PoseDiffCalculator:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"original_frame_pose": ("POSE_KEYPOINT", {}),
|
|
"adjusted_frame_pose": ("POSE_KEYPOINT", {})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("POSE_KEYPOINT",)
|
|
RETURN_NAMES = ("pose_diff",)
|
|
FUNCTION = "calculate_diff"
|
|
CATEGORY = "Snap Processing"
|
|
|
|
def extract_pose(self, data):
|
|
if isinstance(data, list):
|
|
data = data[0]
|
|
try:
|
|
return np.array(data['people'][0]['pose_keypoints_2d'])
|
|
except (IndexError, KeyError, TypeError):
|
|
raise ValueError("Input pose data is invalid or malformed.")
|
|
|
|
def calculate_diff(self, original_frame_pose, adjusted_frame_pose):
|
|
original_pose = self.extract_pose(original_frame_pose)
|
|
adjusted_pose = self.extract_pose(adjusted_frame_pose)
|
|
|
|
# 检查数据有效性
|
|
if original_pose.size == 0 or adjusted_pose.size == 0 or original_pose.size != adjusted_pose.size:
|
|
raise ValueError("Input pose data is invalid or keypoints are mismatched.")
|
|
|
|
# 计算差值
|
|
pose_diff = (adjusted_pose - original_pose).tolist()
|
|
|
|
# 构建差值数据
|
|
diff_result = [{
|
|
"people": [
|
|
{
|
|
"pose_keypoints_2d": pose_diff
|
|
}
|
|
]
|
|
}]
|
|
|
|
return (diff_result,)
|