102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
import cv2
|
|
import time
|
|
from hpe_mp_class import hpe_mp_class
|
|
import json
|
|
import argparse
|
|
from pythonosc import udp_client
|
|
import yaml
|
|
|
|
# Arguments
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--config_path', type=str, default="config.yml", help='path to config file yaml')
|
|
args = parser.parse_args()
|
|
config_path = args.config_path
|
|
|
|
# Config parameters
|
|
data_yaml = yaml.safe_load(open(config_path))
|
|
|
|
address_input = data_yaml["address_input"]
|
|
for i in range(0, 100):
|
|
if address_input == str(i):
|
|
address_input = i
|
|
scale_pose = data_yaml["scale_pose"]
|
|
osc_address = data_yaml["osc_address"]
|
|
osc_port = data_yaml["osc_port"]
|
|
osc_message_address = data_yaml["osc_message_address"]
|
|
output_method = data_yaml["output_method"]
|
|
crop = data_yaml["crop_image"]
|
|
mirror_image = data_yaml["mirror_image"]
|
|
show_image = data_yaml["show_image"]
|
|
apose = data_yaml["apose"]
|
|
world = data_yaml["world"]
|
|
|
|
# OSC client
|
|
client = udp_client.SimpleUDPClient(osc_address, osc_port)
|
|
|
|
# Videocapture
|
|
cap = cv2.VideoCapture(address_input)
|
|
|
|
# Preprocessing parameters
|
|
frame_width = int(crop*cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
frame_height = int(crop*cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
|
|
# FPS variables
|
|
pTime = 0
|
|
cTime = 0
|
|
|
|
# Mediapie class
|
|
mp_cl = hpe_mp_class()
|
|
|
|
while True:
|
|
# Reading frame
|
|
success, img = cap.read()
|
|
|
|
# Image preprocessing
|
|
if crop != 1.0:
|
|
img = cv2.resize(img, (frame_width, frame_height))
|
|
|
|
# Mediapipe
|
|
mp_cl.process(img, scale_pose=scale_pose)
|
|
mp_cl.show(img)
|
|
|
|
# FPS
|
|
cTime = time.time()
|
|
fps = 1. / (cTime - pTime)
|
|
pTime = cTime
|
|
|
|
# Showing
|
|
if show_image:
|
|
if mirror_image:
|
|
img = cv2.flip(img, 1) # mirror
|
|
cv2.putText(img, str(int(fps)), (22, 32), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 0), 2)
|
|
cv2.putText(img, str(int(fps)), (20, 30), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 2)
|
|
cv2.imshow("Main", img)
|
|
|
|
# Output
|
|
if output_method == 'file':
|
|
# JSON
|
|
res = mp_cl.getJSON(apose=apose, world=world)
|
|
with open('hierarchy_data.json', 'w', encoding='utf-8') as f:
|
|
json.dump(res, f, ensure_ascii=False, indent=4)
|
|
else:
|
|
# OSC
|
|
res = mp_cl.getJSON(apose=apose, world=world)
|
|
res_list = []
|
|
for val in res.keys():
|
|
stroka = str(val)
|
|
for val2 in res[val]:
|
|
if val2 == 'visible':
|
|
stroka += " " + str(val2) + " " + str(res[val][val2])
|
|
# res_list.append(str(val) + " " + str(val2) + " " + str(res[val][val2]))
|
|
else:
|
|
for val3 in res[val][val2]:
|
|
stroka += " " + str(val2) + "_" + str(val3) + " " + str(res[val][val2][val3])
|
|
# res_list.append(str(val) + " " + str(val2) + " " + str(val3) + " " + str(res[val][val2][val3]))
|
|
res_list.append(stroka)
|
|
client.send_message(osc_message_address, res_list)
|
|
|
|
# Interface
|
|
key = cv2.waitKey(1)
|
|
if key == 27:
|
|
break
|