175 lines
6.2 KiB
Python
175 lines
6.2 KiB
Python
import numpy as np
|
|
import cv2
|
|
from ..dataset.config import CONFIG
|
|
from ..config import load_object
|
|
from ..mytools.debug_utils import log, mywarn, myerror
|
|
import torch
|
|
from tqdm import tqdm, trange
|
|
|
|
def svd_rot(src, tgt, reflection=False, debug=False):
|
|
# optimum rotation matrix of Y
|
|
A = np.matmul(src.transpose(0, 2, 1), tgt)
|
|
U, s, Vt = np.linalg.svd(A, full_matrices=False)
|
|
V = Vt.transpose(0, 2, 1)
|
|
T = np.matmul(V, U.transpose(0, 2, 1))
|
|
# does the current solution use a reflection?
|
|
have_reflection = np.linalg.det(T) < 0
|
|
|
|
# if that's not what was specified, force another reflection
|
|
V[have_reflection, :, -1] *= -1
|
|
s[have_reflection, -1] *= -1
|
|
T = np.matmul(V, U.transpose(0, 2, 1))
|
|
if debug:
|
|
err = np.linalg.norm(tgt - src @ T.T, axis=1)
|
|
print('[svd] ', err)
|
|
return T
|
|
|
|
def batch_invRodrigues(rot):
|
|
res = []
|
|
for r in rot:
|
|
v = cv2.Rodrigues(r)[0]
|
|
res.append(v)
|
|
res = np.stack(res)
|
|
return res[:, :, 0]
|
|
|
|
class BaseInit:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
def __call__(self, body_model, body_params, infos):\
|
|
return body_params
|
|
|
|
class Remove(BaseInit):
|
|
def __init__(self, key, index) -> None:
|
|
super().__init__()
|
|
self.key = key
|
|
self.index = index
|
|
|
|
def __call__(self, body_model, body_params, infos):
|
|
infos[self.key][..., self.index, :] = 0
|
|
return super().__call__(body_model, body_params, infos)
|
|
|
|
class CheckKeypoints:
|
|
def __init__(self, type) -> None:
|
|
# this class is used to check if the provided keypoints3d
|
|
self.type = type
|
|
self.body_config = CONFIG[type]
|
|
self.hand_config = CONFIG['hand']
|
|
|
|
def __call__(self, body_model, body_params, infos):
|
|
for key in ['keypoints3d', 'handl3d', 'handr3d']:
|
|
if key not in infos.keys(): continue
|
|
keypoints = infos[key]
|
|
conf = keypoints[..., -1]
|
|
keypoints[conf<0.1] = 0
|
|
if key == 'keypoints3d':
|
|
continue
|
|
import ipdb;ipdb.set_trace()
|
|
# limb_length = np.linalg.norm(keypoints[:, , :3], axis=2)
|
|
return body_params
|
|
|
|
class InitRT:
|
|
def __init__(self, torso) -> None:
|
|
self.torso = torso
|
|
|
|
def __call__(self, body_model, body_params, infos):
|
|
keypoints3d = infos['keypoints3d']
|
|
if torch.is_tensor(keypoints3d):
|
|
keypoints3d = keypoints3d.detach().cpu().numpy()
|
|
temp_joints = body_model.keypoints(body_params, return_tensor=False)
|
|
|
|
torso = keypoints3d[..., self.torso, :3].copy()
|
|
torso_temp = temp_joints[..., self.torso, :3].copy()
|
|
# here use the first id of torso as the rotation center
|
|
root, root_temp = torso[..., :1, :], torso_temp[..., :1, :]
|
|
torso = torso - root
|
|
torso_temp = torso_temp - root_temp
|
|
conf = (keypoints3d[..., self.torso, 3] > 0.).all(axis=-1)
|
|
if not conf.all():
|
|
myerror("The torso in frames {} is not valid, please check the 3d keypoints".format(np.where(~conf)))
|
|
if len(torso.shape) == 3:
|
|
R = svd_rot(torso_temp, torso)
|
|
R_flat = R
|
|
T = np.matmul(- root_temp, R.transpose(0, 2, 1)) + root
|
|
else:
|
|
R_flat = svd_rot(torso_temp.reshape(-1, *torso_temp.shape[-2:]), torso.reshape(-1, *torso.shape[-2:]))
|
|
R = R_flat.reshape(*torso.shape[:2], 3, 3)
|
|
T = np.matmul(- root_temp, R.swapaxes(-1, -2)) + root
|
|
for nf in np.where(~conf)[0]:
|
|
# copy previous frames
|
|
mywarn('copy {} from {}'.format(nf, nf-1))
|
|
R[nf] = R[nf-1]
|
|
T[nf] = T[nf-1]
|
|
body_params['Th'] = T[..., 0, :]
|
|
rvec = batch_invRodrigues(R_flat)
|
|
if len(torso.shape) > 3:
|
|
rvec = rvec.reshape(*torso.shape[:2], 3)
|
|
body_params['Rh'] = rvec
|
|
return body_params
|
|
|
|
def __str__(self) -> str:
|
|
return "[Initialize] svd with torso: {}".format(self.torso)
|
|
|
|
class TriangulatorWrapper:
|
|
def __init__(self, module, args):
|
|
self.triangulator = load_object(module, args)
|
|
|
|
def __call__(self, body_model, body_params, infos):
|
|
infos['RT'] = torch.cat([infos['Rc'], infos['Tc']], dim=-1)
|
|
data = {
|
|
'RT': infos['RT'].numpy(),
|
|
}
|
|
for key in self.triangulator.keys:
|
|
if key not in infos.keys():
|
|
continue
|
|
data[key] = infos[key].numpy()
|
|
data[key+'_unproj'] = infos[key+'_unproj'].numpy()
|
|
data[key+'_distort'] = infos[key+'_distort'].numpy()
|
|
results = self.triangulator(data)[0]
|
|
for key, val in results.items():
|
|
if key == 'id': continue
|
|
infos[key] = torch.Tensor(val[None].astype(np.float32))
|
|
body_params = body_model.init_params(nFrames=1, add_scale=True)
|
|
return body_params
|
|
|
|
class CheckRT:
|
|
def __init__(self, T_thres, window):
|
|
self.T_thres = T_thres
|
|
self.window = window
|
|
|
|
def __call__(self, body_model, body_params, infos):
|
|
Th = body_params['Th']
|
|
if len(Th.shape) == 3:
|
|
for nper in range(Th.shape[1]):
|
|
for nf in trange(1, Th.shape[0], desc='Check Th of {}'.format(nper)):
|
|
if nf > self.window:
|
|
tpre = Th[nf-self.window:nf, nper]
|
|
else:
|
|
tpre = Th[:nf, nper]
|
|
tpre = tpre.mean(axis=0)
|
|
tnow = Th[nf , nper]
|
|
dist = np.linalg.norm(tnow - tpre)
|
|
if dist > self.T_thres:
|
|
mywarn('[Check Th] distance in frame {} = {} larger than {}'.format(nf, dist, self.T_thres))
|
|
Th[nf, nper] = tpre
|
|
body_params['Th'] = Th
|
|
return body_params
|
|
|
|
class Scale:
|
|
def __init__(self, keys):
|
|
self.keys = keys
|
|
|
|
def __call__(self, body_model, body_params, infos):
|
|
scale = body_params.pop('scale')[0, 0]
|
|
if scale < 1.1 and scale > 0.9:
|
|
return body_params
|
|
print('scale = ', scale)
|
|
for key in self.keys:
|
|
if key not in infos.keys():
|
|
continue
|
|
infos[key] /= scale
|
|
infos['Tc'] /= scale
|
|
infos['RT'][..., -1] *= scale
|
|
infos['scale'] = scale
|
|
return body_params
|