posetrack.py

import os
import json
import shutil
import subprocess
import numpy as np
import os.path as osp

def run_openpose(
        video_file,
        output_folder,
        vis=False,
        openposetrack_folder='/home/mkocabas/developments/openposetrack'
):
    pwd = os.getcwd()

    os.chdir(openposetrack_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, posetrack_output_folder='/tmp/posetrack_output', display=False):
    # run posetrack on video
    run_openpose(video_file, posetrack_output_folder, vis=display)
    people_dict = read_posetrack_keypoints(posetrack_output_folder)
    shutil.rmtree(posetrack_output_folder)

    # se
    person_id = 0
    max = 0
    for k in people_dict.keys():
        nf = people_dict[k]['frames'].shape[0]
        if nf > max:
            person_id = k
            max = nf

    # person_id = 11
    print(f'Using person {person_id}...')

    joints2d = people_dict[person_id]['joints2d']
    frames = people_dict[person_id]['frames']

    return joints2d, frames