1772 lines
78 KiB
Python
1772 lines
78 KiB
Python
'''
|
||
@ Date: 2022-03-22 16:11:44
|
||
@ Author: Qing Shuai
|
||
@ Mail: s_q@zju.edu.cn
|
||
@ LastEditors: Qing Shuai
|
||
@ LastEditTime: 2022-07-25 11:51:50
|
||
@ FilePath: /EasyMocapPublic/easymocap/multistage/fitting.py
|
||
'''
|
||
# This function provides a realtime fitting interface
|
||
from collections import namedtuple
|
||
from time import time, sleep
|
||
import numpy as np
|
||
import cv2
|
||
import torch
|
||
import copy
|
||
|
||
from ..config.baseconfig import load_object_from_cmd
|
||
from ..mytools.debug_utils import log, mywarn
|
||
from ..mytools import Timer
|
||
from ..config import Config
|
||
from ..mytools.triangulator import iterative_triangulate
|
||
from ..bodymodel.base import Params
|
||
from .torchgeometry import axis_angle_to_euler, euler_to_axis_angle
|
||
|
||
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
|
||
''' Calculates the rotation matrices for a batch of rotation vectors
|
||
Parameters
|
||
----------
|
||
rot_vecs: torch.tensor Nx3
|
||
array of N axis-angle vectors
|
||
Returns
|
||
-------
|
||
R: torch.tensor Nx3x3
|
||
The rotation matrices for the given axis-angle parameters
|
||
'''
|
||
|
||
batch_size = rot_vecs.shape[0]
|
||
device = rot_vecs.device
|
||
|
||
angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
|
||
rot_dir = rot_vecs / angle
|
||
|
||
cos = torch.unsqueeze(torch.cos(angle), dim=1)
|
||
sin = torch.unsqueeze(torch.sin(angle), dim=1)
|
||
|
||
# Bx1 arrays
|
||
rx, ry, rz = torch.split(rot_dir, 1, dim=1)
|
||
K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
|
||
|
||
zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
|
||
K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
|
||
.view((batch_size, 3, 3))
|
||
|
||
ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
|
||
rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
|
||
return rot_mat
|
||
|
||
from scipy.spatial.transform import Rotation
|
||
def aa2euler(aa):
|
||
aa = np.array(aa)
|
||
R = cv2.Rodrigues(aa)[0]
|
||
# res = Rotation.from_dcm(R).as_euler('XYZ', degrees=True)
|
||
res = Rotation.from_matrix(R).as_euler('XYZ', degrees=False)
|
||
return np.round(res, 2).tolist()
|
||
|
||
def rotmat2euler(rot):
|
||
res = Rotation.from_matrix(rot).as_euler('XYZ', degrees=True)
|
||
return res
|
||
|
||
def euler2rotmat(euler):
|
||
res = Rotation.from_euler('XYZ', euler, degrees=True)
|
||
return res.as_matrix()
|
||
|
||
def batch_rodrigues_jacobi(rvec):
|
||
shape = rvec.shape
|
||
rvec = rvec.view(-1, 3)
|
||
device = rvec.device
|
||
dSkew = torch.zeros(3, 9, device=device)
|
||
dSkew[0, 5] = -1
|
||
dSkew[1, 6] = -1
|
||
dSkew[2, 1] = -1
|
||
dSkew[0, 7] = 1
|
||
dSkew[1, 2] = 1
|
||
dSkew[2, 3] = 1
|
||
dSkew = dSkew[None]
|
||
theta = torch.norm(rvec, dim=-1, keepdim=True) + 1e-5
|
||
c = torch.cos(theta)
|
||
s = torch.sin(theta)
|
||
c1 = 1 - c
|
||
itheta = 1 / theta
|
||
r = rvec / theta
|
||
zeros = torch.zeros_like(r[:, :1])
|
||
rx, ry, rz = torch.split(r, 1, dim=1)
|
||
rrt = torch.matmul(r[:, :, None], r[:, None, :])
|
||
skew = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
|
||
.view((r.shape[0], 3, 3))
|
||
I = torch.eye(3, device=rvec.device, dtype=rvec.dtype)[None]
|
||
rot_mat = I + s[:, None] * skew + c1[:, None] * torch.bmm(skew, skew)
|
||
|
||
drrt = torch.stack([
|
||
rx + rx, ry, rz, ry, zeros, zeros, rz, zeros, zeros,
|
||
zeros, rx, zeros, rx, ry + ry, rz, zeros, rz, zeros,
|
||
zeros, zeros, rx, zeros, zeros, ry, rx, ry, rz + rz
|
||
], dim=-1).view((r.shape[0], 3, 9))
|
||
jacobi = torch.zeros((r.shape[0], 3, 9), device=rvec.device, dtype=rvec.dtype)
|
||
for i in range(3):
|
||
ri = r[:, i:i+1]
|
||
a0 = -s * ri
|
||
a1 = (s - 2*c1*itheta)*ri
|
||
a2 = c1 * itheta
|
||
a3 = (c-s*itheta)*ri
|
||
a4 = s * itheta
|
||
jaco = a0[:, None] * I + a1[:, None] * rrt + a2[:, None] * drrt[:, i].view(-1, 3, 3) + a3[:, None] * skew + a4[:, None] * dSkew[:, i].view(-1, 3, 3)
|
||
jacobi[:, i] = jaco.view(-1, 9)
|
||
rot_mat = rot_mat.view(*shape[:-1], 3, 3)
|
||
jacobi = jacobi.view(*shape[:-1], 3, 9)
|
||
return rot_mat, jacobi
|
||
|
||
def getJacobianOfRT(rvec, tvec, joints):
|
||
# joints: (bn, nJ, 3)
|
||
dtype, device = rvec.dtype, rvec.device
|
||
bn, nJoints = joints.shape[:2]
|
||
# jacobiToRvec: (bn, 3, 9) // tested by OpenCV and PyTorch
|
||
Rot, jacobiToRvec = batch_rodrigues_jacobi(rvec)
|
||
I3 = torch.eye(3, dtype=dtype, device=device)[None]
|
||
# jacobiJ_R: (bn, nJ, 3, 3+3+3) => (bn, nJ, 3, 9)
|
||
# // flat by column:
|
||
# // x, 0, 0 | y, 0, 0 | z, 0, 0
|
||
# // 0, x, 0 | 0, y, 0 | 0, z, 0
|
||
# // 0, 0, x | 0, 0, y | 0, 0, z
|
||
jacobi_J_R = torch.zeros((bn, nJoints, 3, 9), dtype=dtype, device=device)
|
||
jacobi_J_R[:, :, 0, :3] = joints
|
||
jacobi_J_R[:, :, 1, 3:6] = joints
|
||
jacobi_J_R[:, :, 2, 6:9] = joints
|
||
# jacobi_J_rvec: (bn, nJ, 3, 3)
|
||
jacobi_J_rvec = torch.matmul(jacobi_J_R, jacobiToRvec[:, None].transpose(-1, -2))
|
||
# if True: # 测试自动梯度
|
||
# def test_func(rvec):
|
||
# Rot = batch_rodrigues(rvec[None])[0]
|
||
# joints_new = joints[0] @ Rot.t()
|
||
# return joints_new
|
||
# jac_J_rvec = torch.autograd.functional.jacobian(test_func, rvec[0])
|
||
# my_j = jacobi_joints_RT[0, ..., :3]
|
||
# jacobi_J_tvec: (bn, nJx3, 3)
|
||
jacobi_J_tvec = I3[None].expand(bn, nJoints, -1, -1)
|
||
jacobi_J_rt = torch.cat([jacobi_J_rvec, jacobi_J_tvec], dim=-1)
|
||
return Rot, jacobiToRvec, jacobi_J_rt
|
||
|
||
class Model:
|
||
rootIdx = 0
|
||
parents = []
|
||
|
||
|
||
INDEX_HALF = [0,1,2,3,4,5,6,7,15,16,17,18]
|
||
|
||
class LowPassFilter:
|
||
def __init__(self):
|
||
self.prev_raw_value = None
|
||
self.prev_filtered_value = None
|
||
|
||
def process(self, value, alpha):
|
||
if self.prev_raw_value is None:
|
||
s = value
|
||
else:
|
||
s = alpha * value + (1.0 - alpha) * self.prev_filtered_value
|
||
self.prev_raw_value = value
|
||
self.prev_filtered_value = s
|
||
return s
|
||
|
||
class OneEuroFilter:
|
||
def __init__(self, mincutoff=1.0, beta=0.0, dcutoff=1.0, freq=30):
|
||
self.freq = freq
|
||
self.mincutoff = mincutoff
|
||
self.beta = beta
|
||
self.dcutoff = dcutoff
|
||
self.x_filter = LowPassFilter()
|
||
self.dx_filter = LowPassFilter()
|
||
|
||
def compute_alpha(self, cutoff):
|
||
te = 1.0 / self.freq
|
||
tau = 1.0 / (2 * np.pi * cutoff)
|
||
return 1.0 / (1.0 + tau / te)
|
||
|
||
def process(self, x):
|
||
prev_x = self.x_filter.prev_raw_value
|
||
dx = 0.0 if prev_x is None else (x - prev_x) * self.freq
|
||
edx = self.dx_filter.process(dx, self.compute_alpha(self.dcutoff))
|
||
cutoff = self.mincutoff + self.beta * np.abs(edx)
|
||
return self.x_filter.process(x, self.compute_alpha(cutoff))
|
||
|
||
class BaseBody:
|
||
def __init__(self, cfg_triangulator, cfg_model, cfg) -> None:
|
||
self.triangulator = load_object_from_cmd(cfg_triangulator, [])
|
||
self.body_model = load_object_from_cmd(cfg_model, ['args.use_pose_blending', False, 'args.device', 'cpu'])
|
||
self.cfg = cfg
|
||
self.register_from_lbs(self.body_model)
|
||
|
||
def register_from_lbs(self, body_model):
|
||
kintree_shape = np.array(self.cfg.shape.kintree)
|
||
self.nJoints = body_model.J_regressor.shape[0]
|
||
self.k_shapeBlend = body_model.j_shapedirs[self.nJoints:]
|
||
self.j_shapeBlend = body_model.j_shapedirs[:self.nJoints]
|
||
self.jacobian_limb_shapes = self.k_shapeBlend[kintree_shape[:, 1]] - self.k_shapeBlend[kintree_shape[:, 0]]
|
||
self.k_template = body_model.j_v_template[self.nJoints:]
|
||
self.j_template = body_model.j_v_template[:self.nJoints]
|
||
self.k_weights = body_model.j_weights[self.nJoints:]
|
||
self.j_weights = body_model.j_weights[:self.nJoints]
|
||
parents = body_model.parents[1:].cpu().numpy()
|
||
child = np.arange(1, parents.shape[0]+1, dtype=np.int64)
|
||
self.kintree = np.stack([parents, child], axis=1)
|
||
self.parents = np.zeros(parents.shape[0]+1, dtype=int) - 1
|
||
self.parents[self.kintree[:, 1]] = self.kintree[:, 0]
|
||
self.rootIdx = 0
|
||
self.time = time()
|
||
|
||
def rotation_matrix_from_3x3(A):
|
||
U, s, Vt = np.linalg.svd(A, full_matrices=False)
|
||
V = Vt.T
|
||
T = np.dot(V, U.T)
|
||
# does the current solution use a reflection?
|
||
have_reflection = np.linalg.det(T) < 0
|
||
|
||
# if that's not what was specified, force another reflection
|
||
if have_reflection:
|
||
V[:,-1] *= -1
|
||
s[-1] *= -1
|
||
T = np.dot(V, U.T)
|
||
return T
|
||
|
||
def svd_rot(src, tgt, reflection=False, debug=True):
|
||
# optimum rotation matrix of Y
|
||
A = np.dot(src.T, tgt)
|
||
U, s, Vt = np.linalg.svd(A, full_matrices=False)
|
||
V = Vt.T
|
||
T = np.dot(V, U.T)
|
||
# does the current solution use a reflection?
|
||
have_reflection = np.linalg.det(T) < 0
|
||
|
||
# if that's not what was specified, force another reflection
|
||
if reflection != have_reflection:
|
||
V[:,-1] *= -1
|
||
s[-1] *= -1
|
||
T = np.dot(V, U.T)
|
||
if debug:
|
||
err = np.linalg.norm(tgt - src @ T.T, axis=1)
|
||
print('[svd] ', err)
|
||
return T
|
||
|
||
def normalize(vector):
|
||
return vector/np.linalg.norm(vector)
|
||
|
||
def rad_from_2vec(vec1, vec2):
|
||
return np.arccos((normalize(vec1)*normalize(vec2)).sum())
|
||
|
||
def smoothing_factor(t_e, cutoff):
|
||
r = 2 * 3.14 * cutoff * t_e
|
||
return r / (r + 1)
|
||
|
||
def exponential_smoothing(a, x, x_prev):
|
||
return a * x + (1 - a) * x_prev
|
||
|
||
FilterResult = namedtuple('FilterResult', ['x', 'dx', 'v', 't'])
|
||
|
||
class MyFilter:
|
||
def __init__(self, key, filled, min_cutoff=1.0, d_cutoff=1.0,
|
||
beta=0.1) -> None:
|
||
self.key = key
|
||
self.fill_result = filled
|
||
self.min_cutoff = min_cutoff
|
||
self.d_cutoff = d_cutoff
|
||
self.beta = beta
|
||
self.init = False
|
||
self.records = []
|
||
self.result = None
|
||
self.counter = 0
|
||
self.data = []
|
||
self.conf = []
|
||
self.smooth_record = []
|
||
|
||
def fill(self, value, conf):
|
||
filled = conf < 0.1
|
||
if filled.sum() >= 1:
|
||
value[filled] = self.fill_result[0][filled]
|
||
if self.key == 'Rh':
|
||
value = rotation_matrix_from_3x3(value.reshape(3, 3))
|
||
value = cv2.Rodrigues(value)[0].reshape(3,)
|
||
if (value < 0).all():
|
||
value = -value
|
||
return value[None]
|
||
|
||
def __call__(self, value, conf):
|
||
self.counter += 1
|
||
x, v = value[0], conf[0]
|
||
if self.key == 'Rh':
|
||
x = cv2.Rodrigues(x)[0].reshape(-1)
|
||
v = np.zeros((9,)) + v[0]
|
||
t = np.zeros_like(x)
|
||
t[v>0.1] = self.counter
|
||
self.data.append(x)
|
||
self.conf.append(v)
|
||
if len(self.smooth_record) == 0:
|
||
if self.key == 'Rh':
|
||
start = x
|
||
else:
|
||
# start = self.fill_result[0]
|
||
start = x
|
||
smoothed = FilterResult(start, np.zeros_like(x), np.zeros_like(x), t)
|
||
self.smooth_record.append(smoothed)
|
||
if len(self.data) < 3:
|
||
return self.fill(x, v)
|
||
data = np.array(self.data)
|
||
conf = np.array(self.conf)
|
||
smoothed = self.smooth_record[-1]
|
||
# 预计的速度
|
||
dx_new = x - smoothed.x
|
||
# 滤波器可见，当前可见
|
||
flag_vis = (smoothed.v > 0.1) & (v > 0.1)
|
||
# - 速度异常的，移除掉，认为当前帧不可见
|
||
flag_outlier = (np.abs(smoothed.dx) > 0.05) & (np.abs(dx_new - smoothed.dx)/(1e-5 + smoothed.dx) > 2.)
|
||
if self.key != 'Rh':
|
||
v[flag_vis&flag_outlier] = 0.
|
||
# 滤波器不可见，当前可见，速度打折，认为是新增的帧
|
||
flag_new = (smoothed.v < 0.1)&(v>0.1)
|
||
dx_new[flag_new] /= 3
|
||
# 滤波器可见，当前不可见，速度使用滤波器的速度
|
||
flag_unvis = (v<0.1) & (conf[-2] < 0.1)
|
||
dx_new[flag_unvis] = smoothed.dx[flag_unvis]
|
||
# 滤波器不可见，当前也不可见，速度清0
|
||
dx_new[(v<0.1)&(smoothed.v<0.1)] = 0.
|
||
# 实际估计出来的速度，这里要去掉不可见的地方
|
||
# 混合的权重使用 0.7, 0.3默认全部使用新的一半
|
||
weight_dx = np.zeros_like(dx_new) + 0.7
|
||
dx_smoothed = smoothed.dx*(1-weight_dx) + dx_new*weight_dx
|
||
smoothed_value = smoothed.x + dx_smoothed
|
||
v_new = smoothed.v.copy()
|
||
v_new = v_new * (1-weight_dx) + v*weight_dx
|
||
t_new = smoothed.t.copy()
|
||
t_new[v>0.1] = t[v>0.1]
|
||
smooth_new = FilterResult(smoothed_value, dx_smoothed, v_new, t_new)
|
||
self.smooth_record.append(smooth_new)
|
||
if self.counter == 1000:
|
||
if self.key == 'poses':
|
||
import matplotlib.pyplot as plt
|
||
xrange = np.arange(0, data.shape[0])
|
||
smoothed = np.array([d.x for d in self.smooth_record])
|
||
for nj in range(data.shape[1]):
|
||
valid = conf[:, nj] > 0.
|
||
plt.scatter(xrange[valid], data[valid, nj])
|
||
# yhat = savgol_filter(data[:, nj], data.shape[0], 3)
|
||
# plt.plot(yhat)
|
||
plt.plot(smoothed)
|
||
plt.show()
|
||
import ipdb;ipdb.set_trace()
|
||
# return self.fill(x, v)
|
||
return self.fill(smooth_new.x, smooth_new.v)
|
||
|
||
def __call__0(self, value, conf):
|
||
self.counter += 1
|
||
x, v = value[0], conf[0]
|
||
if self.key == 'Rh':
|
||
x = cv2.Rodrigues(x)[0].reshape(-1)
|
||
v = np.zeros((9,)) + v[0]
|
||
if self.result is None:
|
||
result = FilterResult(x, np.zeros_like(x), v, (v>0)*self.counter)
|
||
self.result = result
|
||
self.records.append(result)
|
||
return self.fill(result.x, result.v)
|
||
# return self.fill(x, v)
|
||
# 维护一个前一帧的，去除outlier
|
||
prev = self.result
|
||
t = prev.t.copy()
|
||
t[v>0.] = self.counter
|
||
dx = x - prev.x # 这里直接使用与之前的结果的差了，避免多帧不可见，然后速度过大
|
||
MAX_DX = 1.
|
||
WINDOW = 31
|
||
not_valid = ((np.abs(dx) > MAX_DX) & (prev.v > 0.1))|\
|
||
(t-prev.t > WINDOW)
|
||
v[not_valid] = 0.
|
||
x_all = np.stack([r.x for r in self.records[-WINDOW:]])
|
||
v_all = np.stack([r.v for r in self.records[-WINDOW:]])
|
||
dx_all = np.stack([r.dx for r in self.records[-WINDOW:]])
|
||
v_sum = v_all.sum(axis=0)
|
||
dx_mean = (dx_all * v_all).sum(axis=0)/(1e-5 + v_all.sum(axis=0))
|
||
# if (x_all.shape[0] > 30) & (self.counter % 40 == 0):
|
||
if True:
|
||
x_mean = (x_all * v_all).sum(axis=0)/(1e-5 + v_all.sum(axis=0))
|
||
x_pred = x_mean
|
||
dx_pred = np.zeros_like(x_pred)
|
||
elif x_all.shape[0] >= 5:
|
||
# 进行smooth
|
||
axt = np.zeros((2, x_all.shape[1]))
|
||
xrange = np.arange(x_all.shape[0]).reshape(-1, 1)
|
||
A0 = np.hstack([xrange, np.ones((x_all.shape[0], 1))])
|
||
for nj in range(x_all.shape[1]):
|
||
conf = v_all[:, nj:nj+1]
|
||
if (conf>0.).sum() < 3:
|
||
continue
|
||
A = conf * A0
|
||
b = conf * (x_all[:, nj:nj+1])
|
||
est = np.linalg.inv(A.T @ A) @ A.T @ b
|
||
axt[:, nj] = est[:, 0]
|
||
x_all_smoothed = xrange * axt[0:1] + axt[1:]
|
||
x_pred = x_all_smoothed[x_all.shape[0]//2]
|
||
dx_pred = axt[0]
|
||
else:
|
||
x_pred = x_all[x_all.shape[0]//2]
|
||
dx_pred = dx_mean
|
||
if x_all.shape[0] == 1:
|
||
current = FilterResult(x, dx, v, t)
|
||
self.records.append(current)
|
||
self.result = current
|
||
else:
|
||
# dx_hat = (dx * v + dx_mean * v_mean)/(v+v_mean+1e-5)
|
||
# x_pred = x_mean + dx_hat
|
||
# current = FilterResult(x_pred, dx_hat, v, t)
|
||
current = FilterResult(x, dx, v, t)
|
||
self.records.append(current)
|
||
# 使用平均速度模型
|
||
self.result = FilterResult(x_pred, dx_pred, v_sum, t)
|
||
return self.fill(self.result.x, self.result.v)
|
||
|
||
def __call__2(self, value, conf):
|
||
self.counter += 1
|
||
x, v = value[0], conf[0]
|
||
if self.result is None:
|
||
result = FilterResult(x, np.zeros_like(x), v, (v>0)*self.counter)
|
||
self.result = result
|
||
return self.fill(result.x, result.v)
|
||
prev = self.result
|
||
t = prev.t.copy()
|
||
t[v>0.] = self.counter
|
||
# update t
|
||
# dx = (x - prev.x)/(np.maximum(t-prev.t, 1))
|
||
dx = x - prev.x # 这里直接使用与之前的结果的差了，避免多帧不可见，然后速度过大
|
||
dx_ = dx.copy()
|
||
# 判断dx的大小
|
||
large_dx = np.abs(dx) > 0.5
|
||
if large_dx.sum() > 0:
|
||
v[large_dx] = 0.
|
||
t[large_dx] = prev.t[large_dx]
|
||
dx[large_dx] = 0.
|
||
missing_index = ((prev.v > 0.1) & (v < 0.1)) | (t - prev.t > 10)
|
||
if missing_index.sum() > 0:
|
||
print('missing', missing_index)
|
||
new_index = (prev.v < 0.1) & (v > 0.1)
|
||
if new_index.sum() > 0:
|
||
print('new', new_index)
|
||
dx[new_index] = 0.
|
||
weight_dx = v/(1e-5+ 3*prev.v + 1*v)
|
||
weight_x = v/(1e-5+ 3*prev.v + 1*v)
|
||
# 移除速度过大的点
|
||
dx_hat = exponential_smoothing(weight_dx, dx, prev.dx)
|
||
x_pred = prev.x + dx_hat
|
||
x_hat = exponential_smoothing(weight_x, x, x_pred)
|
||
dx_real = x_hat - prev.x
|
||
# consider the unvisible v
|
||
print_val = {
|
||
't_pre': prev.t,
|
||
'x_inp': x,
|
||
'x_pre': prev.x,
|
||
'x_new': x_hat,
|
||
'dx_inp': dx_,
|
||
'dx_pre': prev.dx,
|
||
'dx_new': dx_hat,
|
||
'dx_real': dx_real,
|
||
'v': v,
|
||
'v_pre': prev.v,
|
||
'w_vel': weight_dx,
|
||
'w_x': weight_x
|
||
}
|
||
for key in print_val.keys():
|
||
print('{:7s}'.format(key), end=' ')
|
||
print('')
|
||
for i in range(x.shape[0]):
|
||
for key in print_val.keys():
|
||
print('{:7.2f}'.format(print_val[key][i]), end=' ')
|
||
print('')
|
||
v[missing_index] = prev.v[missing_index] / 1.2 # 衰减系数
|
||
result = FilterResult(x_hat, dx_hat, v, t)
|
||
self.result = result
|
||
return self.fill(result.x, result.v)
|
||
if len(self.records) < 10:
|
||
self.records.append([self.counter, value, conf])
|
||
return self.fill(value[0], conf[0])
|
||
if self.x is None:
|
||
time = np.vstack([x[0] for x in self.records])
|
||
value_pre = np.vstack([x[1] for x in self.records])
|
||
conf_pre = np.vstack([x[2] for x in self.records])
|
||
conf_sum = conf_pre.sum(axis=0)
|
||
value_mean = (value_pre * conf_pre).sum(axis=0)/(conf_sum + 1e-5)
|
||
self.x = value_mean
|
||
self.x_conf = conf_sum
|
||
t_prev = np.zeros_like(self.x, dtype=int) - 1
|
||
t_prev[conf_sum>0] = self.counter
|
||
self.t_prev = t_prev
|
||
# 零速度初始化
|
||
self.d_x = np.zeros_like(self.x)
|
||
return self.fill(self.x, self.x_conf)
|
||
# 假设每帧都传进来的吧
|
||
return self.fill(self.x, self.x_conf)
|
||
x_est, v_est, conf_est = self.x.copy(), self.d_x.copy(), self.x_conf.copy()
|
||
value = value[0]
|
||
conf = conf[0]
|
||
d_x = value - self.x
|
||
t_current = np.zeros_like(self.x, dtype=int) - 1
|
||
t_current[conf>0.] = self.counter
|
||
t_est = t_current - self.t_prev
|
||
# 前一帧有观测，当前帧有观测，两帧之差在10帧以内。正常更新
|
||
flag_vv = (t_current > 0) & (self.t_prev > 0) & \
|
||
(t_current - self.t_prev < 10)
|
||
# 前一帧无观测；当前帧有观测的；判断为新增的
|
||
flag_iv = (self.t_prev < 0) & (t_current > 0)
|
||
weight_vel = smoothing_factor(t_est, self.d_cutoff)
|
||
# 将观测的速度权重置0
|
||
weight_vel[flag_vv] = 0.
|
||
vel_hat = exponential_smoothing(weight_vel, d_x, self.d_x)
|
||
cutoff = self.min_cutoff + self.beta * np.abs(vel_hat)
|
||
weight_value = smoothing_factor(t_est, cutoff)
|
||
# 将观测的数值权重置0
|
||
weight_value[flag_vv] = 0.
|
||
weight_value[flag_iv] = 1. # 当前帧可见的，之前的帧不可见的，直接选择当前帧
|
||
vel_hat[flag_iv] = 0.
|
||
x_hat = exponential_smoothing(weight_value, value, self.x)
|
||
flag_vi = (self.t_prev > 0) & (~flag_vv)
|
||
flag_v = flag_vv | flag_vi | flag_iv
|
||
# 前一帧有观测；当前帧无观测的；判断为丢失的
|
||
x_est[flag_v] = x_hat[flag_v]
|
||
v_est[flag_v] = vel_hat[flag_v]
|
||
conf_est[flag_v] = (self.x_conf + conf)[flag_v]/2
|
||
self.t_prev[flag_v] = self.counter
|
||
self.x = x_est
|
||
self.d_x = v_est
|
||
self.x_conf = conf_est
|
||
return self.fill(x_est, conf_est)
|
||
|
||
class IKBody(BaseBody):
|
||
def __init__(self, key, cfg_triangulator, cfg_model, cfg, debug) -> None:
|
||
super().__init__(cfg_triangulator, cfg_model, cfg)
|
||
self.key = key
|
||
self.frame_index = 0
|
||
self.frame_latest = 0
|
||
self.init = False
|
||
self.records = []
|
||
self.results = []
|
||
self.blank_result = self.make_blank()
|
||
self.fill_result = self.make_fill()
|
||
self.results_newest = self.blank_result
|
||
if True:
|
||
self.lefthand = ManoFitterCPPCache('LEFT')
|
||
self.righthand = ManoFitterCPPCache('RIGHT')
|
||
self.up_vector = 'z'
|
||
self.filter = {}
|
||
for key in ['Rh', 'poses', 'handl', 'handr']:
|
||
self.filter[key] = MyFilter(key, self.fill_result[key])
|
||
|
||
def make_blank(self):
|
||
raise NotImplementedError
|
||
|
||
def make_fill(self):
|
||
raise NotImplementedError
|
||
|
||
def smooth_results(self, params=None):
|
||
results = {'id': 0, 'type': 'smplh_half'}
|
||
for key in ['Rh', 'poses', 'handl', 'handr']:
|
||
value = self.filter[key](params[key], params[key+'_conf'])
|
||
results[key] = value
|
||
for key in ['shapes', 'Th']:
|
||
results[key] = self.blank_result[key]
|
||
return results
|
||
|
||
def smooth_results_old(self, params=None):
|
||
if params is not None:
|
||
self.results.append(params)
|
||
if len(self.results) < 10:
|
||
return params
|
||
else:
|
||
if len(self.results) < 10:
|
||
return self.fill_result
|
||
else:
|
||
params = self.fill_result
|
||
results = {'id': 0}
|
||
if False:
|
||
for key in ['Rh', 'poses', 'handl', 'handr']:
|
||
if not self.filter[key].init:
|
||
import ipdb;ipdb.set_trace()
|
||
else:
|
||
value = self.filter[key](self.results[-1][key])
|
||
if True:
|
||
for key in ['Rh', 'poses', 'handl', 'handr']:
|
||
find = False
|
||
for WINDOW in [10, 20, 40]:
|
||
if WINDOW > len(self.results):
|
||
break
|
||
records = self.results[-WINDOW:]
|
||
value = np.vstack([r[key] for r in records])
|
||
conf = np.vstack([r[key+'_conf'] for r in records])
|
||
valid = conf[..., 0] > 0
|
||
if valid.sum() < WINDOW // 3:
|
||
import ipdb;ipdb.set_trace()
|
||
else:
|
||
value, conf = value[valid], conf[valid]
|
||
mean_value = value.mean(axis=0)
|
||
std_value = value.std(axis=0)
|
||
valid2 = (np.abs(value - mean_value) < WINDOW//3 * std_value).any(axis=-1)
|
||
if valid2.sum() < WINDOW // 4:
|
||
continue
|
||
find = True
|
||
value, conf = value[valid2], conf[valid2]
|
||
conf_sum = conf.sum(axis=0)
|
||
mean = (value*conf).sum(axis=0)/(conf_sum + 1e-5)
|
||
# 计算latest
|
||
break
|
||
if key in ['poses', 'handl', 'handr']:
|
||
conf_sum_p = conf.sum(axis=0)
|
||
mean_previous = (value*conf).sum(axis=0)/(conf_sum_p + 1e-5)
|
||
mean[conf_sum<0.01] = mean_previous[conf_sum<0.01]
|
||
conf_sum[conf_sum<0.01] = conf_sum_p[conf_sum<0.01]
|
||
# 使用fill的填值
|
||
mean[conf_sum<0.01] = self.fill_result[key][0][conf_sum<0.01]
|
||
break
|
||
if find:
|
||
results[key] = mean[None]
|
||
else:
|
||
results[key] = self.fill_result[key]
|
||
if False: # 均值滤波
|
||
for key in ['Rh', 'poses', 'handl', 'handr']:
|
||
if key not in params.keys():
|
||
continue
|
||
if key not in self.cfg.SMOOTH_SIZE.keys():
|
||
results[key] = params[key]
|
||
records = self.results[-self.cfg.SMOOTH_SIZE[key]:]
|
||
value = np.vstack([r[key] for r in records])
|
||
conf = np.vstack([r[key+'_conf'] for r in records])
|
||
conf_sum = conf.sum(axis=0)
|
||
mean = (value*conf).sum(axis=0)/(conf_sum + 1e-5)
|
||
# 计算latest
|
||
if key in ['poses', 'handl', 'handr']:
|
||
records = self.results[-5*self.cfg.SMOOTH_SIZE[key]:]
|
||
value = np.vstack([r[key] for r in records])
|
||
conf = np.vstack([r[key+'_conf'] for r in records])
|
||
conf_sum_p = conf.sum(axis=0)
|
||
mean_previous = (value*conf).sum(axis=0)/(conf_sum_p + 1e-5)
|
||
mean[conf_sum<0.01] = mean_previous[conf_sum<0.01]
|
||
conf_sum[conf_sum<0.01] = conf_sum_p[conf_sum<0.01]
|
||
# 使用fill的填值
|
||
mean[conf_sum<0.01] = self.fill_result[key][0][conf_sum<0.01]
|
||
results[key] = mean[None]
|
||
results['Th'] = self.blank_result['Th']
|
||
results['shapes'] = self.blank_result['shapes']
|
||
return results
|
||
|
||
def get_keypoints3d(self, records, key=None):
|
||
if key is None:
|
||
return np.stack([r[self.key] for r in records])
|
||
else:
|
||
return np.stack([r[key] for r in records])
|
||
|
||
def check_keypoints(self, keypoints3d):
|
||
flag = (keypoints3d[..., -1]>self.cfg.MIN_THRES).sum() > 5
|
||
if len(self.records) > 1:
|
||
pre = self.records[-1]
|
||
k_pre = self.get_keypoints3d([pre])
|
||
dist = np.linalg.norm(keypoints3d[..., :3] - k_pre[..., :3], axis=-1)
|
||
conf = np.sqrt(keypoints3d[..., 3] * k_pre[..., 3])
|
||
dist_mean = (dist * conf).sum()/conf.sum()
|
||
flag = flag and dist_mean < 0.1
|
||
return flag
|
||
|
||
def __call__(self, data):
|
||
self.frame_index += 1
|
||
k3d = self.triangulator(data)[0]
|
||
keypoints3d = self.get_keypoints3d([k3d])
|
||
flag = self.check_keypoints(keypoints3d)
|
||
if not flag:
|
||
mywarn('Missing keypoints {} [{}->{}]'.format(keypoints3d[..., -1].sum(), self.frame_latest, self.frame_index))
|
||
# 1. 初始化过了，但是超出帧数了，清零
|
||
# 2. 没有初始化过，超出了，清零
|
||
if (self.frame_index - self.frame_latest > 10 and self.init) or not self.init:
|
||
mywarn('Missing keypoints, resetting...')
|
||
self.init = False
|
||
self.records = []
|
||
self.results = []
|
||
return [self.fill_result]
|
||
else:
|
||
return [self.smooth_results()]
|
||
elif not self.init: # 暂时还没有初始化，先等待
|
||
if len(self.records) < 10:
|
||
self.records.append(k3d)
|
||
return [self.fill_result]
|
||
self.records.append(k3d)
|
||
flag, params = self.fitting(keypoints3d, self.results_newest)
|
||
if not flag:
|
||
return [self.fill_result]
|
||
self.frame_latest = self.frame_index
|
||
# smooth results
|
||
results = self.smooth_results(params)
|
||
self.results_newest = results
|
||
k3d['type'] = 'body25'
|
||
return [results, k3d]
|
||
|
||
class HalfBodyIK(IKBody):
|
||
def get_keypoints3d(self, records):
|
||
THRES_WRIST = 0.2
|
||
keypoints3d = super().get_keypoints3d(records)
|
||
keypoints3d = keypoints3d[:, INDEX_HALF]
|
||
handl = super().get_keypoints3d(records, key='handl3d')
|
||
handr = super().get_keypoints3d(records, key='handr3d')
|
||
dist_ll = np.linalg.norm(keypoints3d[:, 7, :3] - handl[:, 0, :3], axis=-1)
|
||
dist_rr = np.linalg.norm(keypoints3d[:, 4, :3] - handr[:, 0, :3], axis=-1)
|
||
log('Dist left = {}, right = {}'.format(dist_ll, dist_rr))
|
||
handl[dist_ll>THRES_WRIST] = 0.
|
||
handr[dist_rr>THRES_WRIST] = 0.
|
||
keypoints3d = np.hstack([keypoints3d, handl, handr])
|
||
conf = keypoints3d[..., 3:]
|
||
keypoints3d = np.hstack([(keypoints3d[..., :3] * conf).sum(axis=0)/(1e-5 + conf.sum(axis=0)), conf.min(axis=0)])
|
||
keypoints3d = keypoints3d[None]
|
||
# if (keypoints3d.shape[0] == 10):
|
||
return keypoints3d
|
||
|
||
def _ik_shoulder(self, keypoints3d, params):
|
||
SHOULDER_IDX = [2, 5]
|
||
shoulder = keypoints3d[SHOULDER_IDX[1], :3] - keypoints3d[SHOULDER_IDX[0], :3]
|
||
if self.up_vector == 'x':
|
||
shoulder[..., 0] = 0.
|
||
up_vector = np.array([1., 0., 0.], dtype=np.float32)
|
||
elif self.up_vector == 'z':
|
||
shoulder[..., 2] = 0.
|
||
up_vector = np.array([0., 0., 1.], dtype=np.float32)
|
||
shoulder = shoulder/np.linalg.norm(shoulder, keepdims=True)
|
||
# 限定一下角度范围
|
||
theta = -np.rad2deg(np.arctan2(shoulder[1], shoulder[2]))
|
||
if (theta < 30 or theta > 150) and False:
|
||
return False, params
|
||
front = np.cross(shoulder, up_vector)
|
||
front = front/np.linalg.norm(front, keepdims=True)
|
||
R = np.stack([shoulder, up_vector, front]).T
|
||
Rh = cv2.Rodrigues(R)[0].reshape(1, 3)
|
||
log('Shoulder:{}'.format(Rh))
|
||
params['R'] = R
|
||
params['Rh'] = Rh
|
||
params['Rh_conf'] = np.zeros((1, 3)) + keypoints3d[SHOULDER_IDX, 3].min()
|
||
return True, params
|
||
|
||
def _ik_head(self, keypoints3d, params):
|
||
HEAD_IDX = [0, 8, 9, 10, 11]
|
||
HEAD_ROT_IDX = 0
|
||
est_points = keypoints3d[HEAD_IDX, :3]
|
||
valid = (keypoints3d[HEAD_IDX[0], 3] > self.cfg.MIN_THRES) and (keypoints3d[HEAD_IDX[1:], 3]>self.cfg.MIN_THRES).sum()>=2
|
||
if not valid:
|
||
params['poses_conf'][:, 3*HEAD_ROT_IDX:3*(HEAD_ROT_IDX+1)] = 0.
|
||
return params
|
||
params['poses_conf'][:, 3*HEAD_ROT_IDX:3*(HEAD_ROT_IDX+1)] = keypoints3d[HEAD_IDX, 3].sum()
|
||
|
||
gt_points = self.k_template[HEAD_IDX].numpy()
|
||
gt_points = gt_points - gt_points[:1]
|
||
est_points = est_points - est_points[:1]
|
||
# gt_points = gt_points / np.linalg.norm(gt_points, axis=-1, keepdims=True)
|
||
# est_points = est_points / np.linalg.norm(est_points, axis=-1, keepdims=True)
|
||
|
||
if True:
|
||
R_global = svd_rot(gt_points, est_points)
|
||
R_local = params['R'].T @ R_global
|
||
elif False:
|
||
est_points_inv = est_points @ params['R'].T.T
|
||
R_local = svd_rot(gt_points, est_points_inv)
|
||
else:
|
||
gt_points = gt_points @ params['R'].T
|
||
R_local = svd_rot(gt_points, est_points)
|
||
euler = rotmat2euler(R_local)
|
||
euler[0] = euler[0] - 25
|
||
# log('euler before filter: {}'.format(euler))
|
||
euler[0] = max(min(euler[0], 30), -30)
|
||
euler[1] = max(min(euler[1], 60), -60)
|
||
euler[2] = max(min(euler[2], 45), -45)
|
||
# log('euler after filter: {}'.format(euler))
|
||
R_local = euler2rotmat(euler)
|
||
R_head = cv2.Rodrigues(R_local)[0].reshape(1, 3)
|
||
params['poses'][:, 3*HEAD_ROT_IDX:3*(HEAD_ROT_IDX+1)] = R_head
|
||
return params
|
||
|
||
@staticmethod
|
||
def _rad_from_twovec(keypoints3d, start, mid, end, MIN_THRES):
|
||
start = keypoints3d[start]
|
||
mid = keypoints3d[mid]
|
||
end = keypoints3d[end]
|
||
if isinstance(end, list):
|
||
# dst is a list
|
||
if (end[:, 3] > MIN_THRES).sum() < 2:
|
||
return 0, 0.
|
||
end = np.sum(end * end[:, 3:], axis=0)/(end[:, 3:].sum())
|
||
# use its mean to represent the points
|
||
conf = [start[3], mid[3], end[3]]
|
||
valid = (min(conf) > MIN_THRES).all()
|
||
if not valid:
|
||
return 0, 0.
|
||
conf = sum(conf)
|
||
dir_src = normalize(mid[:3] - start[:3])
|
||
dir_dst = normalize(end[:3] - mid[:3])
|
||
rad = rad_from_2vec(dir_src, dir_dst)
|
||
return conf, rad
|
||
|
||
def _ik_elbow(self, keypoints3d, params):
|
||
for name, info in self.cfg.LEAF.items():
|
||
conf, rad = self._rad_from_twovec(keypoints3d, info.start, info.mid, info.end, self.cfg.MIN_THRES)
|
||
if conf <= 0.: continue
|
||
params['poses_conf'][:, 3*info['index']:3*(info['index']+1)] = conf
|
||
rad = np.clip(rad, *info['ranges'])
|
||
rot = rad*np.array(info['axis']).reshape(1, 3)
|
||
params['poses'][:, 3*info['index']:3*(info['index']+1)] = rot
|
||
return params
|
||
|
||
def _ik_arm(self, keypoints3d, params):
|
||
# forward一遍获得关键点
|
||
# 这里需要确保求解的关节点没有父节点了
|
||
template = self.body_model.keypoints({'poses': params['poses'], 'shapes': params['shapes']}, return_tensor=False)[0]
|
||
for name, info in self.cfg.NODE.items():
|
||
idx = info['children']
|
||
conf = keypoints3d[info['children'], 3]
|
||
if not (conf >self.cfg.MIN_THRES).all():
|
||
continue
|
||
params['poses_conf'][:, 3*info['index']:3*(info['index']+1)] = conf.sum()
|
||
est_points = keypoints3d[idx, :3]
|
||
gt_points = template[idx]
|
||
est_points = est_points - est_points[:1]
|
||
gt_points = gt_points - gt_points[:1]
|
||
R_children = svd_rot(gt_points, est_points)
|
||
R_local = params['R'].T @ R_children
|
||
euler = rotmat2euler(R_local)
|
||
# log('euler {} before filter: {}'.format(name, euler))
|
||
# euler[0] = max(min(euler[0], 90), -90)
|
||
# euler[1] = max(min(euler[1], 90), -90)
|
||
# euler[2] = max(min(euler[2], 90), -90)
|
||
# log('euler {} after filter: {}'.format(name, euler))
|
||
R_local = euler2rotmat(euler)
|
||
params['poses'][:, 3*info['index']:3*(info['index']+1)] = cv2.Rodrigues(R_local)[0].reshape(1, 3)
|
||
return params
|
||
|
||
def _ik_palm(self, keypoints3d, params):
|
||
# template = self.body_model.keypoints({'poses': params['poses'], 'shapes': params['shapes']}, return_tensor=False)[0]
|
||
T_joints, _ = self.body_model.transform({'poses': params['poses'], 'shapes': params['shapes']}, return_vertices=False)
|
||
T_joints = T_joints[0].cpu().numpy()
|
||
for name, info in self.cfg.PALM.items():
|
||
# 计算手掌的朝向
|
||
est_points = keypoints3d[:, :3]
|
||
est_conf = keypoints3d[:, 3]
|
||
if est_conf[info.children].min() < self.cfg.MIN_THRES:
|
||
continue
|
||
# 计算朝向
|
||
dir0 = normalize(est_points[info.children[1]] - est_points[info.children[0]])
|
||
dir1 = normalize(est_points[info.children[-1]] - est_points[info.children[0]])
|
||
normal = normalize(np.cross(dir0, dir1))
|
||
dir_parent = normalize(est_points[info.parent[1]] - est_points[info.parent[0]])
|
||
# 计算夹角
|
||
rad = np.arccos((normal * dir_parent).sum()) - np.pi/2
|
||
rad = np.clip(rad, *info['ranges'])
|
||
rot = rad*np.array(info['axis']).reshape(1, 3)
|
||
params['poses'][:, 3*info['index']:3*(info['index']+1)] = rot
|
||
# 考虑手肘的朝向；这个时候还差一个绕手肘的朝向的方向的旋转；这个旋转是在手肘扭曲之前的
|
||
# 先计算出这个朝向；再转化
|
||
R_parents = params['R'] @ T_joints[info.index, :3, :3]
|
||
normal_canonical = R_parents.T @ normal.reshape(3, 1)
|
||
normal_canonical[0, 0] = 0
|
||
normal_canonical = normalize(normal_canonical)
|
||
# 在canonical下的投影
|
||
# normal_T = np.array([0., -1., 0.])
|
||
# trick: 旋转角度的正弦值等于在z轴上的坐标
|
||
rad = np.arcsin(normal_canonical[2, 0])
|
||
rot_x = np.array([-rad, 0., 0.])
|
||
R_x = cv2.Rodrigues(rot_x)[0]
|
||
R_elbow = cv2.Rodrigues(params['poses'][:, 3*info.index_elbow:3*info.index_elbow+3])[0]
|
||
R_elbow = R_elbow @ R_x
|
||
params['poses'][:, 3*info.index_elbow:3*info.index_elbow+3] = cv2.Rodrigues(R_elbow)[0].reshape(1, 3)
|
||
return params
|
||
|
||
def _ik_hand(self, template, keypoints3d, poses, conf, is_left):
|
||
# 计算手的每一段的置信度
|
||
poses_full = np.zeros((1, 45))
|
||
conf_full = np.zeros((1, 45))
|
||
y_axis = np.array([0., 1., 0.])
|
||
log('_ik for left: {}'.format(is_left))
|
||
for name, info in self.cfg.HAND.LEAF.items():
|
||
conf, rad = self._rad_from_twovec(keypoints3d, *info.ranges, self.cfg.MIN_THRES)
|
||
if conf <= 0.:
|
||
log('- skip: {}'.format(name))
|
||
continue
|
||
# trick: 手的朝向是反的
|
||
rad = - rad
|
||
if info.axis == 'auto':
|
||
template_dir = template[info.ranges[2]] - template[info.ranges[1]]
|
||
# y轴方向设成0
|
||
template_dir[1] = 0.
|
||
template_dir = normalize(template_dir)
|
||
# 计算旋转轴，在与z轴的cross方向上
|
||
rot_vec = normalize(np.cross(template_dir, y_axis)).reshape(1, 3)
|
||
elif info.axis == 'root':
|
||
template_dir0 = template[info.ranges[1]] - template[info.ranges[0]]
|
||
template_dir1 = template[info.ranges[2]] - template[info.ranges[1]]
|
||
template_dir0 = normalize(template_dir0)
|
||
template_dir1 = normalize(template_dir1)
|
||
costheta0 = (template_dir0 *template_dir1).sum()
|
||
# 计算当前的夹角
|
||
est_dir0 = keypoints3d[info.ranges[1], :3] - keypoints3d[info.ranges[0], :3]
|
||
est_dir1 = keypoints3d[info.ranges[2], :3] - keypoints3d[info.ranges[1], :3]
|
||
est_dir0 = normalize(est_dir0)
|
||
est_dir1 = normalize(est_dir1)
|
||
costheta1 = (est_dir0 * est_dir1).sum()
|
||
# trick: 手的旋转角度都是相反的
|
||
rad = - np.arccos(np.clip(costheta1/costheta0, 0., 1.))
|
||
rot_vec = normalize(np.cross(template_dir1, y_axis)).reshape(1, 3)
|
||
log('- get: {}: {:.1f}, {}'.format(name, np.rad2deg(rad), rot_vec))
|
||
poses_full[:, 3*info.index:3*info.index+3] = rot_vec * rad
|
||
conf_full[:, 3*info.index:3*info.index+3] = conf
|
||
# 求解
|
||
usePCA = False
|
||
if usePCA:
|
||
ncomp = 24
|
||
lamb = 0.05
|
||
if is_left:
|
||
A_full = self.body_model.components_full_l[:ncomp].T
|
||
mean_full = self.body_model.mean_full_l
|
||
else:
|
||
A_full = self.body_model.components_full_r[:ncomp].T
|
||
mean_full = self.body_model.mean_full_r
|
||
valid = conf_full[0] > 0.
|
||
A = A_full[valid, :]
|
||
res = (poses_full[:, valid] - mean_full[:, valid]).T
|
||
x = np.linalg.inv(A.T @ A + lamb * np.eye(ncomp)) @ A.T @ res
|
||
poses_full = (A_full @ x).reshape(1, -1) + mean_full
|
||
conf_full = np.zeros_like(poses_full) + valid.sum()
|
||
return poses_full, conf_full
|
||
|
||
def make_blank(self):
|
||
params = self.body_model.init_params(1, ret_tensor=False)
|
||
params['id'] = 0
|
||
params['type'] = 'smplh_half'
|
||
params['Th'][0, 0] = 1.
|
||
params['Th'][0, 1] = -1
|
||
params['Th'][0, 2] = 1.
|
||
return params
|
||
|
||
def make_fill(self):
|
||
params = self.body_model.init_params(1, ret_tensor=False)
|
||
params['id'] = 0
|
||
params['type'] = 'smplh_half'
|
||
params['Rh'][0, 2] = -np.pi/2
|
||
params['handl'] = self.body_model.mean_full_l
|
||
params['handr'] = self.body_model.mean_full_r
|
||
params['poses'][0, 3*4+2] = -np.pi/4
|
||
params['poses'][0, 3*5+2] = np.pi/4
|
||
return params
|
||
|
||
def fitting(self, keypoints3d, results_pre):
|
||
# keypoints3d: (nFrames, nJoints, 4)
|
||
# 根据肩膀计算身体朝向
|
||
if len(keypoints3d.shape) == 3:
|
||
keypoints3d = keypoints3d[0]
|
||
params = self.body_model.init_params(1, ret_tensor=False)
|
||
params['poses_conf'] = np.zeros_like(params['poses'])
|
||
params['handl_conf'] = np.zeros_like(params['handl'])
|
||
params['handr_conf'] = np.zeros_like(params['handr'])
|
||
params['Rh_conf'] = 0.
|
||
params['id'] = 0
|
||
flag, params = self._ik_shoulder(keypoints3d, params)
|
||
if (params['Rh_conf'] <= 0.01).all():
|
||
return False, params
|
||
params = self._ik_head(keypoints3d, params)
|
||
params = self._ik_elbow(keypoints3d, params)
|
||
params = self._ik_arm(keypoints3d, params)
|
||
params = self._ik_palm(keypoints3d, params)
|
||
if False:
|
||
params['handl'], params['handl_conf'] = self._ik_hand(self.k_template[12:12+21].numpy(), keypoints3d[12:12+21], params['handl'], params['handl_conf'], is_left=True)
|
||
params['handr'], params['handr_conf'] = self._ik_hand(self.k_template[12+21:12+21+21].numpy(), keypoints3d[12+21:12+21+21], params['handr'], params['handr_conf'], is_left=False)
|
||
else:
|
||
params_l = self.lefthand(keypoints3d[12:12+21])[0]
|
||
params_r = self.righthand(keypoints3d[12+21:12+21+21])[0]
|
||
# log('[{:06d}] {}'.format(self.frame_index, params_l['poses'][0]))
|
||
# log('[{:06d}] {}'.format(self.frame_index, params_r['poses'][0]))
|
||
ncomp = params_l['poses'].shape[1]
|
||
A_full = self.body_model.components_full_l[:ncomp].T
|
||
mean_full = self.body_model.mean_full_l
|
||
poses_full = (A_full @ params_l['poses'].T).T + mean_full
|
||
params['handl'] = poses_full
|
||
A_full = self.body_model.components_full_r[:ncomp].T
|
||
mean_full = self.body_model.mean_full_r
|
||
poses_full = (A_full @ params_r['poses'].T).T + mean_full
|
||
params['handr'] = poses_full
|
||
params['handl_conf'] = np.ones((1, 45))
|
||
params['handr_conf'] = np.ones((1, 45))
|
||
return True, params
|
||
|
||
class BaseFitter(BaseBody):
|
||
def __init__(self, cfg_triangulator, cfg_model,
|
||
INIT_SIZE, WINDOW_SIZE, FITTING_SIZE, SMOOTH_SIZE,
|
||
init_dict, fix_dict,
|
||
cfg) -> None:
|
||
super().__init__(cfg_triangulator, cfg_model, cfg)
|
||
self.records = []
|
||
self.results = []
|
||
self.INIT_SIZE = INIT_SIZE
|
||
self.WINDOW_SIZE = WINDOW_SIZE
|
||
self.FITTING_SIZE = FITTING_SIZE
|
||
self.SMOOTH_SIZE = SMOOTH_SIZE
|
||
self.time = 0
|
||
self.frame_latest = 0
|
||
self.frame_index = 0
|
||
self.init = False
|
||
self.init_dict = init_dict
|
||
self.fix_dict = fix_dict
|
||
self.identity_cache = {}
|
||
|
||
def get_keypoints3d(self, records):
|
||
raise NotImplementedError
|
||
|
||
def get_init_params(self, nFrames):
|
||
params = self.body_model.init_params(nFrames, ret_tensor=True)
|
||
for key, val in self.init_dict.items():
|
||
if key == 'Rh':
|
||
import cv2
|
||
R = cv2.Rodrigues(params['Rh'][0].cpu().numpy())[0]
|
||
for vec in self.init_dict['Rh']:
|
||
Rrel = cv2.Rodrigues(np.deg2rad(np.array(vec)))[0]
|
||
R = Rrel @ R
|
||
Rh = cv2.Rodrigues(R)[0]
|
||
params['Rh'] = torch.Tensor(Rh).reshape(-1, 3).repeat(nFrames, 1)
|
||
else:
|
||
params[key] = torch.Tensor(val).repeat(nFrames, 1)
|
||
params['id'] = 0
|
||
return params
|
||
|
||
def add_any_reg(self, val, val0, JTJ, JTr, w):
|
||
# loss: (val - val0)
|
||
nVals = val.shape[0]
|
||
if nVals not in self.identity_cache.keys():
|
||
self.identity_cache[nVals] = torch.eye(nVals, device=val.device, dtype=val.dtype)
|
||
identity = self.identity_cache[nVals]
|
||
JTJ += w * identity
|
||
JTr += -w*(val - val0).view(-1, 1)
|
||
return 0
|
||
|
||
def log(self, name, step, delta, res, keys_range=None):
|
||
toc = (time() - self.time)*1000
|
||
norm_d = torch.norm(delta).item()
|
||
norm_f = torch.norm(res).item()
|
||
text = '[{}:{:6.2f}]: step = {:3d}, ||delta|| = {:.4f}, ||res|| = {:.4f}'.format(name, toc, step, norm_d, norm_f)
|
||
print(text)
|
||
self.time = time()
|
||
return norm_d, norm_f
|
||
|
||
def fitShape(self, keypoints3d, params, weight, option):
|
||
kintree = np.array(self.cfg.shape.kintree)
|
||
nShapes = params['shapes'].shape[-1]
|
||
# shapes: (1, 10)
|
||
shapes = params['shapes'].T
|
||
shapes0 = shapes.clone()
|
||
device, dtype = shapes.device, shapes.dtype
|
||
lengths3d_est = torch.norm(keypoints3d[:, kintree[:, 1], :3] - keypoints3d[:, kintree[:, 0], :3], dim=-1)
|
||
conf = torch.sqrt(keypoints3d[:, kintree[:, 1], 3:] * keypoints3d[:, kintree[:, 0], 3:])
|
||
conf = conf.repeat(1, 1, 3).reshape(-1, 1)
|
||
nFrames = keypoints3d.shape[0]
|
||
# jacobian: (nFrames, nLimbs, 3, nShapes)
|
||
jacob_limb_shapes = self.jacobian_limb_shapes[None].repeat(nFrames, 1, 1, 1)
|
||
jacob_limb_shapes = jacob_limb_shapes.reshape(-1, nShapes)
|
||
# 注意：这里乘到雅克比的应该是 sqrt(conf)，这里把两个合并了
|
||
JTJ_limb_shapes = jacob_limb_shapes.t() @ (jacob_limb_shapes * conf)
|
||
lossnorm = 0
|
||
self.time = time()
|
||
for iter_ in range(option.max_iter):
|
||
# perform shape blending
|
||
shape_offset = self.k_shapeBlend @ shapes
|
||
keyShaped = self.k_template + shape_offset[..., 0]
|
||
JTJ = JTJ_limb_shapes
|
||
JTr = torch.zeros((nShapes, 1), device=device, dtype=dtype)
|
||
# 并行添加所有的骨架
|
||
dir = keyShaped[kintree[:, 1]] - keyShaped[kintree[:, 0]]
|
||
dir_normalized = dir / torch.norm(dir, dim=-1, keepdim=True)
|
||
# res: (nFrames, nLimbs, 3)
|
||
res = lengths3d_est[..., None] * dir_normalized[None] - dir[None]
|
||
res = conf * res.reshape(-1, 1)
|
||
JTr += jacob_limb_shapes.t() @ res
|
||
self.add_any_reg(shapes, shapes0, JTJ, JTr, w=weight.init_shape)
|
||
delta = torch.linalg.solve(JTJ, JTr)
|
||
shapes += delta
|
||
norm_d, norm_f = self.log('shape', iter_, delta, res)
|
||
if torch.norm(delta) < option.gtol:
|
||
break
|
||
if iter_ > 0 and abs(norm_f - lossnorm)/norm_f < option.ftol:
|
||
break
|
||
lossnorm = norm_f
|
||
shapes = shapes.t()
|
||
params['shapes'] = shapes
|
||
return params, keyShaped
|
||
|
||
def fitRT(self, keypoints3d, params, weight, option, kpts_index=None):
|
||
keys_optimized = ['Rh', 'Th']
|
||
if kpts_index is not None:
|
||
keypoints3d = keypoints3d[:, kpts_index]
|
||
init_dict = {
|
||
'Rh': params['Rh'].clone(),
|
||
'Th': params['Th'].clone(),
|
||
}
|
||
init_dict['Rot'] = batch_rodrigues(init_dict['Rh'])
|
||
params_dict = {
|
||
'Rh': params['Rh'],
|
||
'Th': params['Th'],
|
||
}
|
||
keys_range = {}
|
||
for ikey, key in enumerate(keys_optimized):
|
||
if ikey == 0:
|
||
keys_range[key] = [0, init_dict[key].shape[-1]]
|
||
else:
|
||
start = keys_range[keys_optimized[ikey-1]][1]
|
||
keys_range[key] = [start, start+init_dict[key].shape[-1]]
|
||
NUM_FRAME = keys_range[keys_optimized[-1]][1]
|
||
kpts = self.body_model.keypoints({'poses': params['poses'], 'shapes': params['shapes']})
|
||
bn = keypoints3d.shape[0]
|
||
conf = keypoints3d[..., -1:].repeat(1, 1, 3).reshape(bn, -1)
|
||
dtype, device = kpts.dtype, kpts.device
|
||
w_joints = 1./keypoints3d.shape[-2] * weight.joints
|
||
self.time = time()
|
||
for iter_ in range(option.max_iter):
|
||
Rh, Th = params_dict['Rh'], params_dict['Th']
|
||
rot, jacobi_R_rvec, jacobi_joints_RT = getJacobianOfRT(Rh, Th, kpts)
|
||
kpts_rt = torch.matmul(kpts, rot.transpose(-1, -2)) + Th[:, None]
|
||
# // loss: J_obs - (R @ jest + T) => -dR/drvec - dT/dtvec - Rx(djest/dtheta)
|
||
jacobi_keypoints = -jacobi_joints_RT
|
||
if kpts_index is not None:
|
||
jacobi_keypoints = jacobi_keypoints[:, kpts_index]
|
||
kpts_rt = kpts_rt[:, kpts_index]
|
||
jacobi_keypoints_flat = jacobi_keypoints.view(bn, -1, jacobi_keypoints.shape[-1])
|
||
JTJ_keypoints = jacobi_keypoints_flat.transpose(-1, -2) @ (jacobi_keypoints_flat * conf[..., None])
|
||
res = conf[..., None] * ((keypoints3d[..., :3] - kpts_rt).view(bn, -1, 1))
|
||
JTr_keypoints = jacobi_keypoints_flat.transpose(-1, -2) @ res
|
||
#
|
||
JTJ = torch.eye(bn*NUM_FRAME, device=device, dtype=dtype) * 1e-4
|
||
JTr = torch.zeros((bn*NUM_FRAME, 1), device=device, dtype=dtype)
|
||
# accumulate loss
|
||
for nf in range(bn):
|
||
JTJ[nf*NUM_FRAME:(nf+1)*NUM_FRAME, nf*NUM_FRAME:(nf+1)*NUM_FRAME] += w_joints * JTJ_keypoints[nf]
|
||
# add regularization for each parameter
|
||
for nf in range(bn):
|
||
for key in keys_optimized:
|
||
start, end = nf*NUM_FRAME + keys_range[key][0], nf*NUM_FRAME + keys_range[key][1]
|
||
if key == 'Rh':
|
||
# 增加初始化的loss
|
||
res_init = rot[nf] - init_dict['Rot'][nf]
|
||
JTJ[start:end, start:end] += weight['init_'+key] * jacobi_R_rvec[nf] @ jacobi_R_rvec[nf].T
|
||
JTr[start:end] += -weight.init_Rh * jacobi_R_rvec[nf] @ res_init.reshape(-1, 1)
|
||
else:
|
||
res_init = Th[nf] - init_dict['Th'][nf]
|
||
JTJ[start:end, start:end] += weight['init_'+key] * torch.eye(3)
|
||
JTr[start:end] += -weight['init_'+key] * res_init.reshape(-1, 1)
|
||
JTr += - w_joints * JTr_keypoints.reshape(-1, 1)
|
||
# solve
|
||
delta = torch.linalg.solve(JTJ, JTr)
|
||
norm_d, norm_f = self.log('pose', iter_, delta, res)
|
||
if torch.norm(delta) < option.gtol:
|
||
break
|
||
if iter_ > 0 and abs(norm_f - lossnorm)/norm_f < option.ftol:
|
||
break
|
||
delta = delta.view(bn, NUM_FRAME)
|
||
lossnorm = norm_f
|
||
for key, _range in keys_range.items():
|
||
if key not in params_dict.keys():continue
|
||
params_dict[key] += delta[:, _range[0]:_range[1]]
|
||
norm_key = torch.norm(delta[:, _range[0]:_range[1]])
|
||
params.update(params_dict)
|
||
return params
|
||
|
||
@staticmethod
|
||
def localTransform(J_shaped, poses, rootIdx, kintree):
|
||
bn = poses.shape[0]
|
||
nThetas = poses.shape[1]//3
|
||
localTrans = torch.eye(4, device=poses.device)[None, None].repeat(bn, nThetas, 1, 1)
|
||
poses_flat = poses.reshape(-1, 3)
|
||
rot_flat = batch_rodrigues(poses_flat)
|
||
rot = rot_flat.view(bn, nThetas, 3, 3)
|
||
localTrans[:, :, :3, :3] = rot
|
||
# set the root
|
||
localTrans[:, rootIdx, :3, 3] = J_shaped[rootIdx].view(1, 3)
|
||
# relTrans: (nKintree, 3)
|
||
relTrans = J_shaped[kintree[:, 1]] - J_shaped[kintree[:, 0]]
|
||
localTrans[:, kintree[:, 1], :3, 3] = relTrans[None]
|
||
return localTrans
|
||
|
||
@staticmethod
|
||
def globalTransform(localTrans, rootIdx, kintree):
|
||
# localTrans: (bn, nJoints, 4, 4)
|
||
globalTrans = localTrans.clone()
|
||
# set the root
|
||
for (parent, child) in kintree:
|
||
globalTrans[:, child] = globalTrans[:, parent] @ localTrans[:, child]
|
||
return globalTrans
|
||
|
||
def jacobi_GlobalTrans_theta(self, poses, j_shaped, rootIdx, kintree,
|
||
device, dtype):
|
||
parents = self.parents
|
||
start = time()
|
||
tic = lambda x: print('-> [{:20s}]: {:.3f}ms'.format(x, 1000*(time()-start)))
|
||
localTrans = self.localTransform(j_shaped, poses, rootIdx, kintree)
|
||
# tic('local trans')
|
||
globalTrans = self.globalTransform(localTrans, rootIdx, kintree)
|
||
# tic('global trans')
|
||
# 计算localTransformation
|
||
poses_flat = poses.view(poses.shape[0], -1, 3)
|
||
# jacobi_R_rvec: (bn, nJ, 3, 9)
|
||
Rot, jacobi_R_rvec = batch_rodrigues_jacobi(poses_flat)
|
||
|
||
bn, nJoints = localTrans.shape[:2]
|
||
dGlobalTrans_template = torch.zeros((bn, nJoints, 4, 4), device=device, dtype=dtype)
|
||
# compute global transformation
|
||
# results: global transformation to each theta: (bn, nJ, 4, 4, nTheta)
|
||
jacobi_GlobalTrans_theta = torch.zeros((bn, nJoints, 4, 4, nJoints*3), device=device, dtype=dtype)
|
||
# tic('rodrigues')
|
||
for djIdx in range(1, nJoints):
|
||
if djIdx in self.cfg.IGNORE_JOINTS: continue
|
||
# // 第djIdx个轴角的第dAxis个维度
|
||
for dAxis in range(3):
|
||
if dAxis in self.cfg.IGNORE_AXIS.get(str(djIdx), []): continue
|
||
# if(model->frozenJoint[3*djIdx+dAxis])continue;
|
||
# // 从上至下堆叠起来
|
||
dGlobalTrans = dGlobalTrans_template.clone()
|
||
# // 将local的映射过来
|
||
dGlobalTrans[:, djIdx, :3, :3] = jacobi_R_rvec[:, djIdx, dAxis].view(bn, 3, 3)
|
||
if djIdx != rootIdx:
|
||
# // d(R0 @ R1)/dt1 = R0 @ dR1/dt1, 这里的R0是上一级所有的累积，因此使用全局的
|
||
parent = parents[djIdx]
|
||
dGlobalTrans[:, djIdx] = globalTrans[:, parent] @ dGlobalTrans[:, djIdx]
|
||
valid = np.zeros(nJoints, dtype=np.bool)
|
||
valid[djIdx] = True
|
||
# tic('current {}'.format(djIdx))
|
||
# // 遍历骨架树: 将对当前theta的导数传递下去
|
||
for (src, dst) in kintree:
|
||
# // 当前处理的关节为子节点的不用考虑
|
||
if dst == djIdx: continue
|
||
# if dst in self.cfg.IGNORE_JOINTS:continue
|
||
valid[dst] = valid[src]
|
||
if valid[src]:
|
||
# // 如果父节点是有效的: d(R0 @ R1)/dt0 = dR0/dt0 @ R1，这里的R1是当前的局部的，因此使用local的
|
||
dGlobalTrans[:, dst] = dGlobalTrans[:, src] @ localTrans[:, dst]
|
||
# tic('forward {}'.format(djIdx))
|
||
jacobi_GlobalTrans_theta[..., 3*djIdx+dAxis] = dGlobalTrans
|
||
# tic('jacobia')
|
||
return globalTrans, jacobi_GlobalTrans_theta
|
||
|
||
def fitPose(self, keypoints3d, params, weight, option, kpts_index=None):
|
||
# preprocess input data
|
||
if kpts_index is not None:
|
||
keypoints3d = keypoints3d[:, kpts_index]
|
||
bn = keypoints3d.shape[0]
|
||
conf = keypoints3d[..., -1:].repeat(1, 1, 3).reshape(bn, -1)
|
||
if (conf > 0.3).sum() < 4:
|
||
print('skip')
|
||
return params
|
||
w_joints = 1./keypoints3d.shape[-2] * weight.joints
|
||
# pre-calculate the shape
|
||
Rh, Th, poses = params['Rh'], params['Th'], params['poses']
|
||
init_dict = {
|
||
'Rh': Rh.clone(),
|
||
'Th': Th.clone(),
|
||
'poses': poses.clone()
|
||
}
|
||
init_dict['Rot'] = batch_rodrigues(init_dict['Rh'])
|
||
zero_dict = {key:torch.zeros_like(val) for key, val in init_dict.items()}
|
||
keys_optimized = ['Rh', 'Th', 'poses']
|
||
keys_range = {}
|
||
for ikey, key in enumerate(keys_optimized):
|
||
if ikey == 0:
|
||
keys_range[key] = [0, init_dict[key].shape[-1]]
|
||
else:
|
||
start = keys_range[keys_optimized[ikey-1]][1]
|
||
keys_range[key] = [start, start+init_dict[key].shape[-1]]
|
||
NUM_FRAME = keys_range[keys_optimized[-1]][1]
|
||
# calculate J_shaped
|
||
shapes_t = params['shapes'].t()
|
||
shape_offset = self.j_shapeBlend @ shapes_t
|
||
# jshaped: (nJoints, 3)
|
||
j_shaped = self.j_template + shape_offset[..., 0]
|
||
shape_offset = self.k_shapeBlend @ shapes_t
|
||
# kshaped: (nJoints, 3)
|
||
k_shaped = self.k_template + shape_offset[..., 0]
|
||
# optimize parameters
|
||
nJoints = j_shaped.shape[0]
|
||
dtype, device = poses.dtype, poses.device
|
||
lossnorm = 0
|
||
self.time = time()
|
||
for iter_ in range(option.max_iter):
|
||
# forward the model
|
||
# 0. poses => full poses
|
||
def tic(name):
|
||
print('[{:20}] {:.3f}ms'.format(name, 1000*(time()-self.time)))
|
||
if 'handl' in params.keys():
|
||
poses_full = self.body_model.extend_poses(poses, params['handl'], params['handr'])
|
||
jacobi_posesful_poses = self.body_model.jacobian_posesfull_poses_
|
||
else:
|
||
poses_full = self.body_model.extend_poses(poses)
|
||
jacobi_posesful_poses = self.body_model.jacobian_posesfull_poses(poses, poses_full)
|
||
# tic('jacobian poses full')
|
||
# 1. poses => local transformation => global transformation(bn, nJ, 4, 4)
|
||
globalTrans, jacobi_GlobalTrans_theta = self.jacobi_GlobalTrans_theta(poses_full, j_shaped, self.rootIdx, self.kintree, device, dtype)
|
||
# tic('global transform')
|
||
# 2. global transformation => relative transformation => final transformation
|
||
relTrans = globalTrans.clone()
|
||
relTrans[..., :3, 3:] -= torch.matmul(globalTrans[..., :3, :3], j_shaped[None, :, :, None])
|
||
|
||
relTrans_weight = torch.einsum('kj,fjab->fkab', self.k_weights, relTrans)
|
||
jacobi_relTrans_theta = jacobi_GlobalTrans_theta.clone()
|
||
# // consider topRight: T - RT0: add -dRT0/dt
|
||
# rot: (bn, nJoints, 3, 3, nThetas) @ (bn, nJoints, 1, 3, 1) => (bn, nJoints, 3, nThetas)
|
||
rot = jacobi_GlobalTrans_theta[..., :3, :3, :]
|
||
j0 = j_shaped.reshape(1, nJoints, 1, 3, 1).expand(bn, -1, -1, -1, -1)
|
||
jacobi_relTrans_theta[..., :3, 3, :] -= torch.sum(rot*j0, dim=-2)
|
||
jacobi_blendtrans_theta = torch.einsum('kj,fjabt->fkabt', self.k_weights, jacobi_relTrans_theta)
|
||
kpts = torch.einsum('fkab,kb->fka', relTrans_weight[..., :3, :3], k_shaped) + relTrans_weight[..., :3, 3]
|
||
# d(RJ0 + J1)/dtheta = d(R)/dtheta @ J0 + dJ1/dtheta
|
||
rot = jacobi_blendtrans_theta[..., :3, :3, :]
|
||
k0 = k_shaped.reshape(1, k_shaped.shape[0], 1, 3, 1).expand(bn, -1, -1, -1, -1)
|
||
# jacobi_keypoints_theta: (bn, nKeypoints, 3, nThetas)
|
||
jacobi_keypoints_theta = torch.sum(rot*k0, dim=-2) + jacobi_blendtrans_theta[..., :3, 3, :]
|
||
# tic('keypoints')
|
||
# // compute the jacobian of R T
|
||
# // loss: J_obs - (R @ jest + T)
|
||
rot, jacobi_R_rvec, jacobi_joints_RT = getJacobianOfRT(Rh, Th, kpts)
|
||
kpts_rt = torch.matmul(kpts, rot.transpose(-1, -2)) + Th[:, None]
|
||
rot_nk = rot[:, None].expand(-1, k_shaped.shape[0], -1, -1)
|
||
jacobi_keypoints_theta = torch.matmul(rot_nk, jacobi_keypoints_theta)
|
||
# compute jacobian of poses
|
||
shape_jacobi = jacobi_keypoints_theta.shape[:-1]
|
||
NUM_THETAS = jacobi_posesful_poses.shape[0]
|
||
jacobi_keypoints_poses = (jacobi_keypoints_theta[..., :NUM_THETAS].view(-1, NUM_THETAS) @ jacobi_posesful_poses).reshape(*shape_jacobi, -1)
|
||
# // loss: J_obs - (R @ jest + T) => -dR/drvec - dT/dtvec - Rx(djest/dtheta)
|
||
jacobi_keypoints = torch.cat([-jacobi_joints_RT, -jacobi_keypoints_poses], dim=-1)
|
||
if kpts_index is not None:
|
||
jacobi_keypoints = jacobi_keypoints[:, kpts_index]
|
||
kpts_rt = kpts_rt[:, kpts_index]
|
||
jacobi_keypoints_flat = jacobi_keypoints.view(bn, -1, jacobi_keypoints.shape[-1])
|
||
# tic('jacobian keypoints')
|
||
JTJ_keypoints = jacobi_keypoints_flat.transpose(-1, -2) @ (jacobi_keypoints_flat * conf[..., None])
|
||
res = conf[..., None] * ((keypoints3d[..., :3] - kpts_rt).view(bn, -1, 1))
|
||
JTr_keypoints = jacobi_keypoints_flat.transpose(-1, -2) @ res
|
||
cache_dict = {
|
||
'Th': Th,
|
||
'Rh': Rh,
|
||
'poses': poses,
|
||
}
|
||
# 计算loss
|
||
# JTJ = torch.zeros((bn*NUM_FRAME, bn*NUM_FRAME), device=device, dtype=dtype)
|
||
JTJ = torch.eye(bn*NUM_FRAME, device=device, dtype=dtype) * 1e-4
|
||
JTr = torch.zeros((bn*NUM_FRAME, 1), device=device, dtype=dtype)
|
||
# add regularization for each parameter
|
||
for nf in range(bn):
|
||
for key in keys_optimized:
|
||
start, end = nf*NUM_FRAME + keys_range[key][0], nf*NUM_FRAME + keys_range[key][1]
|
||
JTJ[start:end, start:end] += weight['reg_{}'.format(key)] * torch.eye(end-start)
|
||
JTr[start:end] += -weight['reg_{}'.format(key)] * cache_dict[key][nf].view(-1, 1)
|
||
# add init for Rh
|
||
if key == 'Rh':
|
||
# 增加初始化的loss
|
||
res_init = rot[nf] - init_dict['Rot'][nf]
|
||
JTJ[start:end, start:end] += weight['init_'+key] * jacobi_R_rvec[nf] @ jacobi_R_rvec[nf].T
|
||
JTr[start:end] += -weight.init_Rh * jacobi_R_rvec[nf] @ res_init.reshape(-1, 1)
|
||
else:
|
||
res_init = cache_dict[key][nf] - init_dict[key][nf]
|
||
JTJ[start:end, start:end] += weight['init_'+key] * torch.eye(end-start)
|
||
JTr[start:end] += -weight['init_'+key] * res_init.reshape(-1, 1)
|
||
# add keypoints loss
|
||
for nf in range(bn):
|
||
JTJ[nf*NUM_FRAME:(nf+1)*NUM_FRAME, nf*NUM_FRAME:(nf+1)*NUM_FRAME] += w_joints * JTJ_keypoints[nf]
|
||
JTr += - w_joints * JTr_keypoints.reshape(-1, 1)
|
||
# tic('add loss')
|
||
delta = torch.linalg.solve(JTJ, JTr)
|
||
# tic('solve')
|
||
norm_d, norm_f = self.log('pose', iter_, delta, res, keys_range)
|
||
if torch.norm(delta) < option.gtol:
|
||
break
|
||
if iter_ > 0 and abs(norm_f - lossnorm)/norm_f < option.ftol:
|
||
break
|
||
delta = delta.view(bn, NUM_FRAME)
|
||
lossnorm = norm_f
|
||
for key, _range in keys_range.items():
|
||
if key not in cache_dict.keys():continue
|
||
cache_dict[key] += delta[:, _range[0]:_range[1]]
|
||
res = {
|
||
'id': params['id'],
|
||
'poses': poses,
|
||
'shapes': params['shapes'],
|
||
'Rh': Rh,
|
||
'Th': Th,
|
||
}
|
||
for key, val in params.items():
|
||
if key not in res.keys():
|
||
res[key] = val
|
||
return res
|
||
|
||
def try_to_init(self, records):
|
||
if self.init:
|
||
return copy.deepcopy(self.params_newest)
|
||
mywarn('>> Initialize')
|
||
keypoints3d = self.get_keypoints3d(records)
|
||
params, keypoints_template = self.fitShape(keypoints3d, self.get_init_params(keypoints3d.shape[0]), self.cfg.shape.weight, self.cfg.shape.option)
|
||
Rot = batch_rodrigues(params['Rh'])
|
||
keypoints_template = torch.matmul(keypoints_template, Rot[0].t())
|
||
if self.cfg.initRT.mean_T:
|
||
conf = keypoints3d[..., 3:]
|
||
T = ((keypoints3d[..., :3] - keypoints_template[None])*conf).sum(dim=-2)/(conf.sum(dim=-2))
|
||
params['Th'] = T
|
||
params = self.fitRT(keypoints3d, params, self.cfg.initRT.weight, self.cfg.initRT.option,
|
||
kpts_index=self.cfg.TORSO_INDEX)
|
||
params = self.fitPose(keypoints3d, params, self.cfg.init_pose.weight, self.cfg.init_pose.option,
|
||
kpts_index=self.cfg.TORSO_INDEX)
|
||
params = self.fitPose(keypoints3d, params, self.cfg.init_pose.weight, self.cfg.init_pose.option,
|
||
kpts_index=self.cfg.BODY_INDEX)
|
||
mywarn('>> Initialize Rh = {}, Th = {}'.format(params['Rh'][0], params['Th'][0]))
|
||
params = Params(params)[-1]
|
||
self.init = True
|
||
return params
|
||
|
||
def fitting(self, params_init, records):
|
||
keypoints3d = self.get_keypoints3d(records[-self.WINDOW_SIZE:])
|
||
params = params_init
|
||
params = self.fitRT(keypoints3d[-self.FITTING_SIZE:], params, self.cfg.RT.weight, self.cfg.RT.option)
|
||
params = self.fitPose(keypoints3d[-self.FITTING_SIZE:], params, self.cfg.pose.weight, self.cfg.pose.option)
|
||
return params
|
||
|
||
def filter_result(self, result):
|
||
poses = result['poses'].reshape(-1, 3)
|
||
# TODO: 这里的xyz是scipy中的XYZ
|
||
euler = axis_angle_to_euler(poses, order='xyz')
|
||
# euler[euler>np.pi] = 0.
|
||
poses = euler_to_axis_angle(euler, order='xyz')
|
||
result['euler'] = euler
|
||
return result
|
||
|
||
def smooth_results(self):
|
||
results_ = {}
|
||
for key in self.results[0].keys():
|
||
if key == 'id': continue
|
||
if key not in self.SMOOTH_SIZE.keys():continue
|
||
results_[key] = np.vstack([r[key] for r in self.results[-self.SMOOTH_SIZE[key]:]])
|
||
results_[key] = np.mean(results_[key], axis=0, keepdims=True)
|
||
results_['id'] = 0
|
||
for key, val in self.fix_dict.items():
|
||
results_[key] = np.array(val)
|
||
# results_['Th'][:] = 0.
|
||
return [results_]
|
||
|
||
def check_keypoints(self, keypoints3d):
|
||
flag = (keypoints3d[..., -1]>0.3).sum() > 5
|
||
if len(self.records) > 1:
|
||
pre = self.records[-1]
|
||
k_pre = self.get_keypoints3d([pre])
|
||
dist = torch.norm(keypoints3d[..., :3] - k_pre[..., :3], dim=-1)
|
||
conf = torch.sqrt(keypoints3d[..., 3] * k_pre[..., 3])
|
||
dist_mean = (dist * conf).sum()/conf.sum()
|
||
flag = flag and dist_mean < 0.1
|
||
return flag
|
||
|
||
def __call__(self, data):
|
||
self.frame_index += 1
|
||
k3d = self.triangulator(data)[0]
|
||
keypoints3d = self.get_keypoints3d([k3d])
|
||
flag = self.check_keypoints(keypoints3d)
|
||
if not flag:
|
||
mywarn('Missing keypoints {} [{}->{}]'.format(keypoints3d[..., -1].sum(), self.frame_latest, self.frame_index))
|
||
if self.frame_index - self.frame_latest > 10 and self.init:
|
||
mywarn('Missing keypoints, resetting...')
|
||
self.init = False
|
||
self.records = []
|
||
self.results = []
|
||
return -1
|
||
self.records.append(k3d)
|
||
if len(self.records) < self.INIT_SIZE:
|
||
return -1
|
||
with Timer('fitting', True):
|
||
params = self.try_to_init(self.records)
|
||
params = self.fitting(params, self.records)
|
||
params = self.filter_result(params)
|
||
self.results.append(params)
|
||
self.params_newest = params
|
||
self.frame_latest = self.frame_index
|
||
return self.smooth_results()
|
||
|
||
class FitterCPPCache:
|
||
def __init__(self) -> None:
|
||
self.init = False
|
||
self.records = []
|
||
self.frame_index = 0
|
||
self.frame_latest = 0
|
||
|
||
def try_to_init(self, keypoints3d):
|
||
if self.init:
|
||
return copy.deepcopy(self.params_newest)
|
||
mywarn('>> Initialize')
|
||
params = self.handmodel.init_params(1)
|
||
if not (keypoints3d[..., -1] > 0.).all(): # 只有一半的点可见
|
||
return params
|
||
params = self.handmodel.fit3DShape(keypoints3d, params)
|
||
params = self.handmodel.init3DRT(keypoints3d[-1:], params)
|
||
params = self.handmodel.fit3DPose(keypoints3d[-1:], params)
|
||
mywarn('>> Initialize Rh = {}, Th = {}'.format(params['Rh'][0], params['Th'][0]))
|
||
self.init = True
|
||
return params
|
||
|
||
def check_keypoints(self, keypoints3d):
|
||
flag = (keypoints3d[..., -1]>0.3).sum() > 5
|
||
if len(self.records) > 1:
|
||
k_pre = self.records[-1][None]
|
||
dist = np.linalg.norm(keypoints3d[..., :3] - k_pre[..., :3], axis=-1)
|
||
conf = np.sqrt(keypoints3d[..., 3] * k_pre[..., 3])
|
||
dist_mean = (dist * conf).sum()/(1e-5+conf.sum())
|
||
print(dist_mean)
|
||
flag = flag and dist_mean < 0.1
|
||
return flag
|
||
|
||
def smooth_results(self, params=None):
|
||
if params is None:
|
||
params = self.params_newest.copy()
|
||
params['poses'] = params['poses'][:, 3:]
|
||
params['id'] = 0
|
||
return [params]
|
||
|
||
def __call__(self, keypoints3d):
|
||
self.frame_index += 1
|
||
flag = self.check_keypoints(keypoints3d)
|
||
if not flag:
|
||
mywarn('Missing keypoints {} [{}->{}]'.format(keypoints3d[..., -1].sum(), self.frame_latest, self.frame_index))
|
||
if self.frame_index - self.frame_latest > 10:
|
||
mywarn('Missing keypoints, resetting...')
|
||
self.init = False
|
||
self.records = []
|
||
return self.smooth_results(self.handmodel.init_params(1))
|
||
self.records.append(keypoints3d)
|
||
with Timer('fitting'):
|
||
params = self.try_to_init(keypoints3d[None])
|
||
params = self.handmodel.fit3DPose(keypoints3d[None], params)
|
||
self.params_newest = params
|
||
self.frame_latest = self.frame_index
|
||
return self.smooth_results()
|
||
class ManoFitterCPPCache(FitterCPPCache):
|
||
def __init__(self, name='LEFT') -> None:
|
||
super().__init__()
|
||
self.handmodel = load_object_from_cmd('config/model/mano_full_cpp.yml', [
|
||
'args.model_path', 'data/bodymodels/manov1.2/MANO_{}.pkl'.format(name),
|
||
'args.regressor_path', 'data/smplx/J_regressor_mano_{}.txt'.format(name),
|
||
])
|
||
|
||
class SMPLFitterCPPCache(FitterCPPCache):
|
||
def __init__(self, name='half') -> None:
|
||
super().__init__()
|
||
self.handmodel = load_object_from_cmd('config/model/mano_full_cpp.yml', [
|
||
'args.model_path', 'data/bodymodels/manov1.2/MANO_{}.pkl'.format(name),
|
||
'args.regressor_path', 'data/smplx/J_regressor_mano_{}.txt'.format(name),
|
||
])
|
||
|
||
class ManoFitterCPP:
|
||
def __init__(self, cfg_triangulator, key) -> None:
|
||
self.handmodel = load_object_from_cmd('config/model/mano_full_cpp.yml', [])
|
||
self.triangulator = load_object_from_cmd(cfg_triangulator, [])
|
||
self.time = 0
|
||
self.frame_latest = 0
|
||
self.frame_index = 0
|
||
self.key = 'handl3d'
|
||
self.init = False
|
||
self.params_newest = None
|
||
self.records, self.results = [], []
|
||
self.INIT_SIZE = 10
|
||
|
||
def get_keypoints3d(self, records, key=None):
|
||
if key is None:
|
||
return np.stack([r[self.key] for r in records])
|
||
else:
|
||
return np.stack([r[key] for r in records])
|
||
|
||
def try_to_init(self, records):
|
||
if self.init:
|
||
return copy.deepcopy(self.params_newest)
|
||
mywarn('>> Initialize')
|
||
keypoints3d = self.get_keypoints3d(records)
|
||
params = self.handmodel.init_params(1)
|
||
params = self.handmodel.fit3DShape(keypoints3d, params)
|
||
params = self.handmodel.init3DRT(keypoints3d[-1:], params)
|
||
params = self.handmodel.fit3DPose(keypoints3d[-1:], params)
|
||
mywarn('>> Initialize Rh = {}, Th = {}'.format(params['Rh'][0], params['Th'][0]))
|
||
self.init = True
|
||
return params
|
||
|
||
def smooth_results(self):
|
||
params = self.params_newest.copy()
|
||
params['poses'] = params['poses'][:, 3:]
|
||
params['id'] = 0
|
||
return [params]
|
||
|
||
def __call__(self, data):
|
||
self.frame_index += 1
|
||
k3d = self.triangulator(data)[0]
|
||
keypoints3d = self.get_keypoints3d([k3d])
|
||
# flag = self.check_keypoints(keypoints3d)
|
||
flag = True
|
||
if not flag:
|
||
mywarn('Missing keypoints {} [{}->{}]'.format(keypoints3d[..., -1].sum(), self.frame_latest, self.frame_index))
|
||
if self.frame_index - self.frame_latest > 10 and self.init:
|
||
mywarn('Missing keypoints, resetting...')
|
||
self.init = False
|
||
self.records = []
|
||
self.results = []
|
||
return -1
|
||
self.records.append(k3d)
|
||
if len(self.records) < self.INIT_SIZE:
|
||
return -1
|
||
with Timer('fitting'):
|
||
params = self.try_to_init(self.records)
|
||
k3d = self.get_keypoints3d(self.records[-1:])
|
||
params = self.handmodel.fit3DPose(k3d, params)
|
||
# params['poses'] = torch.Tensor(params['poses'][:, 3:])
|
||
# params['shapes'] = torch.Tensor(params['shapes'])
|
||
# params['Rh'] = torch.Tensor(params['Rh'])
|
||
# params['Th'] = torch.Tensor(params['Th'])
|
||
# params = self.filter_result(params)
|
||
self.results.append(params)
|
||
self.params_newest = params
|
||
self.frame_latest = self.frame_index
|
||
return self.smooth_results()
|
||
|
||
class BodyFitter(BaseFitter):
|
||
def __init__(self, key, **kwargs):
|
||
super().__init__(**kwargs)
|
||
self.key = key
|
||
|
||
def get_keypoints3d(self, records, key=None):
|
||
if key is None:
|
||
return torch.Tensor(np.stack([r[self.key] for r in records]))
|
||
else:
|
||
return torch.Tensor(np.stack([r[key] for r in records]))
|
||
|
||
class ManoFitter(BodyFitter):
|
||
def __init__(self, **kwargs) -> None:
|
||
super().__init__(**kwargs)
|
||
|
||
class HalfFitter(BodyFitter):
|
||
def __init__(self, **kwargs) -> None:
|
||
super().__init__(**kwargs)
|
||
self.INDEX_HALF = [0,1,2,3,4,5,6,7,15,16,17,18]
|
||
|
||
def get_keypoints3d(self, records):
|
||
THRES_WRIST = 0.05
|
||
keypoints3d = super().get_keypoints3d(records)
|
||
keypoints3d = keypoints3d[:, self.INDEX_HALF]
|
||
handl = super().get_keypoints3d(records, key='handl3d')
|
||
handr = super().get_keypoints3d(records, key='handr3d')
|
||
dist_ll = torch.norm(keypoints3d[:, 7, :3] - handl[:, 0, :3], dim=-1)
|
||
dist_rr = torch.norm(keypoints3d[:, 4, :3] - handr[:, 0, :3], dim=-1)
|
||
handl[dist_ll>THRES_WRIST] = 0.
|
||
handr[dist_rr>THRES_WRIST] = 0.
|
||
keypoints3d = np.hstack([keypoints3d, handl, handr])
|
||
conf = keypoints3d[..., 3:]
|
||
keypoints3d = np.hstack([(keypoints3d[..., :3] * conf).sum(axis=0)/(1e-5 + conf.sum(axis=0)), conf.min(axis=0)])
|
||
keypoints3d = keypoints3d[None]
|
||
# if (keypoints3d.shape[0] == 10):
|
||
return torch.Tensor(keypoints3d)
|
||
|
||
def filter_result(self, result):
|
||
result = super().filter_result(result)
|
||
# 限定一下关节旋转
|
||
# 手肘
|
||
# result['poses'][:, 5*3+1] = np.clip(result['poses'][:, 5*3+1], -2.5, 0.1)
|
||
# result['poses'][:, 6*3+1] = np.clip(result['poses'][:, 5*3+1], -0.1, 2.5)
|
||
# 手腕
|
||
return result
|
||
|
||
class HalfHandFitter(HalfFitter):
|
||
def __init__(self, cfg_handl, cfg_handr, **kwargs) -> None:
|
||
super().__init__(**kwargs)
|
||
self.handl = load_object_from_cmd(cfg_handl, [])
|
||
self.handr = load_object_from_cmd(cfg_handr, [])
|
||
|
||
def get_init_params(self, nFrames):
|
||
params = super().get_init_params(nFrames)
|
||
params_ = self.handl.get_init_params(nFrames)
|
||
params['shapes_handl'] = params_['shapes']
|
||
params['shapes_handr'] = params_['shapes'].clone()
|
||
params['Rh_handl'] = torch.zeros((nFrames, 3))
|
||
params['Rh_handr'] = torch.zeros((nFrames, 3))
|
||
params['Th_handl'] = torch.zeros((nFrames, 3))
|
||
params['Th_handr'] = torch.zeros((nFrames, 3))
|
||
return params
|
||
|
||
def fitPose(self, keypoints3d, params, weight, option):
|
||
keypoints = {
|
||
'handl': keypoints3d[:, -21-21:-21, :],
|
||
'handr': keypoints3d[:, -21:, :]
|
||
}
|
||
for key in ['handl', 'handr']:
|
||
kpts = keypoints[key]
|
||
params_ = {
|
||
'id': 0,
|
||
'Rh': params['Rh_'+key],
|
||
'Th': params['Th_'+key],
|
||
'shapes': params['shapes_'+key],
|
||
'poses': params[key],
|
||
}
|
||
if key == 'handl':
|
||
params_ = self.handl.fitPose(kpts, params_, self.handl.cfg.pose.weight, self.handl.cfg.pose.option)
|
||
else:
|
||
params_ = self.handr.fitPose(kpts, params_, self.handr.cfg.pose.weight, self.handr.cfg.pose.option)
|
||
params['Rh_'+key] = params_['Rh']
|
||
params['Th_'+key] = params_['Th']
|
||
params['shapes_'+key] = params_['shapes']
|
||
params[key] = params_['poses']
|
||
return super().fitPose(keypoints3d, params, weight, option,
|
||
kpts_index=[0,1,2,3,4,5,6,7,8,9,10,11,
|
||
12, 17, 21, 25, 29,
|
||
24, 29, 33, 37, 41])
|
||
|
||
def try_to_init(self, records):
|
||
if self.init:
|
||
return self.params_newest
|
||
params = super().try_to_init(records)
|
||
self.handl.init = False
|
||
self.handr.init = False
|
||
key = 'handl'
|
||
params_ = self.handl.try_to_init(records)
|
||
params['handl'] = params_['poses']
|
||
params['Rh_'+key] = params_['Rh']
|
||
params['Th_'+key] = params_['Th']
|
||
params['shapes_'+key] = params_['shapes']
|
||
key = 'handr'
|
||
params_ = self.handr.try_to_init(records)
|
||
params[key] = params_['poses']
|
||
params['Rh_'+key] = params_['Rh']
|
||
params['Th_'+key] = params_['Th']
|
||
params['shapes_'+key] = params_['shapes']
|
||
return params
|
||
|
||
if __name__ == '__main__':
|
||
from glob import glob
|
||
from os.path import join
|
||
from tqdm import tqdm
|
||
from ..mytools.file_utils import read_json
|
||
from ..config.baseconfig import load_object_from_cmd
|
||
|
||
data= '/nas/datasets/EasyMocap/302'
|
||
# data = '/home/qing/Dataset/handl'
|
||
mode = 'half'
|
||
# data = '/home/qing/DGPU/home/shuaiqing/zju-mocap-mp/female-jump'
|
||
data = '/home/qing/Dataset/desktop/0402/test3'
|
||
k3dnames = sorted(glob(join(data, 'output-keypoints3d', 'keypoints3d', '*.json')))
|
||
if mode == 'handl':
|
||
fitter = load_object_from_cmd('config/recon/fit_manol.yml', [])
|
||
elif mode == 'half':
|
||
fitter = load_object_from_cmd('config/recon/fit_half.yml', [])
|
||
elif mode == 'smpl':
|
||
fitter = load_object_from_cmd('config/recon/fit_smpl.yml', [])
|
||
from easymocap.socket.base_client import BaseSocketClient
|
||
client = BaseSocketClient('0.0.0.0', 9999)
|
||
|
||
for k3dname in tqdm(k3dnames):
|
||
k3ds = read_json(k3dname)
|
||
if mode == 'handl':
|
||
k3ds = np.array(k3ds[0]['handl3d'])
|
||
data = {fitter.key3d: k3ds}
|
||
elif mode == 'half':
|
||
data = {
|
||
'keypoints3d': np.array(k3ds[0]['keypoints3d']),
|
||
'handl3d': np.array(k3ds[0]['handl3d']),
|
||
'handr3d': np.array(k3ds[0]['handr3d'])
|
||
}
|
||
elif mode == 'smpl':
|
||
k3ds = np.array(k3ds[0]['keypoints3d'])
|
||
data = {fitter.key3d: k3ds}
|
||
results = fitter(data)
|
||
if results != -1:
|
||
client.send_smpl(results)
|