332 lines
13 KiB
Python
332 lines
13 KiB
Python
from easymocap.mytools.camera_utils import read_cameras
|
||
from easymocap.mytools.debug_utils import log, myerror, mywarn
|
||
from easymocap.mytools.file_utils import read_json
|
||
from .basedata import ImageDataBase, read_mv_images, find_best_people, find_all_people
|
||
import os
|
||
from os.path import join
|
||
import numpy as np
|
||
import cv2
|
||
from collections import defaultdict
|
||
|
||
panoptic15_in_body15 = [1,0,8,5,6,7,12,13,14,2,3,4,9,10,11]
|
||
|
||
def convert_body15_panoptic15(keypoints):
|
||
k3d_panoptic15 = keypoints[..., panoptic15_in_body15,: ]
|
||
return k3d_panoptic15
|
||
|
||
def convert_panoptic15_body15(keypoints):
|
||
keypoints_b15 = np.zeros_like(keypoints)
|
||
keypoints_b15[..., panoptic15_in_body15, :] = keypoints
|
||
return keypoints_b15
|
||
|
||
def padding_and_stack(datas):
|
||
shapes = {}
|
||
for data in datas:
|
||
if len(data) == 0:
|
||
continue
|
||
for key, value in data.items():
|
||
if key not in shapes.keys():
|
||
shapes[key] = value.shape
|
||
collect = {key: np.zeros((len(datas), *shapes[key])) for key in shapes.keys()}
|
||
for i, data in enumerate(datas):
|
||
for key, value in data.items():
|
||
collect[key][i] = value
|
||
return collect
|
||
|
||
def padding_empty(datas):
|
||
shapes = {}
|
||
for data in datas:
|
||
if len(data) == 0:
|
||
continue
|
||
for key, value in data.items():
|
||
if key not in shapes.keys():
|
||
shapes[key] = value.shape[1:]
|
||
collect = {key: [None for data in datas] for key in shapes.keys()}
|
||
for i, data in enumerate(datas):
|
||
for key, shape in shapes.items():
|
||
if key not in data.keys():
|
||
print('[Dataset] padding empty view {} of {}'.format(i, key))
|
||
collect[key][i] = np.zeros((0, *shape), dtype=np.float32)
|
||
else:
|
||
collect[key][i] = data[key]
|
||
return collect
|
||
|
||
def parse_frames(pafs_frame, H, W):
|
||
# 解析单帧的
|
||
res = {
|
||
'joints': [],
|
||
'pafs': {}
|
||
}
|
||
joints = pafs_frame[1:1+3*25]
|
||
for i in range(25):
|
||
value = np.fromstring(joints[3*i+2], sep=' ').reshape(3, -1).T
|
||
value[:, 0] = value[:, 0] * W
|
||
value[:, 1] = value[:, 1] * H
|
||
res['joints'].append(value.astype(np.float32))
|
||
# parse pafs
|
||
pafs = pafs_frame[1+3*25+1:]
|
||
for npart in range(26):
|
||
label = pafs[3*npart+0].split(' ')[2:]
|
||
label = (int(label[0]), int(label[1]))
|
||
shape = pafs[3*npart+1].split(' ')[2:]
|
||
w, h = int(shape[0]), int(shape[1])
|
||
value = np.fromstring(pafs[3*npart+2], sep=' ').reshape(w, h).astype(np.float32)
|
||
res['pafs'][label] = value
|
||
return res
|
||
|
||
def read_4dassociation(pafs, H, W):
|
||
outputs = []
|
||
# 解析paf文件
|
||
with open(pafs, 'r') as f:
|
||
pafs = f.readlines()
|
||
indices = []
|
||
for i, line in enumerate(pafs):
|
||
if line.startswith('# newframes:'):
|
||
indices.append([i])
|
||
elif line.startswith('# end frames:'):
|
||
indices[-1].append(i)
|
||
print('[Read OpenPose] Totally {} frames'.format(len(indices)))
|
||
for (start, end) in indices:
|
||
pafs_frame = pafs[start+1:end]
|
||
pafs_frame = list(map(lambda x:x.strip(), pafs_frame))
|
||
frames = parse_frames(pafs_frame, H, W)
|
||
outputs.append(frames)
|
||
return outputs
|
||
|
||
class MVDataset(ImageDataBase):
|
||
def __init__(self, root, subs, subs_vis, ranges, read_image=False, reader={}, filter={}) -> None:
|
||
super().__init__(root, subs, ranges, read_image)
|
||
self.subs_vis = subs_vis
|
||
self.length = 0
|
||
for key, value in reader.items():
|
||
if key == 'images':
|
||
self.try_to_extract_images(root, value)
|
||
data, meta = read_mv_images(root, value['root'], value['ext'], subs)
|
||
self.length = len(data)
|
||
elif key == 'image_shape':
|
||
imgnames = self.infos['images'][0]
|
||
shapes = []
|
||
for imgname in imgnames:
|
||
img = cv2.imread(imgname)
|
||
height, width, _ = img.shape
|
||
log('[{}] sub {} shape {}'.format(self.__class__.__name__, imgname, img.shape))
|
||
shapes.append([height, width])
|
||
data = [shapes]
|
||
meta = {}
|
||
elif key == 'annots':
|
||
data, meta = read_mv_images(root, value['root'], value['ext'], subs)
|
||
if self.length > 0:
|
||
if self.length != len(data):
|
||
myerror('annots length {} not equal to images length {}.'.format(len(data), self.length))
|
||
data = data[:self.length]
|
||
else:
|
||
self.length = len(data)
|
||
elif key == 'openpose':
|
||
# 读取open pose
|
||
if len(subs) == 0:
|
||
pafs = sorted(os.listdir(join(root, value['root'])))
|
||
else:
|
||
pafs = [f'{sub}.txt' for sub in subs]
|
||
results = []
|
||
for nv, paf in enumerate(pafs):
|
||
pafname = join(root, value['root'], paf)
|
||
infos = read_4dassociation(pafname, H=self.infos['image_shape'][0][nv][0], W=self.infos['image_shape'][0][nv][1])
|
||
results.append(infos)
|
||
data = [[d[i] for d in results] for i in range(self.length)]
|
||
meta = {}
|
||
elif key == 'cameras':
|
||
if 'with_sub' in value.keys():
|
||
raise NotImplementedError
|
||
else:
|
||
cameras = read_cameras(os.path.join(root, value['root']))
|
||
if 'remove_k3' in value.keys():
|
||
for cam, camera in cameras.items():
|
||
camera['dist'][:, 4] = 0.
|
||
data = [cameras]
|
||
meta = {}
|
||
elif key in ['pelvis']:
|
||
continue
|
||
elif key == 'keypoints3d':
|
||
k3droot = value['root']
|
||
filenames = sorted(os.listdir(k3droot))[:self.length]
|
||
res_key = value.get('key', 'pred')
|
||
data = []
|
||
for filename in filenames:
|
||
results = read_json(join(k3droot, filename))
|
||
if 'pids' not in results.keys():
|
||
# 擅自补全
|
||
results['pids'] = list(range(len(results[res_key])))
|
||
data.append({
|
||
'pids': results['pids'],
|
||
'keypoints3d': np.array(results[res_key], dtype=np.float32)
|
||
})
|
||
if data[-1]['keypoints3d'].shape[-1] == 3:
|
||
mywarn('The input keypoints dont have confidence')
|
||
data[-1]['keypoints3d'] = np.concatenate([data[-1]['keypoints3d'], np.ones_like(data[-1]['keypoints3d'][..., :1])], axis=-1)
|
||
if 'conversion' in value.keys():
|
||
if value['conversion'] == 'panoptic15_to_body15':
|
||
data[-1]['keypoints3d'] = convert_panoptic15_body15(data[-1]['keypoints3d'])
|
||
else:
|
||
raise ValueError(f'Unknown reader: {key}')
|
||
self.infos[key] = data
|
||
self.meta.update(meta)
|
||
self.reader = reader
|
||
self.filter = filter
|
||
if len(self.subs) == 0:
|
||
self.subs = self.meta['subs']
|
||
self.check_frames_length()
|
||
|
||
@staticmethod
|
||
def read_annots(annotnames):
|
||
val = []
|
||
for annname in annotnames:
|
||
annots = read_json(annname)['annots']
|
||
# select the best people
|
||
annots = find_best_people(annots)
|
||
val.append(annots)
|
||
val = padding_and_stack(val)
|
||
return val
|
||
|
||
def filter_openpose(self, candidates, pafs):
|
||
for nv, candview in enumerate(candidates):
|
||
H=self.infos['image_shape'][0][nv][0]
|
||
W=self.infos['image_shape'][0][nv][1]
|
||
for cand in candview:
|
||
if 'border' in self.filter.keys():
|
||
border = self.filter['border'] * max(H, W)
|
||
flag = (cand[:, 0] > border) & (cand[:, 0] < W - border) & (cand[:, 1] > border) & (cand[:, 1] < H - border)
|
||
cand[~flag] = 0
|
||
return candidates, pafs
|
||
|
||
def __getitem__(self, index):
|
||
frame = self.frames[index]
|
||
ret = {}
|
||
for key, value in self.infos.items():
|
||
if len(value) == 1:
|
||
ret[key] = value[0]
|
||
elif frame >= len(value):
|
||
myerror(f'[{self.__class__.__name__}] {key}: index {frame} out of range {len(value)}')
|
||
else:
|
||
ret[key] = value[frame]
|
||
ret_list = defaultdict(list)
|
||
for key, val in ret.items():
|
||
if key == 'annots':
|
||
ret_list[key] = self.read_annots(val)
|
||
elif key == 'cameras':
|
||
for sub in self.subs:
|
||
select = {k: val[sub][k] for k in ['K', 'R', 'T', 'dist', 'P']}
|
||
ret_list[key].append(select)
|
||
ret_list[key] = padding_and_stack(ret_list[key])
|
||
elif key == 'images':
|
||
if self.flag_read_image:
|
||
for i, sub in enumerate(self.subs):
|
||
imgname = val[i]
|
||
if sub in self.subs_vis or self.subs_vis == 'all':
|
||
img = self.read_image(imgname)
|
||
else:
|
||
img = imgname
|
||
ret_list[key].append(img)
|
||
ret_list['imgnames'].append(imgname)
|
||
else:
|
||
ret_list[key] = val
|
||
ret_list['imgnames'] = val
|
||
elif key == 'openpose':
|
||
ret_list[key] = [v['joints'] for v in val]
|
||
# 同时返回PAF
|
||
ret_list[key+'_paf'] = [v['pafs'] for v in val]
|
||
# check一下PAF:
|
||
for nv in range(len(ret_list[key])):
|
||
ret_list[key+'_paf'][nv][(8, 1)] = ret_list[key+'_paf'][nv].pop((1, 8)).T
|
||
ret_list[key], ret_list[key+'_paf'] = self.filter_openpose(ret_list[key], ret_list[key+'_paf'])
|
||
elif key == 'keypoints3d':
|
||
ret_list['keypoints3d'] = val['keypoints3d']
|
||
if 'pids' in val.keys():
|
||
ret_list['pids'] = val['pids']
|
||
else:
|
||
ret_list['pids'] = list(range(len(val['keypoints3d'])))
|
||
elif key in ['image_shape']:
|
||
pass
|
||
else:
|
||
print('[Dataset] Unknown key: {}'.format(key))
|
||
ret_list.update(ret_list.pop('annots', {}))
|
||
for key, val in self.reader.items():
|
||
if key == 'pelvis' and 'annots' in self.reader.keys(): # load pelvis from annots.keypoints
|
||
ret_list[key] = [d[:, val.root_id] for d in ret_list['keypoints']]
|
||
elif key == 'pelvis' and 'openpose' in self.reader.keys():
|
||
ret_list[key] = [d[val.root_id] for d in ret_list['openpose']]
|
||
ret_list['meta'] = {
|
||
'subs': self.subs,
|
||
'index': index,
|
||
'frame': frame,
|
||
'image_shape': ret['image_shape'],
|
||
'imgnames': ret_list['imgnames'],
|
||
}
|
||
return ret_list
|
||
|
||
def check(self, index):
|
||
raise NotImplementedError
|
||
|
||
class MVMP(MVDataset):
|
||
def read_annots(self, annotnames):
|
||
val = []
|
||
for annname in annotnames:
|
||
annots = read_json(annname)['annots']
|
||
# 在这里进行filter，去掉不需要的2D
|
||
annots_valid = []
|
||
for annot in annots:
|
||
flag = True
|
||
if 'bbox_size' in self.filter.keys():
|
||
bbox_size = self.filter['bbox_size']
|
||
bbox = annot['bbox']
|
||
area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
|
||
if area < bbox_size:
|
||
flag = False
|
||
if flag:
|
||
annots_valid.append(annot)
|
||
annots = annots_valid
|
||
# select the best people
|
||
annots = find_all_people(annots)
|
||
val.append(annots)
|
||
val = padding_empty(val)
|
||
return val
|
||
|
||
def check(self, index):
|
||
data = self.__getitem__(index)
|
||
from easymocap.mytools.vis_base import plot_bbox, merge, plot_keypoints_auto
|
||
# check the subs vis
|
||
vis = []
|
||
for nv, sub in enumerate(self.subs):
|
||
if sub not in self.subs_vis:continue
|
||
img = data['images'][nv].copy()
|
||
bbox = data['bbox'][nv]
|
||
kpts = data['keypoints'][nv]
|
||
for i in range(bbox.shape[0]):
|
||
plot_bbox(img, bbox[i], pid=i)
|
||
plot_keypoints_auto(img, kpts[i], pid=i, use_limb_color=False)
|
||
vis.append(img)
|
||
vis = merge(vis)
|
||
cv2.imwrite('debug/{}_{:06d}.jpg'.format(self.__class__.__name__, index), vis)
|
||
|
||
if __name__ == '__main__':
|
||
config = '''
|
||
args:
|
||
root: /nas/ZJUMoCap/Part0/313
|
||
subs: []
|
||
subs_vis: ['01', '07', '13', '19']
|
||
ranges: [0, 100, 1]
|
||
read_image: False
|
||
reader:
|
||
images:
|
||
root: images
|
||
ext: .jpg
|
||
annots:
|
||
root: annots
|
||
ext: .json
|
||
cameras: # 兼容所有帧的相机参数不同的情况
|
||
root: ''
|
||
'''
|
||
import yaml
|
||
config = yaml.load(config, Loader=yaml.FullLoader)
|
||
dataset = MVDataset(**config['args'])
|
||
for i in range(len(dataset)):
|
||
data = dataset[i] |