194 lines
7.2 KiB
Python
194 lines
7.2 KiB
Python
'''
|
|
@ Date: 2021-03-15 12:23:12
|
|
@ Author: Qing Shuai
|
|
@ LastEditors: Qing Shuai
|
|
@ LastEditTime: 2021-05-27 20:50:43
|
|
@ FilePath: /EasyMocapRelease/easymocap/mytools/file_utils.py
|
|
'''
|
|
import os
|
|
import json
|
|
import numpy as np
|
|
from os.path import join
|
|
|
|
mkdir = lambda x:os.makedirs(x, exist_ok=True)
|
|
mkout = lambda x:mkdir(os.path.dirname(x))
|
|
|
|
def read_json(path):
|
|
assert os.path.exists(path), path
|
|
with open(path) as f:
|
|
data = json.load(f)
|
|
return data
|
|
|
|
def save_json(file, data):
|
|
if not os.path.exists(os.path.dirname(file)):
|
|
os.makedirs(os.path.dirname(file))
|
|
with open(file, 'w') as f:
|
|
json.dump(data, f, indent=4)
|
|
|
|
def getFileList(root, ext='.jpg'):
|
|
files = []
|
|
dirs = os.listdir(root)
|
|
while len(dirs) > 0:
|
|
path = dirs.pop()
|
|
fullname = join(root, path)
|
|
if os.path.isfile(fullname) and fullname.endswith(ext):
|
|
files.append(path)
|
|
elif os.path.isdir(fullname):
|
|
for s in os.listdir(fullname):
|
|
newDir = join(path, s)
|
|
dirs.append(newDir)
|
|
files = sorted(files)
|
|
return files
|
|
|
|
def read_annot(annotname, mode='body25'):
|
|
data = read_json(annotname)
|
|
if not isinstance(data, list):
|
|
data = data['annots']
|
|
for i in range(len(data)):
|
|
if 'id' not in data[i].keys():
|
|
data[i]['id'] = data[i].pop('personID')
|
|
if 'keypoints2d' in data[i].keys() and 'keypoints' not in data[i].keys():
|
|
data[i]['keypoints'] = data[i].pop('keypoints2d')
|
|
for key in ['bbox', 'keypoints', 'handl2d', 'handr2d', 'face2d']:
|
|
if key not in data[i].keys():continue
|
|
data[i][key] = np.array(data[i][key])
|
|
if key == 'face2d':
|
|
# TODO: Make parameters, 17 is the offset for the eye brows,
|
|
# etc. 51 is the total number of FLAME compatible landmarks
|
|
data[i][key] = data[i][key][17:17+51, :]
|
|
data[i]['bbox'] = data[i]['bbox'][:5]
|
|
if data[i]['bbox'][-1] < 0.001:
|
|
# print('{}/{} bbox conf = 0, may be error'.format(annotname, i))
|
|
data[i]['bbox'][-1] = 1
|
|
if mode == 'body25':
|
|
data[i]['keypoints'] = data[i]['keypoints']
|
|
elif mode == 'body15':
|
|
data[i]['keypoints'] = data[i]['keypoints'][:15, :]
|
|
elif mode in ['handl', 'handr']:
|
|
data[i]['keypoints'] = np.array(data[i][mode+'2d']).astype(np.float32)
|
|
key = 'bbox_'+mode+'2d'
|
|
if key not in data[i].keys():
|
|
data[i]['bbox'] = np.array(get_bbox_from_pose(data[i]['keypoints'])).astype(np.float32)
|
|
else:
|
|
data[i]['bbox'] = data[i]['bbox_'+mode+'2d'][:5]
|
|
elif mode == 'total':
|
|
data[i]['keypoints'] = np.vstack([data[i][key] for key in ['keypoints', 'handl2d', 'handr2d', 'face2d']])
|
|
elif mode == 'bodyhand':
|
|
data[i]['keypoints'] = np.vstack([data[i][key] for key in ['keypoints', 'handl2d', 'handr2d']])
|
|
elif mode == 'bodyhandface':
|
|
data[i]['keypoints'] = np.vstack([data[i][key] for key in ['keypoints', 'handl2d', 'handr2d', 'face2d']])
|
|
conf = data[i]['keypoints'][..., -1]
|
|
conf[conf<0] = 0
|
|
data.sort(key=lambda x:x['id'])
|
|
return data
|
|
|
|
def array2raw(array, separator=' ', fmt='%.3f'):
|
|
assert len(array.shape) == 2, 'Only support MxN matrix, {}'.format(array.shape)
|
|
res = []
|
|
for data in array:
|
|
res.append(separator.join([fmt%(d) for d in data]))
|
|
|
|
|
|
def myarray2string(array, separator=', ', fmt='%.3f'):
|
|
assert len(array.shape) == 2, 'Only support MxN matrix, {}'.format(array.shape)
|
|
res = ['[']
|
|
for i in range(array.shape[0]):
|
|
res.append(' [{}]'.format(separator.join([fmt%(d) for d in array[i]])))
|
|
if i != array.shape[0] -1:
|
|
res[-1] += ', '
|
|
res.append(' ]')
|
|
return '\r\n'.join(res)
|
|
|
|
def write_common_results(dumpname=None, results=[], keys=[], fmt='%2.3f'):
|
|
format_out = {'float_kind':lambda x: fmt % x}
|
|
out_text = []
|
|
out_text.append('[\n')
|
|
for idata, data in enumerate(results):
|
|
out_text.append(' {\n')
|
|
output = {}
|
|
output['id'] = data['id']
|
|
for key in keys:
|
|
if key not in data.keys():continue
|
|
# BUG: This function will failed if the rows of the data[key] is too large
|
|
# output[key] = np.array2string(data[key], max_line_width=1000, separator=', ', formatter=format_out)
|
|
output[key] = myarray2string(data[key], separator=', ', fmt=fmt)
|
|
for key in output.keys():
|
|
out_text.append(' \"{}\": {}'.format(key, output[key]))
|
|
if key != keys[-1]:
|
|
out_text.append(',\n')
|
|
else:
|
|
out_text.append('\n')
|
|
out_text.append(' }')
|
|
if idata != len(results) - 1:
|
|
out_text.append(',\n')
|
|
else:
|
|
out_text.append('\n')
|
|
out_text.append(']\n')
|
|
if dumpname is not None:
|
|
mkout(dumpname)
|
|
with open(dumpname, 'w') as f:
|
|
f.writelines(out_text)
|
|
else:
|
|
return ''.join(out_text)
|
|
|
|
def write_keypoints3d(dumpname, results):
|
|
# TODO:rewrite it
|
|
keys = ['keypoints3d']
|
|
write_common_results(dumpname, results, keys, fmt='%6.3f')
|
|
|
|
def write_vertices(dumpname, results):
|
|
keys = ['vertices']
|
|
write_common_results(dumpname, results, keys, fmt='%6.3f')
|
|
|
|
def write_smpl(dumpname, results):
|
|
keys = ['Rh', 'Th', 'poses', 'expression', 'shapes']
|
|
write_common_results(dumpname, results, keys)
|
|
|
|
|
|
def get_bbox_from_pose(pose_2d, img=None, rate = 0.1):
|
|
# this function returns bounding box from the 2D pose
|
|
# here use pose_2d[:, -1] instead of pose_2d[:, 2]
|
|
# because when vis reprojection, the result will be (x, y, depth, conf)
|
|
validIdx = pose_2d[:, -1] > 0
|
|
if validIdx.sum() == 0:
|
|
return [0, 0, 100, 100, 0]
|
|
y_min = int(min(pose_2d[validIdx, 1]))
|
|
y_max = int(max(pose_2d[validIdx, 1]))
|
|
x_min = int(min(pose_2d[validIdx, 0]))
|
|
x_max = int(max(pose_2d[validIdx, 0]))
|
|
dx = (x_max - x_min)*rate
|
|
dy = (y_max - y_min)*rate
|
|
# 后面加上类别这些
|
|
bbox = [x_min-dx, y_min-dy, x_max+dx, y_max+dy, 1]
|
|
if img is not None:
|
|
correct_bbox(img, bbox)
|
|
return bbox
|
|
|
|
def correct_bbox(img, bbox):
|
|
# this function corrects the bbox, which is out of image
|
|
w = img.shape[0]
|
|
h = img.shape[1]
|
|
if bbox[2] <= 0 or bbox[0] >= h or bbox[1] >= w or bbox[3] <= 0:
|
|
bbox[4] = 0
|
|
return bbox
|
|
|
|
def merge_params(param_list, share_shape=True):
|
|
output = {}
|
|
for key in ['poses', 'shapes', 'Rh', 'Th', 'expression']:
|
|
if key in param_list[0].keys():
|
|
output[key] = np.vstack([v[key] for v in param_list])
|
|
if share_shape:
|
|
output['shapes'] = output['shapes'].mean(axis=0, keepdims=True)
|
|
return output
|
|
|
|
def select_nf(params_all, nf):
|
|
output = {}
|
|
for key in ['poses', 'Rh', 'Th']:
|
|
output[key] = params_all[key][nf:nf+1, :]
|
|
if 'expression' in params_all.keys():
|
|
output['expression'] = params_all['expression'][nf:nf+1, :]
|
|
if params_all['shapes'].shape[0] == 1:
|
|
output['shapes'] = params_all['shapes']
|
|
else:
|
|
output['shapes'] = params_all['shapes'][nf:nf+1, :]
|
|
return output |