pose2sim/Pose2Sim/Utilities/DLC_to_OpenPose.py at main · perfanalytics/pose2sim · GitHub

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
94
95
96
97
98
#!/usr/bin/env python
# -*- coding: utf-8 -*-


'''
    ########################################################
    ## Convert DeepLabCut h5 files to OpenPose json files ##
    ########################################################

    Translates DeepLabCut (h5) 2D pose estimation files into OpenPose (json) files.
    You may need to install tables: 'pip install tables' or 'conda install pytables'

    Usage:
    python -m DLC_to_OpenPose -i input_h5_file -o output_json_folder
    OR python -m DLC_to_OpenPose -i input_h5_file
    OR from Pose2Sim.Utilities import DLC_to_OpenPose; DLC_to_OpenPose.DLC_to_OpenPose_func(r'input_h5_file', r'output_json_folder')
'''


## INIT
import pandas as pd
import numpy as np
import os
import json
import re
import argparse


## AUTHORSHIP INFORMATION
__author__ = "David Pagnon"
__copyright__ = "Copyright 2021, Pose2Sim"
__credits__ = ["David Pagnon"]
__license__ = "BSD 3-Clause License"
__version__ = "0.9.4"
__maintainer__ = "David Pagnon"
__email__ = "contact@david-pagnon.com"
__status__ = "Development"


## FUNCTIONS
def DLC_to_OpenPose_func(*args):
    '''
    Translates DeepLabCut (h5) 2D pose estimation files into OpenPose (json) files.

    Usage:
    DLC_to_OpenPose -i input_h5_file -o output_json_folder
    OR DLC_to_OpenPose -i input_h5_file
    OR import DLC_to_OpenPose; DLC_to_OpenPose.DLC_to_OpenPose_func(r'input_h5_file', r'output_json_folder')
    '''

    try:
        h5_file_path = os.path.realpath(args[0]['input']) # invoked with argparse
        if args[0]['output'] == None:
            json_folder_path = os.path.splitext(h5_file_path)[0]
        else:
            json_folder_path = os.path.realpath(args[0]['output'])
    except:
        h5_file_path = os.path.realpath(args[0]) # invoked as a function
        try:
            json_folder_path = os.path.realpath(args[1])
        except:
            json_folder_path = os.path.splitext(h5_file_path)[0]

    if not os.path.exists(json_folder_path):
        os.mkdir(json_folder_path)

    # json preparation
    json_dict = {'version':1.3, 'people':[]}
    json_dict['people'] = [{'person_id':[-1],
                    'pose_keypoints_2d': [],
                    'face_keypoints_2d': [],
                    'hand_left_keypoints_2d':[],
                    'hand_right_keypoints_2d':[],
                    'pose_keypoints_3d':[],
                    'face_keypoints_3d':[],
                    'hand_left_keypoints_3d':[],
                    'hand_right_keypoints_3d':[]}]

    # h5 reader
    h5_file = pd.read_hdf(h5_file_path).fillna(0)
    kpt_nb = int(len(h5_file.columns)//3)

    # write each h5 line in json file
    for f, frame in enumerate(h5_file.index):
        h5_line = np.array([[h5_file.iloc[f, 3*k], h5_file.iloc[f, 3*k+1], h5_file.iloc[f, 3*k+2]] for k in range(kpt_nb)]).flatten().tolist()
        json_dict['people'][0]['pose_keypoints_2d'] = h5_line
        json_file = os.path.join(json_folder_path, os.path.splitext(os.path.basename(str(frame).zfill(5)))[0]+'.json')
        with open(json_file, 'w') as js_f:
            js_f.write(json.dumps(json_dict))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', required = True, help='input 2D pose coordinates DeepLabCut h5 file')
    parser.add_argument('-o', '--output', required = False, help='output folder for 2D pose coordinates OpenPose json files')
    args = vars(parser.parse_args())

    DLC_to_OpenPose_func(args)