99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
#!/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.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)
|