127 lines
4.8 KiB
Python
127 lines
4.8 KiB
Python
import json
|
|
import os
|
|
import random
|
|
import numpy as np
|
|
from typing import List
|
|
import math
|
|
import cv2
|
|
import skeleton_lib as skel
|
|
import process_json_file as pjf
|
|
import sys
|
|
sys.path.append('./')
|
|
|
|
def is_normalized(keypoints: List[skel.Keypoint]) -> bool:
|
|
for keypoint in keypoints:
|
|
if not (0 <= keypoint.x <= 1 and 0 <= keypoint.y <= 1):
|
|
return False
|
|
return True
|
|
|
|
def draw_bodypose(canvas: np.ndarray, keypoints: List[skel.Keypoint], limbSeq, colors, xinsr_stick_scaling: bool = False) -> np.ndarray:
|
|
"""
|
|
Draw keypoints and limbs representing body pose on a given canvas.
|
|
|
|
Args:
|
|
canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the body pose.
|
|
keypoints (List[Keypoint]): A list of Keypoint objects representing the body keypoints to be drawn.
|
|
xinsr_stick_scaling (bool): Whether or not scaling stick width for xinsr ControlNet
|
|
|
|
Returns:
|
|
np.ndarray: A 3D numpy array representing the modified canvas with the drawn body pose.
|
|
|
|
Note:
|
|
The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
|
|
"""
|
|
if not is_normalized(keypoints):
|
|
H, W = 1.0, 1.0
|
|
else:
|
|
H, W, _ = canvas.shape
|
|
|
|
CH, CW, _ = canvas.shape
|
|
stickwidth = 2
|
|
|
|
# Ref: https://huggingface.co/xinsir/controlnet-openpose-sdxl-1.0
|
|
max_side = max(CW, CH)
|
|
if xinsr_stick_scaling:
|
|
stick_scale = 1 if max_side < 500 else min(2 + (max_side // 1000), 7)
|
|
else :
|
|
stick_scale = 1
|
|
|
|
if keypoints is None or len(keypoints) == 0:
|
|
return canvas
|
|
|
|
for (k1_index, k2_index), color in zip(limbSeq, colors):
|
|
keypoint1 = keypoints[k1_index]
|
|
keypoint2 = keypoints[k2_index]
|
|
|
|
if keypoint1 is None or keypoint2 is None or keypoint1.confidence == 0 or keypoint2.confidence == 0:
|
|
# if keypoint1 is None or keypoint1.confidence == 0:
|
|
# print(f"keypoint failed: {k1_index}")
|
|
# if keypoint2 is None or keypoint2.confidence == 0:
|
|
# print(f"keypoint failed: {k2_index}")
|
|
continue
|
|
|
|
Y = np.array([keypoint1.x, keypoint2.x]) * float(W)
|
|
X = np.array([keypoint1.y, keypoint2.y]) * float(H)
|
|
mX = np.mean(X)
|
|
mY = np.mean(Y)
|
|
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
|
|
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
|
|
polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth*stick_scale), int(angle), 0, 360, 1)
|
|
cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])
|
|
|
|
for keypoint, color in zip(keypoints, colors):
|
|
if keypoint is None or keypoint.confidence == 0:
|
|
continue
|
|
|
|
x, y = keypoint.x, keypoint.y
|
|
x = int(x * W)
|
|
y = int(y * H)
|
|
cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
|
|
|
|
return canvas
|
|
|
|
def coordinates_to_keypoints(coordinates: list) -> List[skel.Keypoint]:
|
|
keypoints = [skel.Keypoint(coordinates[i], coordinates[i + 1]) for i in range(0, len(coordinates), 3)]
|
|
return keypoints
|
|
|
|
def save_bodypose(width: int, height: int, coordinates: list):
|
|
if not hasattr(save_bodypose, 'counter'):
|
|
save_bodypose.counter = 0 # Initialize the counter attribute
|
|
|
|
canvas = np.zeros((height, width, 3), dtype=np.uint8)
|
|
keypoints = coordinates_to_keypoints(coordinates)
|
|
canvas = draw_bodypose(canvas, keypoints, skel.coco_limbSeq, skel.coco_colors)
|
|
|
|
# Save as body_pose_output0000.png, body_pose_output0001.png, ...
|
|
cv2.imwrite('output/body_pose_output%04d.png' % save_bodypose.counter, canvas)
|
|
save_bodypose.counter += 1 # Increment the counter
|
|
|
|
def main():
|
|
directory = './fixed'
|
|
json_files = [f for f in os.listdir(directory) if f.endswith('.json')]
|
|
if not json_files:
|
|
print("No JSON files found in the directory.")
|
|
return
|
|
|
|
json_file = os.path.join(directory, random.choice(json_files))
|
|
# json_file = './test_output.json'
|
|
image_path = './output/test'
|
|
print(json_file)
|
|
|
|
skeleton_sequences = pjf.array_json_to_Skeleton_Seqences(json_file)
|
|
frame_count = max(len(skeleton_sequences[i].skeletons_frame) for i in range(len(skeleton_sequences)) if skeleton_sequences[i] is not None)
|
|
sliced_list = [skel.get_time_slice_for_Skeleton_Seqences(skeleton_sequences, i) for i in range(frame_count)]
|
|
|
|
for i in range(frame_count):
|
|
sliced = sliced_list[i]
|
|
canvas = np.zeros((360, 640, 3), dtype=np.uint8)
|
|
|
|
for j, skeleton in enumerate(sliced):
|
|
keypoints = skeleton.keypoints
|
|
skeleton_sequences[j].get_frame(i).keypoints = keypoints
|
|
canvas = draw_bodypose(canvas, keypoints, skel.body_25_limbSeq, skel.body_25_colors)
|
|
|
|
cv2.imwrite(image_path + '_' + str(i) + '.png', canvas)
|
|
|
|
if __name__ == '__main__':
|
|
main() |