637 lines
26 KiB
Python
637 lines
26 KiB
Python
'''
|
||
@ Date: 2021-01-13 16:53:55
|
||
@ Author: Qing Shuai
|
||
@ LastEditors: Qing Shuai
|
||
@ LastEditTime: 2021-06-25 15:53:12
|
||
@ FilePath: /EasyMocapRelease/easymocap/dataset/base.py
|
||
'''
|
||
import os
|
||
from os.path import join
|
||
from glob import glob
|
||
import cv2
|
||
import os, sys
|
||
import numpy as np
|
||
|
||
from ..mytools.camera_utils import read_camera, get_fundamental_matrix, Undistort
|
||
from ..mytools import FileWriter, read_annot, getFileList, save_json
|
||
from ..mytools.reader import read_keypoints3d, read_json, read_smpl
|
||
from ..mytools.file_utils import merge_params, select_nf, save_annot
|
||
|
||
def crop_image(img, annot, vis_2d=False, config={}, crop_square=True):
|
||
for det in annot:
|
||
bbox = det['bbox']
|
||
l, t, r, b = det['bbox'][:4]
|
||
if crop_square:
|
||
if b - t > r - l:
|
||
diff = (b - t) - (r - l)
|
||
l -= diff//2
|
||
r += diff//2
|
||
else:
|
||
diff = (r - l) - (b - t)
|
||
t -= diff//2
|
||
b += diff//2
|
||
l = max(0, int(l+0.5))
|
||
t = max(0, int(t+0.5))
|
||
r = min(img.shape[1], int(r+0.5))
|
||
b = min(img.shape[0], int(b+0.5))
|
||
det['bbox'][:4] = [l, t, r, b]
|
||
if vis_2d:
|
||
crop_img = img.copy()
|
||
from easymocap.mytools import plot_keypoints
|
||
plot_keypoints(crop_img, det['keypoints'], pid=det['id'],
|
||
config=config, use_limb_color=True, lw=2)
|
||
else:
|
||
crop_img = img
|
||
crop_img = crop_img[t:b, l:r, :]
|
||
if crop_square:
|
||
crop_img = cv2.resize(crop_img, (256, 256))
|
||
else:
|
||
crop_img = cv2.resize(crop_img, (128, 256))
|
||
det['crop'] = crop_img
|
||
det['img'] = img
|
||
return 0
|
||
|
||
class ImageFolder:
|
||
"""Dataset for image folders"""
|
||
def __init__(self, root, subs=[], out=None, image_root='images', annot_root='annots',
|
||
kpts_type='body15', config={}, no_img=False) -> None:
|
||
self.root = root
|
||
self.image_root = join(root, image_root)
|
||
self.annot_root = join(root, annot_root)
|
||
self.kpts_type = kpts_type
|
||
self.no_img = no_img
|
||
if len(subs) == 0:
|
||
self.imagelist = getFileList(self.image_root, '.jpg')
|
||
self.annotlist = getFileList(self.annot_root, '.json')
|
||
else:
|
||
self.imagelist, self.annotlist = [], []
|
||
for sub in subs:
|
||
images = sorted([join(sub, i) for i in os.listdir(join(self.image_root, sub))])
|
||
annots = sorted([join(sub, i) for i in os.listdir(join(self.annot_root, sub))])
|
||
if len(annots) < len(images):
|
||
print('[WARN] length of annots != lenght of images')
|
||
images = images[:len(annots)]
|
||
self.imagelist.extend(images)
|
||
self.annotlist.extend(annots)
|
||
self.out = out
|
||
self.writer = FileWriter(self.out, config=config)
|
||
self.gtK, self.gtRT = False, False
|
||
|
||
def load_gt_cameras(self):
|
||
cameras = load_cameras(self.root)
|
||
gtCameras = []
|
||
for i, name in enumerate(self.annotlist):
|
||
cam = os.path.dirname(name)
|
||
gtcams = {key:cameras[cam][key].copy() for key in ['K', 'R', 'T', 'dist']}
|
||
gtCameras.append(gtcams)
|
||
self.gtCameras = gtCameras
|
||
|
||
def __len__(self) -> int:
|
||
return len(self.imagelist)
|
||
|
||
def __getitem__(self, index: int):
|
||
imgname = join(self.image_root, self.imagelist[index])
|
||
annname = join(self.annot_root, self.annotlist[index])
|
||
assert os.path.exists(imgname) and os.path.exists(annname), (imgname, annname)
|
||
assert os.path.basename(imgname).split('.')[0] == os.path.basename(annname).split('.')[0], '{}, {}'.format(imgname, annname)
|
||
if not self.no_img:
|
||
img = cv2.imread(imgname)
|
||
else:
|
||
img = None
|
||
annot = read_annot(annname, self.kpts_type)
|
||
return img, annot
|
||
|
||
def camera(self, index=0, annname=None):
|
||
if annname is None:
|
||
annname = join(self.annot_root, self.annotlist[index])
|
||
data = read_json(annname)
|
||
if 'K' not in data.keys():
|
||
height, width = data['height'], data['width']
|
||
# focal = 1.2*max(height, width) # as colmap
|
||
focal = 1.2*min(height, width) # as colmap
|
||
K = np.array([focal, 0., width/2, 0., focal, height/2, 0. ,0., 1.]).reshape(3, 3)
|
||
else:
|
||
K = np.array(data['K']).reshape(3, 3)
|
||
camera = {'K':K ,'R': np.eye(3), 'T': np.zeros((3, 1)), 'dist': np.zeros((1, 5))}
|
||
if self.gtK:
|
||
camera['K'] = self.gtCameras[index]['K']
|
||
if self.gtRT:
|
||
camera['R'] = self.gtCameras[index]['R']
|
||
camera['T'] = self.gtCameras[index]['T']
|
||
# camera['T'][2, 0] = 5. # guess to 5 meters
|
||
camera['RT'] = np.hstack((camera['R'], camera['T']))
|
||
camera['P'] = camera['K'] @ np.hstack((camera['R'], camera['T']))
|
||
return camera
|
||
|
||
def basename(self, nf):
|
||
return self.annotlist[nf].replace('.json', '')
|
||
|
||
def write_keypoints3d(self, results, nf):
|
||
outname = join(self.out, 'keypoints3d', '{}.json'.format(self.basename(nf)))
|
||
self.writer.write_keypoints3d(results, outname)
|
||
|
||
def write_vertices(self, results, nf):
|
||
outname = join(self.out, 'vertices', '{}.json'.format(self.basename(nf)))
|
||
self.writer.write_vertices(results, outname)
|
||
|
||
def write_smpl(self, results, nf):
|
||
outname = join(self.out, 'smpl', '{}.json'.format(self.basename(nf)))
|
||
self.writer.write_smpl(results, outname)
|
||
|
||
def vis_smpl(self, render_data, image, camera, nf):
|
||
outname = join(self.out, 'smpl', '{}.jpg'.format(self.basename(nf)))
|
||
images = [image]
|
||
for key in camera.keys():
|
||
camera[key] = camera[key][None, :, :]
|
||
self.writer.vis_smpl(render_data, images, camera, outname, add_back=True)
|
||
|
||
# class VideoFolder(ImageFolder):
|
||
# "一段视频的图片的文件夹"
|
||
# def __init__(self, root, name, out=None,
|
||
# image_root='images', annot_root='annots',
|
||
# kpts_type='body15', config={}, no_img=False) -> None:
|
||
# self.root = root
|
||
# self.image_root = join(root, image_root, name)
|
||
# self.annot_root = join(root, annot_root, name)
|
||
# self.name = name
|
||
# self.kpts_type = kpts_type
|
||
# self.no_img = no_img
|
||
# self.imagelist = sorted(os.listdir(self.image_root))
|
||
# self.annotlist = sorted(os.listdir(self.annot_root))
|
||
# self.ret_crop = False
|
||
# self.gtK, self.gtRT = False, False
|
||
|
||
def load_annot_all(self, path):
|
||
# 这个不使用personID，只是单纯的罗列一下
|
||
assert os.path.exists(path), '{} not exists!'.format(path)
|
||
results = []
|
||
annnames = sorted(glob(join(path, '*.json')))
|
||
for annname in annnames:
|
||
datas = read_annot(annname, self.kpts_type)
|
||
if self.ret_crop:
|
||
# TODO:修改imgname
|
||
basename = os.path.basename(annname)
|
||
imgname = annname\
|
||
.replace('annots-cpn', 'images')\
|
||
.replace('annots', 'images')\
|
||
.replace('.json', '.jpg')
|
||
assert os.path.exists(imgname), imgname
|
||
img = cv2.imread(imgname)
|
||
crop_image(img, datas)
|
||
results.append(datas)
|
||
return results
|
||
|
||
def load_annot(self, path, pids=[]):
|
||
# 这个根据人的ID预先存一下
|
||
assert os.path.exists(path), '{} not exists!'.format(path)
|
||
results = {}
|
||
annnames = sorted(glob(join(path, '*.json')))
|
||
for annname in annnames:
|
||
nf = int(os.path.basename(annname).replace('.json', ''))
|
||
datas = read_annot(annname, self.kpts_type)
|
||
for data in datas:
|
||
pid = data['id']
|
||
if len(pids) > 0 and pid not in pids:
|
||
continue
|
||
# 注意 这里没有考虑从哪开始的
|
||
if pid not in results.keys():
|
||
results[pid] = {'bboxes': [], 'keypoints2d': []}
|
||
results[pid]['bboxes'].append(data['bbox'])
|
||
results[pid]['keypoints2d'].append(data['keypoints'])
|
||
for pid, val in results.items():
|
||
for key in val.keys():
|
||
val[key] = np.stack(val[key])
|
||
return results
|
||
|
||
def load_smpl(self, path, pids=[]):
|
||
""" load SMPL parameters from files
|
||
|
||
Args:
|
||
path (str): root path of smpl
|
||
pids (list, optional): used person ids. Defaults to [], loading all person.
|
||
"""
|
||
assert os.path.exists(path), '{} not exists!'.format(path)
|
||
results = {}
|
||
smplnames = sorted(glob(join(path, '*.json')))
|
||
for smplname in smplnames:
|
||
nf = int(os.path.basename(smplname).replace('.json', ''))
|
||
datas = read_smpl(smplname)
|
||
for data in datas:
|
||
pid = data['id']
|
||
if len(pids) > 0 and pid not in pids:
|
||
continue
|
||
# 注意 这里没有考虑从哪开始的
|
||
if pid not in results.keys():
|
||
results[pid] = {'body_params': [], 'frames': []}
|
||
results[pid]['body_params'].append(data)
|
||
results[pid]['frames'].append(nf)
|
||
for pid, val in results.items():
|
||
val['body_params'] = merge_params(val['body_params'])
|
||
return results
|
||
|
||
class _VideoBase:
|
||
"""Dataset for single sequence data
|
||
"""
|
||
def __init__(self, image_root, annot_root, out=None, config={}, kpts_type='body15', no_img=False) -> None:
|
||
self.image_root = image_root
|
||
self.annot_root = annot_root
|
||
self.kpts_type = kpts_type
|
||
self.no_img = no_img
|
||
self.config = config
|
||
assert out is not None
|
||
self.out = out
|
||
self.writer = FileWriter(self.out, config=config)
|
||
imgnames = sorted(os.listdir(self.image_root))
|
||
self.imagelist = imgnames
|
||
self.annotlist = sorted(os.listdir(self.annot_root))
|
||
self.nFrames = len(self.imagelist)
|
||
self.undis = False
|
||
self.read_camera()
|
||
|
||
def read_camera(self):
|
||
# 读入相机参数
|
||
annname = join(self.annot_root, self.annotlist[0])
|
||
data = read_json(annname)
|
||
if 'K' not in data.keys():
|
||
height, width = data['height'], data['width']
|
||
focal = 1.2*max(height, width)
|
||
K = np.array([focal, 0., width/2, 0., focal, height/2, 0. ,0., 1.]).reshape(3, 3)
|
||
else:
|
||
K = np.array(data['K']).reshape(3, 3)
|
||
self.camera = {'K':K ,'R': np.eye(3), 'T': np.zeros((3, 1))}
|
||
|
||
def __getitem__(self, index: int):
|
||
imgname = join(self.image_root, self.imagelist[index])
|
||
annname = join(self.annot_root, self.annotlist[index])
|
||
assert os.path.exists(imgname) and os.path.exists(annname)
|
||
assert os.path.basename(imgname).split('.')[0] == os.path.basename(annname).split('.')[0], '{}, {}'.format(imgname, annname)
|
||
if not self.no_img:
|
||
img = cv2.imread(imgname)
|
||
else:
|
||
img = None
|
||
annot = read_annot(annname, self.kpts_type)
|
||
return img, annot
|
||
|
||
def __len__(self) -> int:
|
||
return self.nFrames
|
||
|
||
def write_smpl(self, peopleDict, nf):
|
||
results = []
|
||
for pid, people in peopleDict.items():
|
||
result = {'id': pid}
|
||
result.update(people.body_params)
|
||
results.append(result)
|
||
self.writer.write_smpl(results, nf)
|
||
|
||
def vis_detections(self, image, detections, nf, to_img=True):
|
||
return self.writer.vis_detections([image], [detections], nf,
|
||
key='keypoints', to_img=to_img, vis_id=True)
|
||
|
||
def vis_repro(self, peopleDict, image, annots, nf):
|
||
# 可视化重投影的关键点与输入的关键点
|
||
detections = []
|
||
for pid, data in peopleDict.items():
|
||
keypoints3d = (data.keypoints3d @ self.camera['R'].T + self.camera['T'].T) @ self.camera['K'].T
|
||
keypoints3d[:, :2] /= keypoints3d[:, 2:]
|
||
keypoints3d = np.hstack([keypoints3d, data.keypoints3d[:, -1:]])
|
||
det = {
|
||
'id': pid,
|
||
'repro': keypoints3d
|
||
}
|
||
detections.append(det)
|
||
return self.writer.vis_detections([image], [detections], nf, key='repro',
|
||
to_img=True, vis_id=False)
|
||
|
||
def vis_smpl(self, peopleDict, faces, image, nf, sub_vis=[],
|
||
mode='smpl', extra_data=[], add_back=True,
|
||
axis=np.array([1., 0., 0.]), degree=0., fix_center=None):
|
||
# 为了统一接口，旋转视角的在此处实现，只在单视角的数据中使用
|
||
# 通过修改相机参数实现
|
||
# 相机参数的修正可以通过计算点的中心来获得
|
||
# render the smpl to each view
|
||
render_data = {}
|
||
for pid, data in peopleDict.items():
|
||
render_data[pid] = {
|
||
'vertices': data.vertices, 'faces': faces,
|
||
'vid': pid, 'name': 'human_{}_{}'.format(nf, pid)}
|
||
for iid, extra in enumerate(extra_data):
|
||
render_data[10000+iid] = {
|
||
'vertices': extra['vertices'],
|
||
'faces': extra['faces'],
|
||
'colors': extra['colors'],
|
||
'name': extra['name']
|
||
}
|
||
camera = {}
|
||
for key in self.camera.keys():
|
||
camera[key] = self.camera[key][None, :, :]
|
||
# render another view point
|
||
if np.abs(degree) > 1e-3:
|
||
vertices_all = np.vstack([data.vertices for data in peopleDict.values()])
|
||
if fix_center is None:
|
||
center = np.mean(vertices_all, axis=0, keepdims=True)
|
||
new_center = center.copy()
|
||
new_center[:, 0:2] = 0
|
||
else:
|
||
center = fix_center.copy()
|
||
new_center = fix_center.copy()
|
||
new_center[:, 2] *= 1.5
|
||
direc = np.array(axis)
|
||
rot, _ = cv2.Rodrigues(direc*degree/90*np.pi/2)
|
||
# If we rorate the data, it is like:
|
||
# V = Rnew @ (V0 - center) + new_center
|
||
# = Rnew @ V0 - Rnew @ center + new_center
|
||
# combine with the camera
|
||
# VV = Rc(Rnew @ V0 - Rnew @ center + new_center) + Tc
|
||
# = Rc@Rnew @ V0 + Rc @ (new_center - Rnew@center) + Tc
|
||
blank = np.zeros_like(image, dtype=np.uint8) + 255
|
||
images = [image, blank]
|
||
Rnew = camera['R'][0] @ rot
|
||
Tnew = camera['R'][0] @ (new_center.T - rot @ center.T) + camera['T'][0]
|
||
camera['K'] = np.vstack([camera['K'], camera['K']])
|
||
camera['R'] = np.vstack([camera['R'], Rnew[None, :, :]])
|
||
camera['T'] = np.vstack([camera['T'], Tnew[None, :, :]])
|
||
else:
|
||
images = [image]
|
||
self.writer.vis_smpl(render_data, nf, images, camera, mode, add_back=add_back)
|
||
|
||
def load_cameras(path):
|
||
# 读入相机参数
|
||
intri_name = join(path, 'intri.yml')
|
||
extri_name = join(path, 'extri.yml')
|
||
if os.path.exists(intri_name) and os.path.exists(extri_name):
|
||
cameras = read_camera(intri_name, extri_name)
|
||
cams = cameras.pop('basenames')
|
||
else:
|
||
print('\n\n!!!there is no camera parameters, maybe bug: \n', intri_name, extri_name, '\n')
|
||
cameras = None
|
||
return cameras
|
||
|
||
def numpy_to_list(array, precision=3):
|
||
return np.round(array, precision).tolist()
|
||
|
||
class MVBase:
|
||
""" Dataset for multiview data
|
||
"""
|
||
def __init__(self, root, cams=[], out=None, config={},
|
||
image_root='images', annot_root='annots',
|
||
kpts_type='body15',
|
||
undis=True, no_img=False, filter2d=None) -> None:
|
||
self.root = root
|
||
self.image_root = join(root, image_root)
|
||
self.annot_root = join(root, annot_root)
|
||
self.kpts_type = kpts_type
|
||
self.undis = undis
|
||
self.no_img = no_img
|
||
# use when debug
|
||
self.ret_crop = False
|
||
self.config = config
|
||
# results path
|
||
# the results store keypoints3d
|
||
self.skel_path = None
|
||
self.out = out
|
||
self.writer = FileWriter(self.out, config=config)
|
||
|
||
self.cams = cams
|
||
self.imagelist = {}
|
||
self.annotlist = {}
|
||
for cam in cams: #TODO: 增加start,end
|
||
# ATTN: when image name's frame number is not continuous,
|
||
imgnames = sorted(os.listdir(join(self.image_root, cam)))
|
||
self.imagelist[cam] = imgnames
|
||
if os.path.exists(self.annot_root):
|
||
self.annotlist[cam] = sorted(os.listdir(join(self.annot_root, cam)))
|
||
self.has2d = True
|
||
else:
|
||
self.has2d = False
|
||
nFrames = min([len(val) for key, val in self.imagelist.items()])
|
||
self.nFrames = nFrames
|
||
self.nViews = len(cams)
|
||
self.read_camera(self.root)
|
||
self.filter2d = filter2d
|
||
if filter2d is not None:
|
||
from .filter import make_filter
|
||
self.filter2d = make_filter(filter2d)
|
||
|
||
def read_camera(self, path):
|
||
# 读入相机参数
|
||
intri_name = join(path, 'intri.yml')
|
||
extri_name = join(path, 'extri.yml')
|
||
if os.path.exists(intri_name) and os.path.exists(extri_name):
|
||
self.cameras = read_camera(intri_name, extri_name)
|
||
self.cameras.pop('basenames')
|
||
# 注意：这里的相机参数一定要用定义的，不然只用一部分相机的时候会出错
|
||
cams = self.cams
|
||
self.cameras_for_affinity = [[cam['invK'], cam['R'], cam['T']] for cam in [self.cameras[name] for name in cams]]
|
||
self.Pall = np.stack([self.cameras[cam]['P'] for cam in cams])
|
||
self.Fall = get_fundamental_matrix(self.cameras, cams)
|
||
else:
|
||
print('\n!!!\n!!!there is no camera parameters, maybe bug: \n', intri_name, extri_name, '\n')
|
||
self.cameras = None
|
||
|
||
def undistort(self, images):
|
||
if self.cameras is not None and len(images) > 0:
|
||
images_ = []
|
||
for nv in range(self.nViews):
|
||
mtx = self.cameras[self.cams[nv]]['K']
|
||
dist = self.cameras[self.cams[nv]]['dist']
|
||
if images[nv] is not None:
|
||
frame = cv2.undistort(images[nv], mtx, dist, None)
|
||
else:
|
||
frame = None
|
||
images_.append(frame)
|
||
else:
|
||
images_ = images
|
||
return images_
|
||
|
||
def undis_det(self, lDetections):
|
||
for nv in range(len(lDetections)):
|
||
camera = self.cameras[self.cams[nv]]
|
||
for det in lDetections[nv]:
|
||
det['bbox'] = Undistort.bbox(det['bbox'], K=camera['K'], dist=camera['dist'])
|
||
keypoints = det['keypoints']
|
||
det['keypoints'] = Undistort.points(keypoints=keypoints, K=camera['K'], dist=camera['dist'])
|
||
return lDetections
|
||
|
||
def select_person(self, annots_all, index, pid):
|
||
annots = {'bbox': [], 'keypoints': []}
|
||
for nv, cam in enumerate(self.cams):
|
||
data = [d for d in annots_all[nv] if d['id'] == pid]
|
||
if len(data) == 1:
|
||
data = data[0]
|
||
bbox = data['bbox']
|
||
keypoints = data['keypoints']
|
||
else:
|
||
if self.verbose:print('not found pid {} in frame {}, view {}'.format(self.pid, index, nv))
|
||
keypoints = np.zeros((self.config['nJoints'], 3))
|
||
bbox = np.array([0, 0, 100., 100., 0.])
|
||
annots['bbox'].append(bbox)
|
||
annots['keypoints'].append(keypoints)
|
||
for key in ['bbox', 'keypoints']:
|
||
annots[key] = np.stack(annots[key])
|
||
return annots
|
||
|
||
def __getitem__(self, index: int):
|
||
images, annots = [], []
|
||
for cam in self.cams:
|
||
imgname = join(self.image_root, cam, self.imagelist[cam][index])
|
||
assert os.path.exists(imgname), imgname
|
||
if self.has2d:
|
||
annname = join(self.annot_root, cam, self.annotlist[cam][index])
|
||
assert os.path.exists(annname), annname
|
||
assert self.imagelist[cam][index].split('.')[0] == self.annotlist[cam][index].split('.')[0]
|
||
annot = read_annot(annname, self.kpts_type)
|
||
else:
|
||
annot = []
|
||
if not self.no_img:
|
||
img = cv2.imread(imgname)
|
||
images.append(img)
|
||
else:
|
||
img = None
|
||
images.append(None)
|
||
if self.filter2d is not None:
|
||
annot_valid = []
|
||
for ann in annot:
|
||
if self.filter2d(**ann):
|
||
annot_valid.append(ann)
|
||
annot = annot_valid
|
||
annot = self.filter2d.nms(annot)
|
||
if self.ret_crop:
|
||
crop_image(img, annot, True, self.config)
|
||
annots.append(annot)
|
||
if self.undis:
|
||
images = self.undistort(images)
|
||
annots = self.undis_det(annots)
|
||
return images, annots
|
||
|
||
def __len__(self) -> int:
|
||
return self.nFrames
|
||
|
||
def vis_detections(self, images, lDetections, nf, mode='detec', to_img=True, sub_vis=[]):
|
||
outname = join(self.out, mode, '{:06d}.jpg'.format(nf))
|
||
if len(sub_vis) != 0:
|
||
valid_idx = [self.cams.index(i) for i in sub_vis]
|
||
images = [images[i] for i in valid_idx]
|
||
lDetections = [lDetections[i] for i in valid_idx]
|
||
return self.writer.vis_keypoints2d_mv(images, lDetections, outname=outname, vis_id=True)
|
||
|
||
def basename(self, nf):
|
||
return '{:06d}'.format(nf)
|
||
|
||
def write_keypoints2d(self, lDetections, nf):
|
||
for nv in range(len(lDetections)):
|
||
cam = self.cams[nv]
|
||
annname = join(self.annot_root, cam, self.annotlist[cam][nf])
|
||
outname = join(self.out, 'keypoints2d', cam, self.annotlist[cam][nf])
|
||
annot_origin = read_json(annname)
|
||
annots = lDetections[nv]
|
||
results = []
|
||
for annot in annots:
|
||
results.append({
|
||
'personID': annot['id'],
|
||
'bbox': numpy_to_list(annot['bbox'], 2),
|
||
'keypoints': numpy_to_list(annot['keypoints'], 2)
|
||
})
|
||
annot_origin['annots'] = results
|
||
save_annot(outname, annot_origin)
|
||
|
||
def write_keypoints3d(self, results, nf):
|
||
outname = join(self.out, 'keypoints3d', self.basename(nf)+'.json')
|
||
self.writer.write_keypoints3d(results, outname)
|
||
|
||
def write_vertices(self, results, nf):
|
||
outname = join(self.out, 'vertices', '{}.json'.format(self.basename(nf)))
|
||
self.writer.write_vertices(results, outname)
|
||
|
||
def write_smpl(self, results, nf, mode='smpl'):
|
||
outname = join(self.out, mode, self.basename(nf)+'.json')
|
||
self.writer.write_smpl(results, outname)
|
||
|
||
def vis_smpl(self, peopleDict, faces, images, nf, sub_vis=[],
|
||
mode='smpl', extra_data=[], extra_mesh=[],
|
||
add_back=True, camera_scale=1, cameras=None):
|
||
# render the smpl to each view
|
||
render_data = {}
|
||
for pid, data in peopleDict.items():
|
||
render_data[pid] = {
|
||
'vertices': data.vertices, 'faces': faces,
|
||
'vid': pid, 'name': 'human_{}_{}'.format(nf, pid)}
|
||
for iid, extra in enumerate(extra_data):
|
||
render_data[10000+iid] = {
|
||
'vertices': extra['vertices'],
|
||
'faces': extra['faces'],
|
||
'name': extra['name']
|
||
}
|
||
if 'colors' in extra.keys():
|
||
render_data[10000+iid]['colors'] = extra['colors']
|
||
elif 'vid' in extra.keys():
|
||
render_data[10000+iid]['vid'] = extra['vid']
|
||
|
||
if len(sub_vis) == 0:
|
||
sub_vis = self.cams
|
||
|
||
images = [images[self.cams.index(cam)] for cam in sub_vis]
|
||
if cameras is None:
|
||
cameras = {'K': [], 'R':[], 'T':[]}
|
||
for key in cameras.keys():
|
||
cameras[key] = [self.cameras[cam][key] for cam in sub_vis]
|
||
for key in cameras.keys():
|
||
cameras[key] = np.stack([self.cameras[cam][key] for cam in sub_vis])
|
||
# 根据camera_back参数，控制相机向后退的距离
|
||
# 相机的光心的位置: -R.T @ T
|
||
if False:
|
||
R = cameras['R']
|
||
T = cameras['T']
|
||
cam_center = np.einsum('bij,bjk->bik', -R.transpose(0, 2, 1), T)
|
||
# 相机的朝向: R @ [0, 0, 1]
|
||
zdir = np.array([0., 0., 1.]).reshape(-1, 3, 1)
|
||
direction = np.einsum('bij,bjk->bik', R, zdir)
|
||
cam_center = cam_center - direction * 1
|
||
# 更新过后的相机的T: - R @ C
|
||
Tnew = - np.einsum('bij,bjk->bik', R, cam_center)
|
||
cameras['T'] = Tnew
|
||
else:
|
||
cameras['K'][:, 0, 0] /= camera_scale
|
||
cameras['K'][:, 1, 1] /= camera_scale
|
||
return self.writer.vis_smpl(render_data, nf, images, cameras, mode, add_back=add_back, extra_mesh=extra_mesh)
|
||
|
||
def read_skeleton(self, start, end):
|
||
keypoints3ds = []
|
||
for nf in range(start, end):
|
||
skelname = join(self.out, 'keypoints3d', '{:06d}.json'.format(nf))
|
||
skeletons = read_keypoints3d(skelname)
|
||
skeleton = [i for i in skeletons if i['id'] == self.pid]
|
||
assert len(skeleton) == 1, 'There must be only 1 keypoints3d, id = {} in {}'.format(self.pid, skelname)
|
||
keypoints3ds.append(skeleton[0]['keypoints3d'])
|
||
keypoints3ds = np.stack(keypoints3ds)
|
||
return keypoints3ds
|
||
|
||
def read_skel(self, nf, path=None, mode='none'):
|
||
if path is None:
|
||
path = self.skel_path
|
||
assert path is not None, 'please set the skeleton path'
|
||
if mode == 'a4d':
|
||
outname = join(path, '{}.txt'.format(nf))
|
||
assert os.path.exists(outname), outname
|
||
skels = readReasultsTxt(outname)
|
||
elif mode == 'none':
|
||
outname = join(path, '{:06d}.json'.format(nf))
|
||
assert os.path.exists(outname), outname
|
||
skels = readResultsJson(outname)
|
||
else:
|
||
import ipdb; ipdb.set_trace()
|
||
return skels
|
||
|
||
def read_smpl(self, nf, path=None):
|
||
if path is None:
|
||
path = self.skel_path
|
||
assert path is not None, 'please set the skeleton path'
|
||
outname = join(path, '{:06d}.json'.format(nf))
|
||
assert os.path.exists(outname), outname
|
||
datas = read_json(outname)
|
||
outputs = []
|
||
for data in datas:
|
||
for key in ['Rh', 'Th', 'poses', 'shapes']:
|
||
data[key] = np.array(data[key])
|
||
outputs.append(data)
|
||
return outputs
|