Spaces:
Running
on
L40S
Running
on
L40S
File size: 2,859 Bytes
2252f3d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import os
import json
import shutil
import subprocess
import numpy as np
import os.path as osp
def run_openpose(
video_file,
output_folder,
staf_folder,
vis=False,
):
pwd = os.getcwd()
os.chdir(staf_folder)
render = 1 if vis else 0
display = 2 if vis else 0
cmd = [
'build/examples/openpose/openpose.bin', '--model_pose', 'BODY_21A', '--tracking', '1',
'--render_pose',
str(render), '--video', video_file, '--write_json', output_folder, '--display',
str(display)
]
print('Executing', ' '.join(cmd))
subprocess.call(cmd)
os.chdir(pwd)
def read_posetrack_keypoints(output_folder):
people = dict()
for idx, result_file in enumerate(sorted(os.listdir(output_folder))):
json_file = osp.join(output_folder, result_file)
data = json.load(open(json_file))
# print(idx, data)
for person in data['people']:
person_id = person['person_id'][0]
joints2d = person['pose_keypoints_2d']
if person_id in people.keys():
people[person_id]['joints2d'].append(joints2d)
people[person_id]['frames'].append(idx)
else:
people[person_id] = {
'joints2d': [],
'frames': [],
}
people[person_id]['joints2d'].append(joints2d)
people[person_id]['frames'].append(idx)
for k in people.keys():
people[k]['joints2d'] = np.array(people[k]['joints2d']).reshape(
(len(people[k]['joints2d']), -1, 3)
)
people[k]['frames'] = np.array(people[k]['frames'])
return people
def run_posetracker(video_file, staf_folder, posetrack_output_folder='/tmp', display=False):
posetrack_output_folder = os.path.join(
posetrack_output_folder, f'{os.path.basename(video_file)}_posetrack'
)
# run posetrack on video
run_openpose(video_file, posetrack_output_folder, vis=display, staf_folder=staf_folder)
people_dict = read_posetrack_keypoints(posetrack_output_folder)
shutil.rmtree(posetrack_output_folder)
return people_dict
|