135 lines
5.1 KiB
Python
135 lines
5.1 KiB
Python
import json
|
|
import os
|
|
import random
|
|
from typing import List
|
|
import cv2
|
|
import numpy as np
|
|
import urllib
|
|
import skeleton_lib as skel
|
|
import process_json_file as pjf
|
|
from urllib import request
|
|
from requests_toolbelt.multipart.encoder import MultipartEncoder
|
|
import sys
|
|
import hashlib
|
|
sys.path.append('./')
|
|
|
|
server_address = "localhost:8188"
|
|
|
|
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, pid: str) -> None:
|
|
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 = skel.draw_bodypose(canvas, keypoints, skel.coco_limbSeq, skel.coco_colors)
|
|
|
|
# Save as body_pose_output0000.png, body_pose_output0001.png, ...
|
|
output_dir = 'output'
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
image_path = 'output/body_pose_output%04d.png' % save_bodypose.counter
|
|
cv2.imwrite(image_path, canvas)
|
|
save_bodypose.counter += 1 # Increment the counter
|
|
|
|
return image_path
|
|
|
|
def save_bodypose_mulit(width: int, height: int, coordinates_list: list, pid: str) -> None:
|
|
if not hasattr(save_bodypose_mulit, 'counter'):
|
|
save_bodypose_mulit.counter = 0 # Initialize the counter attribute
|
|
|
|
canvas = np.zeros((height, width, 3), dtype=np.uint8)
|
|
for coordinates in coordinates_list:
|
|
keypoints = coordinates_to_keypoints(coordinates)
|
|
canvas = skel.draw_bodypose(canvas, keypoints, skel.coco_limbSeq, skel.coco_colors)
|
|
|
|
# Save as body_pose_output0000.png, body_pose_output0001.png, ...
|
|
output_dir = 'output'
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
|
|
image_path = 'output/body_pose_output_multi%04d.png' % save_bodypose_mulit.counter
|
|
cv2.imwrite(image_path, canvas)
|
|
save_bodypose_mulit.counter += 1 # Increment the counter
|
|
|
|
return image_path
|
|
|
|
def queue_prompt(prompt, server_address):
|
|
p = {"prompt": prompt}
|
|
data = json.dumps(p).encode('utf-8')
|
|
req = request.Request("http://{}/prompt".format(server_address), data=data)
|
|
request.urlopen(req)
|
|
|
|
def upload_image(input_image, name, server_address, image_type="input", overwrite=True):
|
|
# Check if input_image is a valid file path
|
|
if isinstance(input_image, str) and os.path.isfile(input_image):
|
|
file = open(input_image, 'rb')
|
|
close_file = True
|
|
else:
|
|
file = input_image
|
|
close_file = False
|
|
|
|
try:
|
|
multipart_data = MultipartEncoder(
|
|
fields={
|
|
'image': (name, file, 'image/png'),
|
|
'type': image_type,
|
|
'overwrite': str(overwrite).lower()
|
|
}
|
|
)
|
|
|
|
data = multipart_data
|
|
headers = {'Content-Type': multipart_data.content_type}
|
|
request = urllib.request.Request("http://{}/upload/image".format(server_address), data=data, headers=headers)
|
|
with urllib.request.urlopen(request) as response:
|
|
return response.read()
|
|
finally:
|
|
if close_file:
|
|
file.close()
|
|
|
|
def upload_image_circular_queue(image_path, size, unqiue_id, server_address):
|
|
# create a dict in this function to store the counter for each unique_id, key is the unique_id, value is the counter
|
|
if not hasattr(upload_image_circular_queue, 'id_counter_dict'):
|
|
upload_image_circular_queue.id_counter_dict = {}
|
|
|
|
if unqiue_id not in upload_image_circular_queue.id_counter_dict:
|
|
upload_image_circular_queue.id_counter_dict[unqiue_id] = 0
|
|
|
|
image_name = hashlib.sha256((unqiue_id + str(upload_image_circular_queue.id_counter_dict[unqiue_id])).encode('utf-8')).hexdigest() + ".png"
|
|
upload_image_circular_queue.id_counter_dict[unqiue_id] += 1 % size
|
|
upload_image(image_path, image_name, server_address)
|
|
|
|
return image_name
|
|
|
|
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 = skel.draw_bodypose(canvas, keypoints, skel.body_25_limbSeq, skel.body_25_colors)
|
|
|
|
cv2.imwrite(image_path + '_' + str(i) + '.png', canvas)
|
|
|
|
if __name__ == '__main__':
|
|
main() |