44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
"""
|
|
Used to process keypoints generated from ControlNet extension.
|
|
"""
|
|
import json
|
|
from typing import List, Tuple
|
|
|
|
def process_keypoints(nums: List[float], width: int, height: int) -> List[List[float]]:
|
|
if not nums:
|
|
return []
|
|
|
|
assert len(nums) % 3 == 0
|
|
|
|
def find_min(nums: float):
|
|
return min(num for num in nums if num > 0)
|
|
|
|
base_x = find_min(nums[::3])
|
|
base_y = find_min(nums[1::3])
|
|
|
|
normalized = all(abs(num) <= 1.0 for num in nums)
|
|
x_factor = width if normalized else 1.0
|
|
y_factor = height if normalized else 1.0
|
|
|
|
return [
|
|
[(x-base_x) * x_factor, (y-base_y) * y_factor, c]
|
|
for x, y, c in zip(nums[::3], nums[1::3], nums[2::3])
|
|
]
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
with open('pose_hand_fixed.json', 'r') as f:
|
|
pose = json.load(f)
|
|
person = pose["people"][0]
|
|
|
|
with open('output_pose.json', 'w') as f:
|
|
width = pose['canvas_width']
|
|
height = pose['canvas_height']
|
|
|
|
json.dump({
|
|
'left_hand': process_keypoints(person.get('hand_left_keypoints_2d'), width, height),
|
|
'right_hand': process_keypoints(person.get('hand_right_keypoints_2d'), width, height),
|
|
'face': process_keypoints(person.get('face_keypoints_2d'), width, height),
|
|
}, f, indent=4)
|