228 lines
9.2 KiB
Python
228 lines
9.2 KiB
Python
|
'''
|
|||
|
|
@ Date: 2022-07-15 19:25:33
|
|||
|
|
@ Author: Qing Shuai
|
|||
|
|
@ Mail: s_q@zju.edu.cn
|
|||
|
|
@ LastEditors: Qing Shuai
|
|||
|
|
@ LastEditTime: 2022-07-15 21:00:40
|
|||
|
|
@ FilePath: /EasyMocapPublic/easymocap/neuralbody/dataset/mirror.py
|
|||
|
|
'''
|
|||
|
|
from os.path import join
|
|||
|
|
import numpy as np
|
|||
|
|
import cv2
|
|||
|
|
from tqdm import trange
|
|||
|
|
import copy
|
|||
|
|
from .mvbase import BaseDataset, read_json, get_bounds
|
|||
|
|
from ...multistage.mirror import calc_mirror_transform
|
|||
|
|
import torch
|
|||
|
|
from .utils_sample import AABBwMask
|
|||
|
|
|
|||
|
|
def mirror_params(params, mirror, T0=np.eye(4, dtype=np.float32)):
|
|||
|
|
params = params.copy()
|
|||
|
|
M = np.eye(4, dtype=np.float32)
|
|||
|
|
M[:3] = calc_mirror_transform(torch.Tensor(mirror))[0].numpy()
|
|||
|
|
T1 = M @ T0
|
|||
|
|
T1[0, :] *= -1
|
|||
|
|
Rh = cv2.Rodrigues(params['Rh'])[0]
|
|||
|
|
Th = params['Th'].T
|
|||
|
|
Rnew = T1[:3, :3] @ Rh
|
|||
|
|
Tnew = T1[:3, :3] @ Th + T1[:3, 3:]
|
|||
|
|
params['Rh'] = cv2.Rodrigues(Rnew)[0].reshape(1, 3)
|
|||
|
|
params['Th'] = Tnew.T
|
|||
|
|
params['vertices'] = params['vertices'] @ T1[:3, :3].T + T1[:3, 3:].T
|
|||
|
|
return params
|
|||
|
|
|
|||
|
|
class MirrorDataset(BaseDataset):
|
|||
|
|
def get_allnames(self, root, subs, ranges, image_args):
|
|||
|
|
assert len(subs) == 1, 'Only support 1 sub-folder'
|
|||
|
|
infos = []
|
|||
|
|
index = 0
|
|||
|
|
nv = 0
|
|||
|
|
sub = subs[0]
|
|||
|
|
camera = self.cameras[sub].copy()
|
|||
|
|
K = camera['K'].copy()
|
|||
|
|
K[:2] *= image_args.scale
|
|||
|
|
camera['K'] = K
|
|||
|
|
# load mirror
|
|||
|
|
for nnf, nf in enumerate(trange(*ranges, desc='cache parameters')):
|
|||
|
|
annot = read_json(join(root, 'output-smpl-3d','smpl', sub, '{:06d}.json'.format(nf)))
|
|||
|
|
mirror = np.array(annot['mirror'], dtype=np.float32)
|
|||
|
|
vertices = read_json(join(root, 'output-smpl-3d','vertices', sub, '{:06d}.json'.format(nf)))
|
|||
|
|
annots = annot['annots']
|
|||
|
|
imgname = join(root, image_args.root, sub, '{:06d}{}'.format(nf, image_args.ext))
|
|||
|
|
info = {
|
|||
|
|
'imgname': imgname,
|
|||
|
|
'sub': sub,
|
|||
|
|
'frame': nf,
|
|||
|
|
'nf': nnf,
|
|||
|
|
'nv': nv,
|
|||
|
|
'index': index,
|
|||
|
|
'camera': camera,
|
|||
|
|
'flip_lr': False,
|
|||
|
|
'pid': 0
|
|||
|
|
}
|
|||
|
|
annots[0]['vertices'] = vertices[0]['vertices']
|
|||
|
|
annots[1]['vertices'] = vertices[0]['vertices']
|
|||
|
|
info_mirror = info.copy()
|
|||
|
|
for annot in annots:
|
|||
|
|
for key, val in annot.items():
|
|||
|
|
if key == 'id':continue
|
|||
|
|
annot[key] = np.array(val, dtype=np.float32)
|
|||
|
|
info['params'] = annots[0]
|
|||
|
|
info_mirror['params'] = mirror_params(annots[0], mirror)
|
|||
|
|
info_mirror['flip_lr'] = True
|
|||
|
|
info_mirror['pid'] = 1
|
|||
|
|
info_mirror['nv'] = 1
|
|||
|
|
info['bounds'] = get_bounds(info['params']['vertices'], delta=0.1)
|
|||
|
|
info_mirror['bounds'] = get_bounds(info_mirror['params']['vertices'], delta=0.1)
|
|||
|
|
info['params']['R'] = cv2.Rodrigues(info['params']['Rh'])[0]
|
|||
|
|
info_mirror['params']['R'] = cv2.Rodrigues(info_mirror['params']['Rh'])[0]
|
|||
|
|
infos.append(info)
|
|||
|
|
infos.append(info_mirror)
|
|||
|
|
return infos
|
|||
|
|
|
|||
|
|
def get_objects(self, root, info, object_keys, object_args):
|
|||
|
|
sub, frame, pid = info['sub'], info['frame'], info['pid']
|
|||
|
|
mskname = join(root, object_args['human_0'].args.reader.mask.root, sub, '{:06d}_{}.png'.format(frame, pid))
|
|||
|
|
msk = cv2.imread(mskname, 0)
|
|||
|
|
msk = self.scale_and_undistort(msk, info, undis=False)
|
|||
|
|
if info['flip_lr']:
|
|||
|
|
msk = cv2.flip(msk, 1)
|
|||
|
|
msk = msk > 0
|
|||
|
|
obj = AABBwMask(split=self.split, bounds=info['bounds'],
|
|||
|
|
mask=msk,
|
|||
|
|
label=None,
|
|||
|
|
dilate=False, # 这里假设mask很准，或者使用了patch
|
|||
|
|
rate_body=0.85)
|
|||
|
|
for key in ['R', 'Rh', 'Th', 'vertices', 'poses', 'shapes']:
|
|||
|
|
obj.feature[key] = info['params'][key]
|
|||
|
|
vertices_canonical = (info['params']['vertices'] - info['params']['Th']) @ info['params']['R'].T.T
|
|||
|
|
obj.feature['bounds_canonical'] = get_bounds(vertices_canonical,
|
|||
|
|
delta=object_args['human_0'].args.reader.vertices.padding)
|
|||
|
|
return {'human_0': obj}
|
|||
|
|
|
|||
|
|
def read_image(self, imgname, image_args, info, isgray=False, skip_mask=False, mask_global='_0.png'):
|
|||
|
|
if info['flip_lr']:
|
|||
|
|
return super().read_image(imgname, image_args, info, isgray, skip_mask, mask_global='_1.png')
|
|||
|
|
else:
|
|||
|
|
return super().read_image(imgname, image_args, info, isgray, skip_mask, mask_global)
|
|||
|
|
|
|||
|
|
def augment_rotation(self, info, rot, rotvec=[0., 1., 0.]):
|
|||
|
|
# method 1: rotate on its self
|
|||
|
|
rot = rot / 180 * np.pi
|
|||
|
|
rotvec = np.array([rotvec],dtype=np.float32)
|
|||
|
|
R = cv2.Rodrigues(rot*rotvec)[0]
|
|||
|
|
info = copy.deepcopy(info)
|
|||
|
|
params = info['params']
|
|||
|
|
T0 = params['Th']
|
|||
|
|
R0 = cv2.Rodrigues(params['Rh'])[0]
|
|||
|
|
v0 = params['vertices']
|
|||
|
|
v1 = (v0 - T0) @ R.T + T0
|
|||
|
|
params['vertices'] = v1
|
|||
|
|
R1 = R @ R0
|
|||
|
|
params['Rh'] = cv2.Rodrigues(R1)[0].reshape(1, 3)
|
|||
|
|
params['R'] = R1
|
|||
|
|
info['bounds'] = get_bounds(params['vertices'], delta=0.1)
|
|||
|
|
return info
|
|||
|
|
|
|||
|
|
class MirrorDatasetDemo(MirrorDataset):
|
|||
|
|
def __init__(self, keyframes, **cfg):
|
|||
|
|
self.keyframes = keyframes
|
|||
|
|
# 注意：初始化里面调用了get_allnames，所以要预先写进去
|
|||
|
|
super().__init__(**cfg)
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return super().__len__()//2
|
|||
|
|
|
|||
|
|
def get_allnames(self, root, subs, ranges, image_args):
|
|||
|
|
infos = super().get_allnames(root, subs, ranges, image_args)
|
|||
|
|
infos_new = []
|
|||
|
|
for nf in range(len(infos)//2):
|
|||
|
|
infos_l = infos[2*nf]
|
|||
|
|
infos_r = infos[2*nf+1]
|
|||
|
|
if nf in self.keyframes:
|
|||
|
|
for rot in range(0, 360, 4):
|
|||
|
|
infos_r_ = self.augment_rotation(infos_r, rot, rotvec=[0,1,0])
|
|||
|
|
infos_new.append(infos_l)
|
|||
|
|
infos_new.append(infos_r_)
|
|||
|
|
else:
|
|||
|
|
infos_new.append(infos_l)
|
|||
|
|
infos_new.append(infos_r)
|
|||
|
|
return infos_new
|
|||
|
|
|
|||
|
|
def __getitem__(self, index):
|
|||
|
|
left = super().__getitem__(2*index)
|
|||
|
|
right = super().__getitem__(2*index+1)
|
|||
|
|
return {'left': left, 'right': right, 'meta': {'type': 'mirror'}}
|
|||
|
|
|
|||
|
|
class MirrorDatasetDemoCool(BaseDataset):
|
|||
|
|
def __init__(self, **cfg):
|
|||
|
|
super().__init__(**cfg)
|
|||
|
|
|
|||
|
|
def get_allnames(self, root, subs, ranges, image_args):
|
|||
|
|
assert len(subs) == 1, 'Only support 1 sub-folder'
|
|||
|
|
infos = []
|
|||
|
|
index = 0
|
|||
|
|
nv = 0
|
|||
|
|
sub = subs[0]
|
|||
|
|
camera = self.cameras[sub].copy()
|
|||
|
|
K = camera['K'].copy()
|
|||
|
|
K[:2] *= image_args.scale
|
|||
|
|
camera['K'] = K
|
|||
|
|
# load mirror
|
|||
|
|
for nnf, nf in enumerate(trange(*ranges, desc='cache parameters')):
|
|||
|
|
annot = read_json(join(root, 'output-smpl-3d','smpl', sub, '{:06d}.json'.format(nf)))
|
|||
|
|
mirror = np.array(annot['mirror'], dtype=np.float32)
|
|||
|
|
vertices = read_json(join(root, 'output-smpl-3d','vertices', sub, '{:06d}.json'.format(nf)))
|
|||
|
|
# only use outer person
|
|||
|
|
annots = annot['annots'][:1]
|
|||
|
|
imgname = join(root, image_args.root, sub, '{:06d}{}'.format(nf, image_args.ext))
|
|||
|
|
info = {
|
|||
|
|
'imgname': imgname,
|
|||
|
|
'sub': sub,
|
|||
|
|
'frame': nf,
|
|||
|
|
'nf': nnf,
|
|||
|
|
'nv': nv,
|
|||
|
|
'index': index,
|
|||
|
|
'camera': camera,
|
|||
|
|
'flip_lr': False,
|
|||
|
|
'pid': 0
|
|||
|
|
}
|
|||
|
|
annots[0]['vertices'] = vertices[0]['vertices']
|
|||
|
|
for annot in annots:
|
|||
|
|
for key, val in annot.items():
|
|||
|
|
if key == 'id':continue
|
|||
|
|
annot[key] = np.array(val, dtype=np.float32)
|
|||
|
|
info['params'] = annots[0]
|
|||
|
|
# TODO: augment the rotation
|
|||
|
|
# method 1: rotate on its self
|
|||
|
|
rot = nnf / 180 * np.pi
|
|||
|
|
rotvec = np.array([[0., 1., 0.]],dtype=np.float32)
|
|||
|
|
R = cv2.Rodrigues(rot*rotvec)[0]
|
|||
|
|
T0 = info['params']['Th']
|
|||
|
|
R0 = cv2.Rodrigues(info['params']['Rh'])[0]
|
|||
|
|
v0 = info['params']['vertices']
|
|||
|
|
v1 = (v0 - T0) @ R.T + T0
|
|||
|
|
info['params']['vertices'] = v1
|
|||
|
|
R1 = R @ R0
|
|||
|
|
info['params']['Rh'] = cv2.Rodrigues(R1)[0].reshape(1, 3)
|
|||
|
|
info['params']['R'] = R1
|
|||
|
|
info['bounds'] = get_bounds(info['params']['vertices'], delta=0.1)
|
|||
|
|
infos.append(info)
|
|||
|
|
return infos
|
|||
|
|
|
|||
|
|
def get_objects(self, root, info, object_keys, object_args):
|
|||
|
|
sub, frame, pid = info['sub'], info['frame'], info['pid']
|
|||
|
|
mskname = join(root, object_args['human_0'].args.reader.mask.root, sub, '{:06d}_{}.png'.format(frame, pid))
|
|||
|
|
msk = cv2.imread(mskname, 0)
|
|||
|
|
msk = self.scale_and_undistort(msk, info, undis=False)
|
|||
|
|
if info['flip_lr']:
|
|||
|
|
msk = cv2.flip(msk, 1)
|
|||
|
|
msk = msk > 0
|
|||
|
|
obj = AABBwMask(split=self.split, bounds=info['bounds'],
|
|||
|
|
mask=msk,
|
|||
|
|
label=None,
|
|||
|
|
dilate=True,
|
|||
|
|
rate_body=0.85)
|
|||
|
|
for key in ['R', 'Rh', 'Th', 'vertices', 'poses', 'shapes']:
|
|||
|
|
obj.feature[key] = info['params'][key]
|
|||
|
|
return {'human_0': obj}
|