Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+309
@@ -0,0 +1,309 @@
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from body_model import MANO_JOINTS, run_mano
|
||||
from geometry.rotation import (
|
||||
rotation_matrix_to_angle_axis,
|
||||
angle_axis_to_rotation_matrix,
|
||||
)
|
||||
from util.logger import Logger
|
||||
from util.tensor import move_to, detach_all
|
||||
|
||||
# from .helpers import estimate_initial_trans
|
||||
from .params import CameraParams
|
||||
|
||||
J_HAND = len(MANO_JOINTS) - 1 # no root
|
||||
|
||||
|
||||
class BaseSceneModel(nn.Module):
|
||||
"""
|
||||
Scene model of sequences of human poses.
|
||||
All poses are in their own INDEPENDENT camera reference frames.
|
||||
A basic class mostly for testing purposes.
|
||||
|
||||
Parameters:
|
||||
batch_size: number of sequences to optimize
|
||||
seq_len: length of the sequences
|
||||
body_model: MANO hand model
|
||||
pose_prior: VPoser model
|
||||
fit_gender: gender of model (optional)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
# fit_gender="neutral",
|
||||
use_init=False,
|
||||
opt_cams=False,
|
||||
opt_scale=True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
B, T = batch_size, seq_len
|
||||
self.batch_size = batch_size
|
||||
self.seq_len = seq_len
|
||||
|
||||
self.body_model = body_model
|
||||
self.hand_mean = self.body_model.hand_mean
|
||||
|
||||
self.pose_prior = pose_prior
|
||||
if self.pose_prior is not None:
|
||||
self.latent_pose_dim = self.pose_prior.latentD
|
||||
|
||||
self.num_betas = body_model.num_betas
|
||||
|
||||
self.use_init = use_init
|
||||
print("USE INIT", use_init)
|
||||
self.opt_scale = opt_scale
|
||||
self.opt_cams = opt_cams
|
||||
print("OPT SCALE", self.opt_scale)
|
||||
print("OPT CAMERAS", self.opt_cams)
|
||||
print("Batch size: ", batch_size)
|
||||
self.params = CameraParams(batch_size)
|
||||
|
||||
def initialize(self, obs_data, cam_data):
|
||||
Logger.log("Initializing scene model with observed data")
|
||||
|
||||
# initialize cameras
|
||||
self.params.set_cameras(
|
||||
cam_data,
|
||||
opt_scale=self.opt_scale,
|
||||
opt_cams=self.opt_cams,
|
||||
opt_focal=self.opt_cams,
|
||||
)
|
||||
|
||||
# initialize body params
|
||||
B, T = self.batch_size, self.seq_len
|
||||
device = next(iter(cam_data.values())).device
|
||||
init_betas = torch.mean(obs_data["init_body_shape"], dim=1) # torch.zeros(B, self.num_betas, device=device)
|
||||
|
||||
if self.use_init and "init_body_pose" in obs_data:
|
||||
init_pose = obs_data["init_body_pose"][:, :, :J_HAND, :]
|
||||
init_pose_latent = self.pose2latent(init_pose)
|
||||
else:
|
||||
raise ValueError
|
||||
# init_pose = torch.zeros(B, T, J_HAND, 3, device=device)
|
||||
# init_pose_latent = torch.zeros(B, T, self.latent_pose_dim, device=device)
|
||||
|
||||
# transform into world frame (T, 3, 3), (T, 3)
|
||||
R_w2c, t_w2c = cam_data["cam_R"], cam_data["cam_t"]
|
||||
R_c2w = R_w2c.transpose(-1, -2)
|
||||
t_c2w = -torch.einsum("tij,tj->ti", R_c2w, t_w2c)
|
||||
|
||||
if self.use_init and "init_root_orient" in obs_data:
|
||||
init_rot = obs_data["init_root_orient"] # (B, T, 3)
|
||||
init_rot_mat = angle_axis_to_rotation_matrix(init_rot)
|
||||
print(R_c2w.shape, init_rot_mat.shape)
|
||||
init_rot_mat = torch.einsum("tij,btjk->btik", R_c2w, init_rot_mat)
|
||||
init_rot = rotation_matrix_to_angle_axis(init_rot_mat)
|
||||
else:
|
||||
raise ValueError
|
||||
# init_rot = (
|
||||
# torch.tensor([np.pi, 0, 0], dtype=torch.float32)
|
||||
# .reshape(1, 1, 3)
|
||||
# .repeat(B, T, 1)
|
||||
# )
|
||||
|
||||
init_trans = torch.zeros(B, T, 3, device=device)
|
||||
if self.use_init and "init_trans" in obs_data:
|
||||
# must offset by the root location before applying camera to world transform
|
||||
is_right = obs_data["is_right"]
|
||||
pred_data = self.pred_mano(init_trans, init_rot, init_pose, is_right, init_betas)
|
||||
root_loc = pred_data["joints3d"][..., 0, :] # (B, T, 3)
|
||||
init_trans = obs_data["init_trans"] # (B, T, 3)
|
||||
init_trans = (
|
||||
torch.einsum("tij,btj->bti", R_c2w, init_trans + root_loc)
|
||||
+ t_c2w[None]
|
||||
- root_loc
|
||||
)
|
||||
else:
|
||||
# initialize trans with reprojected joints
|
||||
# pred_data = self.pred_mano(init_trans, init_rot, init_pose, init_betas)
|
||||
# init_trans = estimate_initial_trans(
|
||||
# init_pose,
|
||||
# pred_data["joints3d_op"],
|
||||
# obs_data["joints2d"],
|
||||
# obs_data["intrins"][:, 0],
|
||||
# )
|
||||
raise ValueError
|
||||
|
||||
self.params.set_param("init_body_pose", init_pose)
|
||||
self.params.set_param("latent_pose", init_pose_latent)
|
||||
self.params.set_param("betas", init_betas)
|
||||
self.params.set_param("trans", init_trans)
|
||||
self.params.set_param("root_orient", init_rot)
|
||||
self.params.set_param("is_right", is_right, requires_grad=False)
|
||||
|
||||
# Store initial latent pose in obs_data for pose prior loss
|
||||
obs_data["init_latent_pose"] = init_pose_latent.detach()
|
||||
# print(init_pose.shape, init_pose_latent.shape, init_betas.shape, init_trans.shape, init_rot.shape, is_right.shape)
|
||||
# raise ValueError
|
||||
|
||||
def get_optim_result(self, **kwargs):
|
||||
"""
|
||||
Collect predicted outputs (latent_pose, trans, root_orient, betas, body pose) into dict
|
||||
"""
|
||||
print('running get_optim_result in base_scene.py...')
|
||||
res = self.params.get_dict()
|
||||
if "latent_pose" in res:
|
||||
res["pose_body"] = self.latent2pose(self.params.latent_pose).detach()
|
||||
|
||||
# add the cameras
|
||||
res["cam_R"], res["cam_t"], _, _ = self.params.get_cameras()
|
||||
res["intrins"] = self.params.intrins
|
||||
return {"world": res}
|
||||
|
||||
def latent2pose(self, latent_pose):
|
||||
"""
|
||||
Converts VPoser latent embedding to aa body pose.
|
||||
latent_pose : B x T x D
|
||||
body_pose : B x T x J*3
|
||||
"""
|
||||
if self.pose_prior is not None:
|
||||
B, T, _ = latent_pose.size()
|
||||
d_latent = self.pose_prior.latentD
|
||||
latent_pose = latent_pose.reshape((-1, d_latent))
|
||||
body_pose = self.pose_prior.decode(latent_pose, output_type="matrot")
|
||||
body_pose = rotation_matrix_to_angle_axis(
|
||||
body_pose.reshape((B * T * J_HAND, 3, 3))
|
||||
).reshape((B, T, J_HAND * 3))
|
||||
return body_pose + self.hand_mean
|
||||
else:
|
||||
return latent_pose
|
||||
|
||||
def pose2latent(self, body_pose):
|
||||
"""
|
||||
Encodes aa body pose to VPoser latent space.
|
||||
body_pose : B x T x J*3
|
||||
latent_pose : B x T x D
|
||||
"""
|
||||
if self.pose_prior is not None:
|
||||
B, T = body_pose.shape[:2]
|
||||
body_pose = body_pose.reshape((-1, J_HAND * 3))
|
||||
latent_pose_distrib = self.pose_prior.encode(body_pose - self.hand_mean)
|
||||
d_latent = self.pose_prior.latentD
|
||||
latent_pose = latent_pose_distrib.mean.reshape((B, T, d_latent))
|
||||
return latent_pose
|
||||
else:
|
||||
return body_pose
|
||||
|
||||
def pred_mano(self, trans, root_orient, body_pose, is_right, betas):
|
||||
"""
|
||||
Forward pass of the MANO model and populates pred_data accordingly with
|
||||
joints3d, verts3d, points3d.
|
||||
|
||||
trans : B x T x 3
|
||||
root_orient : B x T x 3
|
||||
body_pose : B x T x J*3
|
||||
betas : B x D
|
||||
"""
|
||||
|
||||
mano_out = run_mano(self.body_model, trans, root_orient, body_pose, is_right, betas=betas)
|
||||
joints3d, points3d = mano_out["joints"], mano_out["vertices"]
|
||||
|
||||
# select desired joints and vertices
|
||||
# joints3d_body = joints3d
|
||||
# joints3d_op = joints3d.clone()
|
||||
# joints3d_op = joints3d[:, :, self.smpl2op_map, :]
|
||||
# # hacky way to get hip joints that align with ViTPose keypoints
|
||||
# # this could be moved elsewhere in the future (and done properly)
|
||||
# joints3d_op[:, :, [9, 12]] = (
|
||||
# joints3d_op[:, :, [9, 12]]
|
||||
# + 0.25 * (joints3d_op[:, :, [9, 12]] - joints3d_op[:, :, [12, 9]])
|
||||
# + 0.5
|
||||
# * (
|
||||
# joints3d_op[:, :, [8]]
|
||||
# - 0.5 * (joints3d_op[:, :, [9, 12]] + joints3d_op[:, :, [12, 9]])
|
||||
# )
|
||||
# )
|
||||
verts3d = points3d #[:, :, KEYPT_VERTS, :]
|
||||
|
||||
return {
|
||||
"is_right": is_right,
|
||||
"points3d": points3d, # all vertices
|
||||
"verts3d": verts3d, # keypoint vertices
|
||||
"joints3d": joints3d, # smpl joints
|
||||
"joints3d_op": joints3d, # OP joints
|
||||
"l_faces": mano_out["l_faces"], # index array of faces
|
||||
"r_faces": mano_out["r_faces"], # index array of faces
|
||||
"body_pose": mano_out["body_pose"]
|
||||
}
|
||||
|
||||
def pred_params_mano(self, is_right, reproj=True):
|
||||
body_pose = self.latent2pose(self.params.latent_pose)
|
||||
pred_data = self.pred_mano(
|
||||
self.params.trans, self.params.root_orient, body_pose, is_right, self.params.betas
|
||||
)
|
||||
# from geometry.mesh import save_mesh_scenes, vertices_to_trimesh
|
||||
# LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353)
|
||||
# print(pred_data.keys())
|
||||
# verts, l_faces, r_faces, is_right = pred_data['verts3d'], pred_data['l_faces'], pred_data['l_faces'], pred_data['is_right']
|
||||
# T = len(verts[0])
|
||||
# print(f"{T} mesh frames")
|
||||
# print(verts.shape, l_faces.shape, is_right.shape, is_right)
|
||||
# times = list(range(0, T, 1))
|
||||
# os.makedirs('./mesh/')
|
||||
# for t in times:
|
||||
|
||||
# tmesh = vertices_to_trimesh(verts[0][t].detach().cpu().numpy(), l_faces.detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
# tmesh.export(os.path.join('./mesh/', f'{str(t).zfill(6)}_0.obj'))
|
||||
|
||||
# tmesh = vertices_to_trimesh(verts[1][t].detach().cpu().numpy(), r_faces.detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
# tmesh.export(os.path.join('./mesh/', f'{str(t).zfill(6)}_1.obj'))
|
||||
# print(os.path.join('./mesh/', f'{str(t).zfill(6)}_1.obj'))
|
||||
|
||||
#################################
|
||||
# print(self.params.root_orient.shape, self.params.init_body_pose.shape, body_pose.shape, self.params.trans.shape, self.params.betas.shape)
|
||||
# hand_model = mano.load(
|
||||
# model_path= '/vol/bitbucket/zy3023/code/hand/slam-hand/_DATA/data/mano',# '/vol/bitbucket/zy3023/code/hand/slam-hand/VPoser/pretrained/Vposer_right_mirrored',
|
||||
# is_right=True,
|
||||
# num_pca_comps=45,
|
||||
# batch_size=1,
|
||||
# flat_hand_mean=False,
|
||||
# ).cuda()
|
||||
|
||||
# output = hand_model(
|
||||
# betas=self.params.betas[0][None].repeat(180,1),
|
||||
# global_orient=self.params.root_orient[0],
|
||||
# hand_pose=self.params.init_body_pose[0].reshape(180, -1) - self.hand_mean,
|
||||
# transl=self.params.trans[0],
|
||||
# return_verts=True,
|
||||
# return_tips=False
|
||||
# )
|
||||
# np.save('x.npy', self.params.init_body_pose[0].reshape(180, -1).detach().cpu().numpy())
|
||||
|
||||
# out_path = './out'
|
||||
# h_meshes = hand_model.hand_meshes(output)
|
||||
# if not os.path.exists(out_path):
|
||||
# os.mkdir(out_path + '_after')
|
||||
# os.mkdir(out_path)
|
||||
|
||||
# for i in range(180):
|
||||
# #print(f"export to {out_path}/{str(i).zfill(6)}.ply")
|
||||
# h_meshes[i].export(f"{out_path}/{str(i).zfill(6)}.ply")
|
||||
|
||||
# # print(body_pose.shape, body_pose[0])
|
||||
# print(self.hand_mean)
|
||||
# output = hand_model(
|
||||
# betas=self.params.betas[0][None].repeat(180,1),
|
||||
# global_orient=self.params.root_orient[0],
|
||||
# hand_pose=body_pose[0].reshape(180, -1) - self.hand_mean,
|
||||
# transl=self.params.trans[0],
|
||||
# return_verts=True,
|
||||
# return_tips=False
|
||||
# )
|
||||
|
||||
# h_meshes = hand_model.hand_meshes(output)
|
||||
# for i in range(180):
|
||||
# #print(f"export to {out_path}_after/{str(i).zfill(6)}.ply")
|
||||
# h_meshes[i].export(f"{out_path}_after/{str(i).zfill(6)}.ply")
|
||||
# exit()
|
||||
|
||||
return pred_data
|
||||
+399
@@ -0,0 +1,399 @@
|
||||
import os
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as torch_f
|
||||
|
||||
SNAP_PARENT = [
|
||||
0, # 0's parent
|
||||
0, # 1's parent
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
0, # 5's parent
|
||||
5,
|
||||
6,
|
||||
7,
|
||||
0, # 9's parent
|
||||
9,
|
||||
10,
|
||||
11,
|
||||
0, # 13's parent
|
||||
13,
|
||||
14,
|
||||
15,
|
||||
0, # 17's parent
|
||||
17,
|
||||
18,
|
||||
19,
|
||||
]
|
||||
|
||||
JOINT_ROOT_IDX = 9
|
||||
|
||||
REF_BONE_LINK = (0, 9) # mid mcp
|
||||
|
||||
# bone indexes in 20 bones setting
|
||||
ID_ROOT_bone = np.array([0, 4, 8, 12, 16]) # ROOT_bone from wrist to MCP
|
||||
ID_PIP_bone = np.array([1, 5, 9, 13, 17]) # PIP_bone from MCP to PIP
|
||||
ID_DIP_bone = np.array([2, 6, 10, 14, 18]) # DIP_bone from PIP to DIP
|
||||
ID_TIP_bone = np.array([3, 7, 11, 15, 19]) # TIP_bone from DIP to TIP
|
||||
|
||||
def plot_hull(theta, hull):
|
||||
del_rdp_hull = hull.detach().cpu().numpy()
|
||||
theta = theta.detach().cpu().numpy()
|
||||
|
||||
fig = plt.figure()
|
||||
figManager = plt.get_current_fig_manager()
|
||||
figManager.window.showMaximized()
|
||||
|
||||
ax = fig.add_subplot(111)
|
||||
ax.scatter(theta[:, 0], theta[:, 1], s=10, c='r')
|
||||
ax.set_xlabel("flexion")
|
||||
ax.set_ylabel("abduction")
|
||||
|
||||
plt.plot(del_rdp_hull[:, 0],
|
||||
del_rdp_hull[:, 1],
|
||||
'-yo', linewidth=2)
|
||||
|
||||
plt.xticks(np.arange(-3, 4, 1))
|
||||
plt.yticks(np.arange(-2, 2, 0.5))
|
||||
plt.show()
|
||||
|
||||
|
||||
def two_norm(a):
|
||||
'''
|
||||
|
||||
Args:
|
||||
a: B*M*2 or B*M*3
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
return torch.norm(a, dim=-1)
|
||||
|
||||
|
||||
def one_norm(a):
|
||||
'''
|
||||
|
||||
Args:
|
||||
a: B*M*2 or B*M*3
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
return torch.norm(a, dim=-1, p=1)
|
||||
|
||||
|
||||
def calculate_joint_angle_loss(thetas, hulls):
|
||||
'''
|
||||
|
||||
Args:
|
||||
Theta: B*15*2
|
||||
hulls: list
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
|
||||
loss = torch.Tensor([0]).cuda()
|
||||
for i in range(15):
|
||||
# print("i=",i)
|
||||
hull = hulls[i] # (M*2)
|
||||
theta = thetas[:, i] # (B*2)
|
||||
hull = torch.cat((hull, hull[0].unsqueeze(0)), dim=0)
|
||||
|
||||
v = (hull[1:] - hull[:-1]).unsqueeze(0) # (M-1)*2
|
||||
w = - hull[:-1].unsqueeze(0) + theta.unsqueeze(1).repeat(1, hull[:-1].shape[0], 1) # B*(M-1)*2
|
||||
|
||||
cross_product_2d = w[:, :, 0] * v[:, :, 1] - w[:, :, 1] * v[:, :, 0]
|
||||
tmp = torch.sum(cross_product_2d < 1e-6, dim=-1)
|
||||
|
||||
is_outside = (tmp != (hull.shape[0] - 1))
|
||||
if not torch.sum(is_outside):
|
||||
sub_loss = torch.Tensor([0]).cuda()
|
||||
else:
|
||||
outside_theta = theta[is_outside]
|
||||
outside_theta = outside_theta.unsqueeze(1).repeat(1, hull[:-1].shape[0], 1)
|
||||
w_outside = - hull[:-1].unsqueeze(0) + outside_theta # B*(M-1)*2
|
||||
t = torch.clamp(inner_product(w_outside, v) / (two_norm(v) ** 2), min=0, max=1).unsqueeze(2)
|
||||
p = hull[:-1] + t * v
|
||||
|
||||
D = one_norm(torch.cos(outside_theta) - torch.cos(p)) + one_norm(torch.sin(outside_theta) - torch.sin(p))
|
||||
sub_loss = torch.sum(torch.min(D, dim=-1)[0])
|
||||
|
||||
vis = 0
|
||||
if vis:
|
||||
print(theta)
|
||||
plot_hull(theta, hull)
|
||||
|
||||
loss += sub_loss
|
||||
|
||||
loss /= (15 * thetas.shape[0])
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
def angle_between(v1, v2):
|
||||
epsilon = 1e-7
|
||||
cos = torch_f.cosine_similarity(v1, v2, dim=-1).clamp(-1 + epsilon, 1 - epsilon) # (B)
|
||||
theta = torch.acos(cos) # (B)
|
||||
return theta
|
||||
|
||||
|
||||
def normalize(vec):
|
||||
return torch_f.normalize(vec, p=2, dim=-1)
|
||||
|
||||
|
||||
def inner_product(x1, x2):
|
||||
return torch.sum(x1 * x2, dim=-1)
|
||||
|
||||
|
||||
def cross_product(x1, x2):
|
||||
return torch.cross(x1, x2, dim=-1)
|
||||
|
||||
|
||||
def axangle2mat_torch(axis, angle, is_normalized=False):
|
||||
""" Rotation matrix for rotation angle `angle` around `axis`
|
||||
Parameters
|
||||
----------
|
||||
axis : [B,M, 3] element sequence
|
||||
vector specifying axis for rotation.
|
||||
angle :[B,M, ] scalar
|
||||
angle of rotation in radians.
|
||||
is_normalized : bool, optional
|
||||
True if `axis` is already normalized (has norm of 1). Default False.
|
||||
Returns
|
||||
-------
|
||||
mat : array shape (B, M,3,3)
|
||||
rotation matrix for specified rotation
|
||||
Notes
|
||||
-----
|
||||
From: http://en.wikipedia.org/wiki/Rotation_matrix#Axis_and_angle
|
||||
"""
|
||||
B = axis.shape[0]
|
||||
M = axis.shape[1]
|
||||
|
||||
if not is_normalized:
|
||||
norm_axis = axis.norm(p=2, dim=-1, keepdim=True)
|
||||
normed_axis = axis / norm_axis
|
||||
else:
|
||||
normed_axis = axis
|
||||
x, y, z = normed_axis[:, :, 0], normed_axis[:, :, 1], normed_axis[:, :, 2]
|
||||
c = torch.cos(angle)
|
||||
s = torch.sin(angle)
|
||||
C = 1 - c
|
||||
|
||||
xs = x * s;
|
||||
ys = y * s;
|
||||
zs = z * s # noqa
|
||||
xC = x * C;
|
||||
yC = y * C;
|
||||
zC = z * C # noqa
|
||||
xyC = x * yC;
|
||||
yzC = y * zC;
|
||||
zxC = z * xC # noqa
|
||||
|
||||
TMP = torch.stack([x * xC + c, xyC - zs, zxC + ys, xyC + zs, y * yC + c, yzC - xs, zxC - ys, yzC + xs, z * zC + c],
|
||||
dim=-1)
|
||||
return TMP.reshape(B, M, 3, 3)
|
||||
|
||||
|
||||
def interval_loss(value, min, max):
|
||||
'''
|
||||
calculate interval loss
|
||||
Args:
|
||||
value: B*M
|
||||
max: M
|
||||
min: M
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
|
||||
batch_3d_size = value.shape[0]
|
||||
|
||||
min = min.repeat(value.shape[0], 1)
|
||||
max = max.repeat(value.shape[0], 1)
|
||||
|
||||
loss1 = torch.max(min - value, torch.Tensor([0]).cuda())
|
||||
loss2 = torch.max(value - max, torch.Tensor([0]).cuda())
|
||||
|
||||
loss = (loss1 + loss2).sum()
|
||||
|
||||
loss /= (batch_3d_size * value.shape[1])
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
class BMCLoss:
|
||||
def __init__(
|
||||
self,
|
||||
lambda_bl=0.,
|
||||
lambda_rb=0.,
|
||||
lambda_a=0.,
|
||||
):
|
||||
self.lambda_bl = lambda_bl
|
||||
self.lambda_rb = lambda_rb
|
||||
self.lambda_a = lambda_a
|
||||
|
||||
self.lp = os.path.join(os.path.dirname(__file__), "../../_DATA/BMC")
|
||||
# resolved_path = os.path.abspath(self.lp)
|
||||
# print(f"Resolved path: {resolved_path}")
|
||||
# print(f"Current working directory: {os.getcwd()}")
|
||||
# exit()
|
||||
|
||||
self.bone_len_max = np.load(os.path.join(self.lp, "bone_len_max.npy"))
|
||||
self.bone_len_min = np.load(os.path.join(self.lp, "bone_len_min.npy"))
|
||||
self.rb_curvatures_max = np.load(os.path.join(self.lp, "curvatures_max.npy"))
|
||||
self.rb_curvatures_min = np.load(os.path.join(self.lp, "curvatures_min.npy"))
|
||||
self.rb_PHI_max = np.load(os.path.join(self.lp, "PHI_max.npy"))
|
||||
self.rb_PHI_min = np.load(os.path.join(self.lp, "PHI_min.npy"))
|
||||
|
||||
self.joint_angle_limit = np.load(os.path.join(self.lp, "CONVEX_HULLS.npy"),
|
||||
allow_pickle=True)
|
||||
LEN_joint_angle_limit = len(self.joint_angle_limit)
|
||||
|
||||
self.bl_max = torch.from_numpy(self.bone_len_max).float().cuda()
|
||||
self.bl_min = torch.from_numpy(self.bone_len_min).float().cuda()
|
||||
self.rb_curvatures_max = torch.from_numpy(self.rb_curvatures_max).float().cuda()
|
||||
self.rb_curvatures_min = torch.from_numpy(self.rb_curvatures_min).float().cuda()
|
||||
self.rb_PHI_max = torch.from_numpy(self.rb_PHI_max).float().cuda()
|
||||
self.rb_PHI_min = torch.from_numpy(self.rb_PHI_min).float().cuda()
|
||||
|
||||
self.joint_angle_limit = [torch.from_numpy(self.joint_angle_limit[i]).float().cuda() for i in
|
||||
range(LEN_joint_angle_limit)]
|
||||
|
||||
def compute_loss(self, ori_joints, valid_mask):
|
||||
'''
|
||||
|
||||
Args:
|
||||
joints: B*21*3
|
||||
|
||||
Returns:
|
||||
|
||||
'''
|
||||
|
||||
joints = ori_joints[valid_mask]
|
||||
print(joints.shape, ori_joints.shape)
|
||||
batch_size = joints.shape[0]
|
||||
final_loss = torch.Tensor([0]).cuda()
|
||||
|
||||
BMC_losses = {"bmc_bl": torch.Tensor([0]).cuda(), "bmc_rb": torch.Tensor([0]).cuda(), "bmc_a": torch.Tensor([0]).cuda()}
|
||||
|
||||
if (self.lambda_bl < 1e-6) and (self.lambda_rb < 1e-6) and (self.lambda_a < 1e-6):
|
||||
return final_loss, BMC_losses
|
||||
|
||||
ALL_bones = [
|
||||
(
|
||||
joints[:, i, :] -
|
||||
joints[:, SNAP_PARENT[i], :]
|
||||
) for i in range(21)
|
||||
]
|
||||
ALL_bones = torch.stack(ALL_bones[1:], dim=1) # (B,20,3)
|
||||
ROOT_bones = ALL_bones[:, ID_ROOT_bone] # (B,5,3)
|
||||
PIP_bones = ALL_bones[:, ID_PIP_bone]
|
||||
DIP_bones = ALL_bones[:, ID_DIP_bone]
|
||||
TIP_bones = ALL_bones[:, ID_TIP_bone]
|
||||
|
||||
ALL_Z_axis = normalize(ALL_bones)
|
||||
PIP_Z_axis = ALL_Z_axis[:, ID_ROOT_bone]
|
||||
DIP_Z_axis = ALL_Z_axis[:, ID_PIP_bone]
|
||||
TIP_Z_axis = ALL_Z_axis[:, ID_DIP_bone]
|
||||
|
||||
normals = normalize(cross_product(ROOT_bones[:, 1:], ROOT_bones[:, :-1]))
|
||||
|
||||
# compute loss of bone length
|
||||
bl_loss = torch.Tensor([0]).cuda()
|
||||
if self.lambda_bl:
|
||||
bls = two_norm(ALL_bones) # (B,20,1)
|
||||
bl_loss = interval_loss(value=bls, min=self.bl_min, max=self.bl_max)
|
||||
final_loss += self.lambda_bl * bl_loss
|
||||
BMC_losses["bmc_bl"] = bl_loss
|
||||
|
||||
# compute loss of Root bones
|
||||
rb_loss = torch.Tensor([0]).cuda()
|
||||
if self.lambda_rb:
|
||||
edge_normals = torch.zeros_like(ROOT_bones).cuda() # (B,5,3)
|
||||
edge_normals[:, [0, 4]] = normals[:, [0, 3]]
|
||||
edge_normals[:, 1:4] = normalize(normals[:, 1:4] + normals[:, :3])
|
||||
|
||||
curvatures = inner_product(edge_normals[:, 1:] - edge_normals[:, :4],
|
||||
ROOT_bones[:, 1:] - ROOT_bones[:, :4]) / \
|
||||
(two_norm(ROOT_bones[:, 1:] - ROOT_bones[:, :4]) ** 2)
|
||||
PHI = angle_between(ROOT_bones[:, :4], ROOT_bones[:, 1:]) # (B)
|
||||
|
||||
rb_loss = interval_loss(value=curvatures, min=self.rb_curvatures_min, max=self.rb_curvatures_max) + \
|
||||
interval_loss(value=PHI, min=self.rb_PHI_min, max=self.rb_PHI_max)
|
||||
final_loss += self.lambda_rb * rb_loss
|
||||
BMC_losses["bmc_rb"] = rb_loss
|
||||
|
||||
# compute loss of Joint angles
|
||||
a_loss = torch.Tensor([0]).cuda()
|
||||
if self.lambda_a:
|
||||
# PIP bones
|
||||
PIP_X_axis = torch.zeros([batch_size, 5, 3]).cuda() # (B,5,3)
|
||||
PIP_X_axis[:, [0, 1, 4], :] = -normals[:, [0, 1, 3]]
|
||||
PIP_X_axis[:, 2:4] = -normalize(normals[:, 2:4] + normals[:, 1:3]) # (B,2,3)
|
||||
PIP_Y_axis = normalize(cross_product(PIP_Z_axis, PIP_X_axis)) # (B,5,3)
|
||||
|
||||
PIP_bones_xz = PIP_bones - inner_product(PIP_bones, PIP_Y_axis).unsqueeze(2) * PIP_Y_axis
|
||||
PIP_theta_flexion = angle_between(PIP_bones_xz, PIP_Z_axis) # in global coordinate (B)
|
||||
PIP_theta_abduction = angle_between(PIP_bones_xz, PIP_bones) # in global coordinate (B)
|
||||
# x-component of the bone vector
|
||||
tmp = inner_product(PIP_bones, PIP_X_axis)
|
||||
PIP_theta_flexion = torch.where(tmp < 1e-6, -PIP_theta_flexion, PIP_theta_flexion)
|
||||
# y-component of the bone vector
|
||||
tmp = torch.sum((PIP_bones * PIP_Y_axis), dim=-1)
|
||||
PIP_theta_abduction = torch.where(tmp < 1e-6, -PIP_theta_abduction, PIP_theta_abduction)
|
||||
|
||||
temp_axis = normalize(cross_product(PIP_Z_axis, PIP_bones))
|
||||
temp_alpha = angle_between(PIP_Z_axis, PIP_bones) # alpha belongs to [pi/2, pi]
|
||||
temp_R = axangle2mat_torch(axis=temp_axis, angle=temp_alpha, is_normalized=True)
|
||||
|
||||
# DIP bones
|
||||
DIP_X_axis = torch.matmul(temp_R, PIP_X_axis.unsqueeze(3)).squeeze()
|
||||
DIP_Y_axis = torch.matmul(temp_R, PIP_Y_axis.unsqueeze(3)).squeeze()
|
||||
|
||||
DIP_bones_xz = DIP_bones - inner_product(DIP_bones, DIP_Y_axis).unsqueeze(2) * DIP_Y_axis
|
||||
DIP_theta_flexion = angle_between(DIP_bones_xz, DIP_Z_axis) # in global coordinate
|
||||
DIP_theta_abduction = angle_between(DIP_bones_xz, DIP_bones) # in global coordinate
|
||||
# x-component of the bone vector
|
||||
tmp = inner_product(DIP_bones, DIP_X_axis)
|
||||
DIP_theta_flexion = torch.where(tmp < 0, -DIP_theta_flexion, DIP_theta_flexion)
|
||||
# y-component of the bone vector
|
||||
tmp = inner_product(DIP_bones, DIP_Y_axis)
|
||||
DIP_theta_abduction = torch.where(tmp < 0, -DIP_theta_abduction, DIP_theta_abduction)
|
||||
|
||||
temp_axis = normalize(cross_product(DIP_Z_axis, DIP_bones))
|
||||
temp_alpha = angle_between(DIP_Z_axis, DIP_bones) # alpha belongs to [pi/2, pi]
|
||||
temp_R = axangle2mat_torch(axis=temp_axis, angle=temp_alpha, is_normalized=True)
|
||||
|
||||
# TIP bones
|
||||
TIP_X_axis = torch.matmul(temp_R, DIP_X_axis.unsqueeze(3)).squeeze()
|
||||
TIP_Y_axis = torch.matmul(temp_R, DIP_Y_axis.unsqueeze(3)).squeeze()
|
||||
TIP_bones_xz = TIP_bones - inner_product(TIP_bones, TIP_Y_axis).unsqueeze(2) * TIP_Y_axis
|
||||
|
||||
TIP_theta_flexion = angle_between(TIP_bones_xz, TIP_Z_axis) # in global coordinate
|
||||
TIP_theta_abduction = angle_between(TIP_bones_xz, TIP_bones) # in global coordinate
|
||||
# x-component of the bone vector
|
||||
tmp = inner_product(TIP_bones, TIP_X_axis)
|
||||
TIP_theta_flexion = torch.where(tmp < 1e-6, -TIP_theta_flexion, TIP_theta_flexion)
|
||||
# y-component of the bone vector
|
||||
tmp = inner_product(TIP_bones, TIP_Y_axis)
|
||||
TIP_theta_abduction = torch.where(tmp < 1e-6, -TIP_theta_abduction, TIP_theta_abduction)
|
||||
|
||||
# ALL
|
||||
ALL_theta_flexion = torch.cat((PIP_theta_flexion, DIP_theta_flexion, TIP_theta_flexion), dim=-1)
|
||||
ALL_theta_abduction = torch.cat((PIP_theta_abduction, DIP_theta_abduction, TIP_theta_abduction), dim=-1)
|
||||
ALL_theta = torch.stack((ALL_theta_flexion, ALL_theta_abduction), dim=-1)
|
||||
|
||||
a_loss = calculate_joint_angle_loss(ALL_theta, self.joint_angle_limit)
|
||||
final_loss += self.lambda_a * a_loss
|
||||
|
||||
BMC_losses["bmc_a"] = a_loss
|
||||
|
||||
return final_loss, BMC_losses
|
||||
+180
@@ -0,0 +1,180 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from body_model import OP_EDGE_LIST
|
||||
from geometry.rotation import batch_rodrigues
|
||||
from geometry.plane import (
|
||||
parse_floor_plane,
|
||||
compute_plane_intersection,
|
||||
fit_plane,
|
||||
)
|
||||
|
||||
|
||||
def compute_world2prior(floor_plane, trans, root_orient, origin):
|
||||
"""
|
||||
Computes rotation and translation from the camera frame to the canonical coordinate system
|
||||
used by the motion and initial state priors.
|
||||
- floor_plane : B x 3
|
||||
- trans : B x 3
|
||||
- root_orient : B x 3
|
||||
- origin: B x 3 desired origin (first joint of human model)
|
||||
- returns R (B, 3, 3), t (B, 3), root_height (B)
|
||||
"""
|
||||
B = trans.size(0)
|
||||
if floor_plane.size(1) == 3:
|
||||
floor_plane_4d = parse_floor_plane(floor_plane)
|
||||
else:
|
||||
floor_plane_4d = floor_plane
|
||||
floor_normal = floor_plane_4d[:, :3]
|
||||
floor_trans, _ = compute_plane_intersection(trans, -floor_normal, floor_plane_4d)
|
||||
|
||||
# compute prior frame axes within the camera frame
|
||||
# up is the floor_plane normal
|
||||
up_axis = floor_normal
|
||||
# right is body -x direction projected to floor plane
|
||||
root_orient_mat = batch_rodrigues(root_orient)
|
||||
body_right = -root_orient_mat[:, :, 0]
|
||||
floor_body_right, s = compute_plane_intersection(trans, body_right, floor_plane_4d)
|
||||
right_axis = floor_body_right - floor_trans
|
||||
# body right may not actually intersect - in this case must negate axis because we have the -x
|
||||
right_axis = torch.where(s.reshape((B, 1)) < 0, -right_axis, right_axis)
|
||||
right_axis = right_axis / torch.norm(right_axis, dim=1, keepdim=True)
|
||||
# forward is their cross product
|
||||
fwd_axis = torch.linalg.cross(up_axis, right_axis)
|
||||
fwd_axis = fwd_axis / torch.norm(fwd_axis, dim=1, keepdim=True)
|
||||
|
||||
# prior frame is right, fwd, up
|
||||
prior_R = torch.stack([right_axis, fwd_axis, up_axis], dim=2)
|
||||
world2prior_R = prior_R.transpose(-1, -2)
|
||||
|
||||
# translation takes translation to origin plus offset to the floor
|
||||
world2prior_t = -trans
|
||||
|
||||
# compute the distance from root to ground plane
|
||||
_, s_root = compute_plane_intersection(origin, -floor_normal, floor_plane_4d)
|
||||
root_height = s_root.reshape(B, 1)
|
||||
|
||||
if False:
|
||||
# apply transform to origin
|
||||
origin_prior = torch.matmul(world2prior_R, (origin - trans)[..., None])[..., 0]
|
||||
root_height = s_root.reshape(B, 1) - origin_prior[:, 2:3]
|
||||
|
||||
return world2prior_R, world2prior_t, root_height
|
||||
|
||||
|
||||
def find_cliques(edges):
|
||||
cliques = []
|
||||
N = edges.shape[0]
|
||||
empty = np.ones(N, dtype=bool)
|
||||
while empty.sum() > 0:
|
||||
idcs = np.where(empty)[0]
|
||||
i = idcs[0]
|
||||
cur_clique = [i]
|
||||
for j in idcs[1:]:
|
||||
if all(edges[c, j] for c in cur_clique):
|
||||
cur_clique.append(j)
|
||||
empty[cur_clique] = False
|
||||
cliques.append(cur_clique)
|
||||
return cliques
|
||||
|
||||
|
||||
def estimate_floor_planes(
|
||||
smpl_joints, valid_mask, thresh=0.5, group=False, flatten=False
|
||||
):
|
||||
"""
|
||||
:param smpl_joints (B, T, J, 3)
|
||||
:param valid_mask (B, T)
|
||||
:param thresh (optional, default 0.5) threshold for separating
|
||||
:param group (optional, default False) whether to group the planes
|
||||
:param flatten (optional, default False) whether to estimate one plane for everyone
|
||||
returns floor_plane (C, 3), floor_idcs (B,)
|
||||
"""
|
||||
B, T = valid_mask.shape
|
||||
device = valid_mask.device
|
||||
feet_joints = smpl_joints[..., 10:12, :] # (B, T, 2, 3)
|
||||
points = [feet_joints[b, valid_mask[b]].reshape(-1, 3) for b in range(B)]
|
||||
|
||||
if flatten:
|
||||
planes = fit_plane(torch.cat(points, dim=0))[None]
|
||||
labels = torch.zeros(B, device=device, dtype=torch.long)
|
||||
return planes[..., :3] * planes[..., 3:], labels
|
||||
|
||||
# estimate floors independently
|
||||
planes = torch.stack([fit_plane(pts) for pts in points], dim=0) # (B, 4)
|
||||
normals, offsets = planes[..., :3], planes[..., 3:] # (B, 3), (B, 1)
|
||||
sgn = torch.sign(offsets)
|
||||
normals, offsets = sgn * normals, sgn * offsets
|
||||
planes = normals * offsets # (B, 3)
|
||||
labels = torch.arange(B, device=device)
|
||||
|
||||
if not group:
|
||||
return planes, labels # (B, 3), (B,)
|
||||
|
||||
# asymetric, project each plane vector onto the other normals
|
||||
diffs = offsets - (planes[None] * normals[:, None]).sum(dim=-1)
|
||||
print(diffs)
|
||||
edges = diffs < thresh
|
||||
edges = edges.T & edges
|
||||
clusters = find_cliques(edges) # list of indices of clusters
|
||||
planes = []
|
||||
labels = torch.zeros(B, dtype=torch.long, device=device)
|
||||
for c, idcs in enumerate(clusters):
|
||||
plane = fit_plane(torch.cat([points[i] for i in idcs], dim=0))
|
||||
planes.append(plane)
|
||||
labels[idcs] = c
|
||||
planes = torch.stack(planes, dim=0) # (C, 4)
|
||||
return planes[..., :3] * planes[..., 3:], labels
|
||||
|
||||
|
||||
def estimate_initial_trans(body_pose, joints3d_op, joints2d_op, focal):
|
||||
"""
|
||||
use focal length and bone lengths to approximate distance from camera
|
||||
(based on PROX https://github.com/mohamedhassanmus/prox/blob/master/prox/fitting.py)
|
||||
:param body_pose
|
||||
:param joints3d_op OP keypoints in 3d
|
||||
:param joints2d_op OP keypoints in 2d
|
||||
"""
|
||||
B, T, N, _ = joints2d_op.shape
|
||||
joints2d_obs = joints2d_op[:, :, :, :2] # (B, T, N, 3)
|
||||
joints2d_conf = joints2d_op[:, :, :, 2]
|
||||
|
||||
# find least-occluded 2d frame
|
||||
num_2d_vis = torch.sum(joints2d_conf > 0.0, dim=2)
|
||||
best_2d_idx = torch.max(num_2d_vis, dim=1)[1]
|
||||
|
||||
# calculate bone lengths and confidence in each bone length
|
||||
bone3d = []
|
||||
bone2d = []
|
||||
conf2d = []
|
||||
for pair in OP_EDGE_LIST:
|
||||
diff3d = torch.norm(
|
||||
joints3d_op[:, 0, pair[0], :] - joints3d_op[:, 0, pair[1], :],
|
||||
dim=1,
|
||||
) # does not change over time
|
||||
diff2d = torch.norm(
|
||||
joints2d_obs[:, :, pair[0], :] - joints2d_obs[:, :, pair[1], :],
|
||||
dim=2,
|
||||
)
|
||||
minconf2d = torch.min(
|
||||
joints2d_conf[:, :, pair[0]], joints2d_conf[:, :, pair[1]]
|
||||
)
|
||||
bone3d.append(diff3d)
|
||||
bone2d.append(diff2d)
|
||||
conf2d.append(minconf2d)
|
||||
|
||||
bone3d = torch.stack(bone3d, dim=1)
|
||||
bone2d = torch.stack(bone2d, dim=2)
|
||||
bone2d = bone2d[np.arange(B), best_2d_idx, :]
|
||||
conf2d = torch.stack(conf2d, dim=2)
|
||||
conf2d = conf2d[np.arange(B), best_2d_idx, :]
|
||||
|
||||
# mean over all bones
|
||||
mean_bone3d = torch.mean(bone3d, dim=1)
|
||||
mean_bone2d = torch.mean(bone2d * (conf2d > 0.0), dim=1)
|
||||
|
||||
# approx z based on ratio
|
||||
init_z = focal * (mean_bone3d / mean_bone2d)
|
||||
init_trans = torch.zeros(B, T, 3)
|
||||
init_trans[:, :, 2] = init_z.unsqueeze(1).expand(B, T).detach()
|
||||
|
||||
return init_trans
|
||||
+1456
File diff suppressed because it is too large
Load Diff
+727
@@ -0,0 +1,727 @@
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from body_model import MANO_JOINTS, KEYPT_VERTS, smpl_to_openpose
|
||||
from geometry.rotation import batch_rodrigues, rotation_matrix_to_angle_axis
|
||||
from util.logger import Logger
|
||||
from util.tensor import (
|
||||
detach_all,
|
||||
scatter_intervals,
|
||||
select_intervals,
|
||||
get_scatter_mask,
|
||||
)
|
||||
|
||||
from .base_scene import BaseSceneModel
|
||||
from .helpers import compute_world2prior, estimate_floor_planes
|
||||
|
||||
|
||||
J_BODY = len(MANO_JOINTS) - 1 # no root
|
||||
CONTACT_ORDERING = [
|
||||
"hips",
|
||||
"leftLeg",
|
||||
"rightLeg",
|
||||
"leftFoot",
|
||||
"rightFoot",
|
||||
"leftToeBase",
|
||||
"rightToeBase",
|
||||
"leftHand",
|
||||
"rightHand",
|
||||
]
|
||||
CONTACT_INDS = [MANO_JOINTS[jname] for jname in CONTACT_ORDERING]
|
||||
CONTACT_THRESH = 0.5
|
||||
|
||||
|
||||
class MovingSceneModel(BaseSceneModel):
|
||||
"""
|
||||
Scene model of moving people in a shared global reference frame
|
||||
|
||||
Parameters:
|
||||
batch_size: number of sequences to optimize
|
||||
seq_len: length of the sequences
|
||||
pose_prior: VPoser model
|
||||
motion_prior: humor model
|
||||
init_motion_prior: dict of GMM params to use for prior on initial motion state
|
||||
shared_floor: (default True) if true, sequences are in the same reference frame
|
||||
fit_gender: (optional) gender of SMPL model
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
motion_prior,
|
||||
init_motion_prior=None,
|
||||
fit_gender="neutral",
|
||||
use_init=False,
|
||||
opt_cams=False,
|
||||
opt_scale=True,
|
||||
cam_graph=False,
|
||||
est_floor=True,
|
||||
floor_type="shared",
|
||||
async_tracks=True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
batch_size,
|
||||
seq_len,
|
||||
body_model,
|
||||
pose_prior,
|
||||
fit_gender=fit_gender,
|
||||
use_init=use_init,
|
||||
opt_cams=opt_cams,
|
||||
opt_scale=opt_scale,
|
||||
cam_graph=cam_graph,
|
||||
)
|
||||
assert motion_prior is not None
|
||||
assert motion_prior.model_data_config in [
|
||||
"mano+joints",
|
||||
"mano+joints+contacts",
|
||||
], "Only mano+joints motion prior configuration is supported!"
|
||||
|
||||
self.motion_prior = motion_prior
|
||||
self.init_motion_prior = init_motion_prior
|
||||
|
||||
# need latent dynamics sequence as well
|
||||
self.latent_motion_dim = self.motion_prior.latent_size
|
||||
self.cond_prior = self.motion_prior.use_conditional_prior
|
||||
assert self.cond_prior
|
||||
|
||||
# the frame chosen to use for the initial state (first frame by default)
|
||||
self.async_tracks = async_tracks
|
||||
self.register_buffer("track_start", torch.zeros(self.batch_size))
|
||||
self.register_buffer("track_end", torch.ones(self.batch_size) * self.seq_len)
|
||||
|
||||
self.shared_floor = floor_type == "shared"
|
||||
self.group_floor = floor_type == "group"
|
||||
self.floor_type = floor_type
|
||||
print("FLOOR TYPE", floor_type)
|
||||
self.est_floor = est_floor
|
||||
|
||||
@property
|
||||
def is_motion_active(self):
|
||||
return hasattr(self.params, "latent_motion")
|
||||
|
||||
def initialize(self, obs_data, cam_data, param_dict, data_fps):
|
||||
"""
|
||||
we need to also optimize for floor and world scale
|
||||
"""
|
||||
Logger.log("Initializing moving scene model with observed data")
|
||||
|
||||
self.params.set_cameras(
|
||||
cam_data,
|
||||
opt_scale=self.opt_scale,
|
||||
opt_cams=self.opt_cams,
|
||||
opt_focal=self.opt_cams,
|
||||
**param_dict,
|
||||
)
|
||||
# self.init_floor(obs_data, param_dict)
|
||||
self.init_first_state(obs_data, param_dict, data_fps)
|
||||
|
||||
# def init_floor(self, obs_data, param_dict):
|
||||
# if self.est_floor or self.group_floor:
|
||||
# with torch.no_grad():
|
||||
# smpl_preds = self.pred_mano(
|
||||
# param_dict["trans"],
|
||||
# param_dict["root_orient"],
|
||||
# self.latent2pose(param_dict["latent_pose"]),
|
||||
# param_dict["betas"],
|
||||
# )
|
||||
# floor_plane, floor_idcs = estimate_floor_planes(
|
||||
# smpl_preds["joints3d"].detach(),
|
||||
# obs_data["vis_mask"] > 0,
|
||||
# group=self.group_floor,
|
||||
# flatten=self.shared_floor,
|
||||
# )
|
||||
# if self.group_floor and not self.est_floor:
|
||||
# # don't use the estimated floor as initial plane
|
||||
# floor_plane = obs_data["floor_plane"][: len(floor_plane)]
|
||||
# else: # fixed shared or separate floors
|
||||
# num_floors = 1 if self.shared_floor else self.batch_size
|
||||
# floor_plane = obs_data["floor_plane"][:num_floors]
|
||||
# if self.shared_floor:
|
||||
# floor_idcs = torch.zeros(
|
||||
# self.batch_size, dtype=torch.long, device=floor_plane.device
|
||||
# )
|
||||
# else:
|
||||
# floor_idcs = torch.arange(
|
||||
# self.batch_size, dtype=torch.long, device=floor_plane.device
|
||||
# )
|
||||
|
||||
# Logger.log(f"ESTIMATED FLOORS: {str(floor_plane.detach().cpu())}")
|
||||
# Logger.log(f"FLOOR IDCS: {str(floor_idcs.detach().cpu())}")
|
||||
# self.params.set_param("floor_plane", floor_plane.float().detach())
|
||||
# self.params.set_param(
|
||||
# "floor_idcs", floor_idcs.long().detach(), requires_grad=False
|
||||
# )
|
||||
|
||||
def init_first_state(self, obs_data, param_dict, data_fps):
|
||||
"""
|
||||
initialize the latent motion and first state of each track
|
||||
using per-frame trajectories of trans, rot, latent_pose, betas
|
||||
and observed data
|
||||
"""
|
||||
B, T = self.batch_size, self.seq_len
|
||||
param_dict = detach_all(param_dict)
|
||||
|
||||
# select the valid segments of each track; pad the shorter tracks with final element
|
||||
if self.async_tracks and "track_interval" in obs_data:
|
||||
print("asynchronous tracks")
|
||||
interval = obs_data["track_interval"] # (B, 2)
|
||||
start, end = interval[:, 0], interval[:, 1]
|
||||
self.track_start, self.track_end = start, end
|
||||
param_dict = select_dict_segments(
|
||||
param_dict, start, end, names=["trans", "root_orient", "latent_pose"]
|
||||
)
|
||||
|
||||
trans = param_dict["trans"] # (B, T, 3)
|
||||
root_orient = param_dict["root_orient"] # (B, T, 3)
|
||||
latent_pose = param_dict["latent_pose"] # (B, T, D)
|
||||
betas = param_dict["betas"] # (B, b)
|
||||
|
||||
Logger.log(f"SEL TRACKS {trans.shape}, {root_orient.shape}")
|
||||
|
||||
# save each track's first appearance
|
||||
self.params.set_param("trans", trans[:, :1])
|
||||
Logger.log(f"INITIAL TRANS {trans[:, :1].detach().cpu()}")
|
||||
self.params.set_param("root_orient", root_orient[:, :1])
|
||||
self.params.set_param("latent_pose", latent_pose[:, :1])
|
||||
self.params.set_param("betas", betas)
|
||||
|
||||
# pass the current pose estimates through the motion prior
|
||||
self.data_fps = data_fps
|
||||
body_pose = self.latent2pose(latent_pose)
|
||||
init_latent = self.infer_latent_motion(
|
||||
trans, root_orient, body_pose, betas, data_fps
|
||||
).detach() # (B, T-1, D)
|
||||
self.params.set_param("latent_motion", init_latent)
|
||||
|
||||
# estimate velocities in prior frame and save initial state
|
||||
trans_vel, joints_vel, root_orient_vel = self.estimate_prior_velocities(
|
||||
trans, root_orient, body_pose, betas, data_fps
|
||||
)
|
||||
self.params.set_param("trans_vel", trans_vel[:, :1].detach())
|
||||
self.params.set_param("joints_vel", joints_vel[:, :1].detach())
|
||||
self.params.set_param("root_orient_vel", root_orient_vel[:, :1].detach())
|
||||
|
||||
def get_optim_result(self, num_steps=-1):
|
||||
res = super().get_optim_result()
|
||||
if not self.is_motion_active:
|
||||
return res
|
||||
|
||||
num_steps = self.seq_len if num_steps < 0 else num_steps
|
||||
prior_res, world_res = self.rollout_latent_motion(
|
||||
self.params.latent_motion[:, : num_steps - 1]
|
||||
)
|
||||
world_res = self.synchronize_preds(world_res, num_steps)
|
||||
res["world"].update(world_res)
|
||||
|
||||
prior_res = self.synchronize_preds(prior_res, num_steps)
|
||||
res["prior"] = prior_res
|
||||
return res
|
||||
|
||||
def estimate_prior_velocities(self, trans, root_orient, body_pose, betas, data_fps):
|
||||
"""
|
||||
compute velocities of trajectory in prior frame (not world frame)
|
||||
Transforms the first of each track into the prior frame
|
||||
:param trans (B, T, 3)
|
||||
:param root_orient (B, T, 3)
|
||||
:param body_pose (B, T, Dp)
|
||||
:param betas (B, Db)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
mano_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
return estimate_velocities(
|
||||
trans, root_orient, mano_results["joints3d"], data_fps
|
||||
)
|
||||
|
||||
def update_world2prior(self, trans, root_orient, body_pose, betas):
|
||||
B = trans.shape[0]
|
||||
mano_data = self.pred_mano(
|
||||
trans[:, :1], root_orient[:, :1], body_pose[:, :1], betas
|
||||
)
|
||||
floor_plane = self.params.floor_plane[self.params.floor_idcs]
|
||||
R, t, height = compute_world2prior(
|
||||
floor_plane,
|
||||
trans[:, 0],
|
||||
root_orient[:, 0],
|
||||
mano_data["joints3d"][:, 0, 0],
|
||||
)
|
||||
self.world2prior_R = R # (B, 3, 3)
|
||||
self.world2prior_t = t # (B, 3)
|
||||
self.world2prior_root_height = height # (B)
|
||||
|
||||
def apply_world2prior(self, trans, root_orient, body_pose, betas, inverse=False):
|
||||
"""
|
||||
Applies the world2prior tranformation to trans, root_orient
|
||||
If T=1, this function assumes they are at key_frame_idx,
|
||||
which we need to compute the offset from the origin
|
||||
:param trans (B, T, 3)
|
||||
:param root_orient (B, T, 3)
|
||||
:param body_pose (B, T, J, 3)
|
||||
:param betas (B, b)
|
||||
:param inverse (bool) optional, default False
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
# (B, 3, 3), (B, 3), (B)
|
||||
R, t, root_height = (
|
||||
self.world2prior_R,
|
||||
self.world2prior_t,
|
||||
self.world2prior_root_height,
|
||||
)
|
||||
R_time = R.unsqueeze(1).expand((B, T, 3, 3))
|
||||
t_time = t.unsqueeze(1).expand((B, T, 3))
|
||||
root_orient_mat = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
B, T, 3, 3
|
||||
)
|
||||
|
||||
if inverse:
|
||||
R_time = R_time.transpose(-1, -2)
|
||||
|
||||
root_orient_mat = torch.matmul(R_time, root_orient_mat)
|
||||
root_orient = rotation_matrix_to_angle_axis(
|
||||
root_orient_mat.reshape(B * T, 3, 3)
|
||||
).reshape(B, T, 3)
|
||||
|
||||
if inverse:
|
||||
# transform so first frame is at origin
|
||||
trans = trans - trans[:, :1, :]
|
||||
|
||||
# rotates to world frame
|
||||
trans = torch.matmul(R_time, trans[..., None])[..., 0]
|
||||
# translate to world frame
|
||||
trans = trans - t_time
|
||||
|
||||
return trans, root_orient
|
||||
|
||||
# first transform so the trans of key frame is at origin
|
||||
trans = trans + t_time
|
||||
# then rotate to canonical frame
|
||||
trans = torch.matmul(R_time, trans[..., None])[..., 0]
|
||||
# compute the root height after transforming
|
||||
cur_smpl_data = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
cur_root_height = cur_smpl_data["joints3d"][:, 0, 0, 2:3]
|
||||
# then apply floor offset so the root joint is at root_height
|
||||
height_diff = root_height - cur_root_height
|
||||
trans_offset = torch.cat(
|
||||
[torch.zeros((B, 2)).to(height_diff), height_diff], axis=1
|
||||
)
|
||||
trans = trans + trans_offset.reshape((B, 1, 3))
|
||||
|
||||
return trans, root_orient
|
||||
|
||||
def convert_prior_rot_inputs(self, root_orient, body_pose):
|
||||
# convert rots
|
||||
# body pose and root orient are both in aa
|
||||
B, T = root_orient.shape[:2]
|
||||
if (
|
||||
self.motion_prior.in_rot_rep == "mat"
|
||||
or self.motion_prior.in_rot_rep == "6d"
|
||||
):
|
||||
root_orient_in = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
(B, T, 9)
|
||||
)
|
||||
body_pose_in = batch_rodrigues(body_pose.reshape(-1, 3)).reshape(
|
||||
(B, T, J_BODY * 9)
|
||||
)
|
||||
if self.motion_prior.in_rot_rep == "6d":
|
||||
root_orient_in = root_orient[:, :, :6]
|
||||
body_pose_in = body_pose.reshape((B, T, J_BODY, 9))[:, :, :, :6].reshape(
|
||||
(B, T, J_BODY * 6)
|
||||
)
|
||||
return root_orient_in, body_pose_in
|
||||
|
||||
def convert_prior_rot_outputs(self, out_dict):
|
||||
keys = ["root_orient", "pose_body"]
|
||||
for k in keys:
|
||||
out = out_dict[k]
|
||||
B, T = out.shape[:2]
|
||||
out_dict[k] = rotation_matrix_to_angle_axis(
|
||||
out.reshape((-1, 3, 3))
|
||||
).reshape((B, T, -1))
|
||||
return out_dict
|
||||
|
||||
def infer_latent_motion(
|
||||
self, trans, root_orient, body_pose, betas, data_fps, full_forward_pass=False
|
||||
):
|
||||
"""
|
||||
By default, gets a sequence of z's from the current SMPL optim params.
|
||||
|
||||
If full_forward_pass is true, also samples from the posterior and feeds
|
||||
through the motion prior decoder to get all terms needed to calculate the ELBO.
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
h = 1.0 / data_fps
|
||||
|
||||
# need to first transform into canonical coordinate frame
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
|
||||
mano_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
joints = mano_results["joints3d"] # (B, T, len(MANO_JOINTS), 3)
|
||||
trans_vel, joints_vel, root_orient_vel = estimate_velocities(
|
||||
trans, root_orient, joints, data_fps
|
||||
)
|
||||
|
||||
root_orient_in, body_pose_in = self.convert_prior_rot_inputs(
|
||||
root_orient, body_pose
|
||||
)
|
||||
joints_in = joints.reshape((B, T, -1))
|
||||
joints_vel_in = joints_vel.reshape((B, T, -1))
|
||||
|
||||
seq_dict = {
|
||||
"trans": trans,
|
||||
"trans_vel": trans_vel,
|
||||
"root_orient": root_orient_in,
|
||||
"root_orient_vel": root_orient_vel,
|
||||
"pose_body": body_pose_in,
|
||||
"joints": joints_in,
|
||||
"joints_vel": joints_vel_in,
|
||||
}
|
||||
|
||||
infer_results = self.motion_prior.infer_global_seq(
|
||||
seq_dict, full_forward_pass=full_forward_pass
|
||||
)
|
||||
if full_forward_pass:
|
||||
# return both the given motion and the one from the forward pass
|
||||
# make sure rotations are matrix
|
||||
# NOTE: assumes seq_dict is same thing we want to compute loss on
|
||||
# Need to change if multiple future steps.
|
||||
if self.motion_prior.in_rot_rep != "mat":
|
||||
seq_dict["trans"] = batch_rodrigues(root_orient.reshape(-1, 3)).reshape(
|
||||
(B, T, 9)
|
||||
)
|
||||
seq_dict["pose_body"] = batch_rodrigues(
|
||||
body_pose.reshape(-1, 3)
|
||||
).reshape((B, T, J_BODY * 9))
|
||||
# do not need initial step anymore since output will be T-1
|
||||
for k, v in seq_dict.items():
|
||||
seq_dict[k] = v[:, 1:]
|
||||
for k in infer_results.keys():
|
||||
if k != "posterior_distrib" and k != "prior_distrib":
|
||||
infer_results[k] = infer_results[k][
|
||||
:, :, 0
|
||||
] # only want first output step
|
||||
infer_results = (seq_dict, infer_results)
|
||||
else:
|
||||
prior_z, posterior_z = infer_results
|
||||
infer_results = posterior_z[0] # mean of the approximate posterior
|
||||
|
||||
return infer_results
|
||||
|
||||
def rollout_mano_steps(self, num_steps=-1):
|
||||
# rollout motion given initial state
|
||||
num_steps = self.seq_len if num_steps < 0 else num_steps
|
||||
latent_motion = self.params.latent_motion[:, : num_steps - 1]
|
||||
|
||||
# roll out tracks from their first appearances
|
||||
# NOTE rolled out tracks are unsynced in time
|
||||
res, world_res = self.rollout_latent_motion(
|
||||
latent_motion, return_prior=self.cond_prior
|
||||
)
|
||||
|
||||
# get the smpl predictions for the unsynced tracks
|
||||
preds = self.pred_mano(
|
||||
res["trans"], res["root_orient"], res["pose_body"], res["betas"]
|
||||
)
|
||||
world_preds = self.pred_mano(
|
||||
world_res["trans"],
|
||||
world_res["root_orient"],
|
||||
world_res["pose_body"],
|
||||
world_res["betas"],
|
||||
)
|
||||
|
||||
# pass along relevant results
|
||||
preds["joints3d_rollout"] = res["joints"]
|
||||
preds["latent_pose"] = self.pose2latent(res["pose_body"])
|
||||
preds["joints3d_init"] = preds["joints3d"][:, :1]
|
||||
|
||||
if "contacts" in res:
|
||||
preds["contacts"] = res["contacts"]
|
||||
if "contacts_conf" in res:
|
||||
preds["contacts_conf"] = res["contacts_conf"]
|
||||
|
||||
# synchronize the tracklets to the same timesteps
|
||||
preds = self.synchronize_preds(preds, num_steps)
|
||||
world_preds = self.synchronize_preds(world_preds, num_steps)
|
||||
|
||||
# pass along unsynchronized motion latents
|
||||
preds["latent_motion"] = latent_motion
|
||||
if "cond_prior" in res:
|
||||
preds["cond_prior"] = res["cond_prior"]
|
||||
|
||||
return preds, world_preds
|
||||
|
||||
def rollout_latent_motion(
|
||||
self,
|
||||
latent_motion,
|
||||
return_prior=False,
|
||||
return_vel=False,
|
||||
num_steps=-1,
|
||||
canonicalize_input=False,
|
||||
):
|
||||
"""
|
||||
From the stored initial state, rolls out a sequence from the latent_motion vector
|
||||
NOTE: the stored initial states are not at the same time step, but we return all
|
||||
predicted sequences starting from time 0.
|
||||
|
||||
If latent_motion is None, samples num_steps into the future sequence from the prior.
|
||||
using the mean of the prior rather than random samples.
|
||||
|
||||
If canonicalize_input is True, transform the initial state into the local canonical
|
||||
frame before roll out
|
||||
"""
|
||||
# get the first frame state
|
||||
# NOTE: first states do not occur at same time step
|
||||
trans, root_orient, betas = (
|
||||
self.params.trans,
|
||||
self.params.root_orient,
|
||||
self.params.betas,
|
||||
)
|
||||
trans_vel, root_orient_vel, joints_vel = (
|
||||
self.params.trans_vel,
|
||||
self.params.root_orient_vel,
|
||||
self.params.joints_vel,
|
||||
)
|
||||
body_pose = self.latent2pose(self.params.latent_pose)
|
||||
|
||||
B = trans.size(0)
|
||||
is_sampling = latent_motion is None
|
||||
Tm1 = num_steps if latent_motion is None else latent_motion.size(1)
|
||||
if is_sampling and Tm1 <= 0:
|
||||
Logger.log("num_steps must be positive to sample!")
|
||||
exit()
|
||||
|
||||
# need to first transform initial state into canonical coordinate frame
|
||||
self.update_world2prior(trans, root_orient, body_pose, betas)
|
||||
trans, root_orient = self.apply_world2prior(
|
||||
trans, root_orient, body_pose, betas
|
||||
)
|
||||
|
||||
smpl_results = self.pred_mano(trans, root_orient, body_pose, betas)
|
||||
joints = smpl_results["joints3d"] # (B, T, len(MANO_JOINTS), 3)
|
||||
|
||||
# update to correct rotations for input
|
||||
root_orient_in, body_pose_in = self.convert_prior_rot_inputs(
|
||||
root_orient, body_pose
|
||||
)
|
||||
joints_in = joints.reshape((B, 1, -1))
|
||||
joints_vel_in = joints_vel.reshape((B, 1, -1))
|
||||
|
||||
rollout_in_dict = {
|
||||
"trans": trans,
|
||||
"root_orient": root_orient_in,
|
||||
"pose_body": body_pose_in,
|
||||
"joints": joints_in,
|
||||
"trans_vel": trans_vel,
|
||||
"root_orient_vel": root_orient_vel,
|
||||
"joints_vel": joints_vel_in,
|
||||
}
|
||||
|
||||
roll_output = self.motion_prior.roll_out(
|
||||
None,
|
||||
rollout_in_dict,
|
||||
Tm1,
|
||||
z_seq=latent_motion,
|
||||
return_prior=return_prior,
|
||||
return_z=is_sampling,
|
||||
canonicalize_input=canonicalize_input,
|
||||
gender=[self.fit_gender] * B,
|
||||
betas=betas.reshape((B, 1, -1)),
|
||||
)
|
||||
|
||||
pred_dict = prior_out = None
|
||||
if return_prior:
|
||||
pred_dict, prior_out = roll_output
|
||||
else:
|
||||
pred_dict = roll_output
|
||||
|
||||
pred_dict = self.convert_prior_rot_outputs(pred_dict)
|
||||
|
||||
# concat with initial state
|
||||
trans_out = torch.cat([trans, pred_dict["trans"]], dim=1)
|
||||
root_orient_out = torch.cat([root_orient, pred_dict["root_orient"]], dim=1)
|
||||
body_pose_out = torch.cat([body_pose, pred_dict["pose_body"]], dim=1)
|
||||
joints_out = torch.cat(
|
||||
[joints, pred_dict["joints"].reshape((B, Tm1, -1, 3))], dim=1
|
||||
)
|
||||
out_dict = {
|
||||
"trans": trans_out,
|
||||
"root_orient": root_orient_out,
|
||||
"pose_body": body_pose_out,
|
||||
"joints": joints_out,
|
||||
"betas": betas,
|
||||
}
|
||||
if return_vel:
|
||||
out_dict["trans_vel"] = torch.cat(
|
||||
[trans_vel, pred_dict["trans_vel"]], dim=1
|
||||
)
|
||||
out_dict["root_orient_vel"] = torch.cat(
|
||||
[root_orient_vel, pred_dict["root_orient_vel"]], dim=1
|
||||
)
|
||||
out_dict["joints_vel"] = torch.cat(
|
||||
[joints_vel, pred_dict["joints_vel"].reshape((B, Tm1, -1, 3))],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
if return_prior: # return the mean and var of distribution stacked
|
||||
pm, pv = prior_out
|
||||
out_dict["cond_prior"] = torch.stack([pm, pv], dim=-1)
|
||||
|
||||
if self.motion_prior.model_data_config == "smpl+joints+contacts":
|
||||
pred_contacts = pred_dict["contacts"]
|
||||
# get binary classification
|
||||
contact_conf = torch.sigmoid(pred_contacts)
|
||||
pred_contacts = (contact_conf > CONTACT_THRESH).to(torch.float)
|
||||
# expand to full body
|
||||
full_contact_conf = torch.zeros((B, Tm1, len(MANO_JOINTS))).to(contact_conf)
|
||||
full_contact_conf[:, :, CONTACT_INDS] = (
|
||||
full_contact_conf[:, :, CONTACT_INDS] + contact_conf
|
||||
)
|
||||
full_contacts = torch.zeros((B, Tm1, len(MANO_JOINTS))).to(pred_contacts)
|
||||
full_contacts[:, :, CONTACT_INDS] = (
|
||||
full_contacts[:, :, CONTACT_INDS] + pred_contacts
|
||||
)
|
||||
# repeat first entry for t0
|
||||
full_contact_conf = torch.cat(
|
||||
[full_contact_conf[:, 0:1], full_contact_conf], dim=1
|
||||
)
|
||||
full_contacts = torch.cat([full_contacts[:, 0:1], full_contacts], dim=1)
|
||||
out_dict["contacts_conf"] = full_contact_conf
|
||||
out_dict["contacts"] = full_contacts
|
||||
|
||||
if is_sampling:
|
||||
out_dict["z"] = pred_dict["z"]
|
||||
|
||||
cam_dict = {
|
||||
"trans": out_dict["trans"],
|
||||
"root_orient": out_dict["root_orient"],
|
||||
"pose_body": out_dict["pose_body"],
|
||||
"betas": out_dict["betas"],
|
||||
}
|
||||
|
||||
# also must return trans and root orient in camera frame
|
||||
cam_dict["trans"], cam_dict["root_orient"] = self.apply_world2prior(
|
||||
out_dict["trans"],
|
||||
out_dict["root_orient"],
|
||||
out_dict["pose_body"],
|
||||
out_dict["betas"],
|
||||
inverse=True,
|
||||
)
|
||||
return out_dict, cam_dict
|
||||
|
||||
def synchronize_preds(self, pred_dict, seg_len):
|
||||
"""
|
||||
synchronize predictions in time, scatter predictions into [0, seq_len)
|
||||
"""
|
||||
if not self.async_tracks:
|
||||
# predictions are already synchronized
|
||||
return pred_dict, None
|
||||
|
||||
# return time-synchronized predictions
|
||||
start = self.track_start
|
||||
end = torch.clip(start + seg_len, max=self.track_end)
|
||||
T = end.max()
|
||||
pred_dict = scatter_dict_segments(pred_dict, start, end, T)
|
||||
pred_mask = get_scatter_mask(start, end, T)
|
||||
pred_dict["track_mask"] = pred_mask
|
||||
return pred_dict
|
||||
|
||||
|
||||
def scatter_dict_segments(data_dict, start, end, T=None, names=None):
|
||||
"""
|
||||
the tracks as they are stored are synchronized by time step
|
||||
Uses the start and end to synchronize output predictions
|
||||
"""
|
||||
sync_dict = data_dict.copy()
|
||||
min_len = (end - start).min()
|
||||
if names is None:
|
||||
names = data_dict.keys()
|
||||
for name in names:
|
||||
val = data_dict[name]
|
||||
if not isinstance(val, torch.Tensor) or val.ndim < 3 or val.shape[1] < min_len:
|
||||
continue
|
||||
sync_dict[name] = scatter_intervals(val, start, end, T)
|
||||
return sync_dict
|
||||
|
||||
|
||||
def select_dict_segments(data_dict, start, end, names=None):
|
||||
out_data_dict = data_dict.copy()
|
||||
min_len = (end - start).min()
|
||||
if names is None:
|
||||
names = data_dict.keys()
|
||||
for name in names:
|
||||
val = data_dict[name]
|
||||
if not isinstance(val, torch.Tensor) or val.ndim < 3 or val.shape[1] < min_len:
|
||||
continue
|
||||
# only select for time-series observations
|
||||
out_data_dict[name] = select_intervals(val, start, end)
|
||||
return out_data_dict
|
||||
|
||||
|
||||
def estimate_velocities(trans, root_orient, joints3d, data_fps):
|
||||
"""
|
||||
Estimates velocity inputs to the motion prior.
|
||||
- trans (B, T, 3) root translation
|
||||
- root_orient (B, T, 3) aa root orientation
|
||||
- joints3d (B, T, len(MANO_JOINTS), 3) joints3d of SMPL prediction
|
||||
"""
|
||||
B, T, _ = trans.size()
|
||||
h = 1.0 / data_fps
|
||||
trans_vel = estimate_linear_velocity(trans, h)
|
||||
joints_vel = estimate_linear_velocity(joints3d, h)
|
||||
root_orient_mat = batch_rodrigues(root_orient.reshape((-1, 3))).reshape(
|
||||
(B, T, 3, 3)
|
||||
)
|
||||
root_orient_vel = estimate_angular_velocity(root_orient_mat, h)
|
||||
return trans_vel, joints_vel, root_orient_vel
|
||||
|
||||
|
||||
def estimate_linear_velocity(data_seq, h):
|
||||
"""
|
||||
Given some batched data sequences of T timesteps in the shape (B, T, ...), estimates
|
||||
the velocity for the middle T-2 steps using a second order central difference scheme.
|
||||
The first and last frames are with forward and backward first-order
|
||||
differences, respectively
|
||||
- h : step size
|
||||
"""
|
||||
# first steps is forward diff (t+1 - t) / h
|
||||
init_vel = (data_seq[:, 1:2] - data_seq[:, :1]) / h
|
||||
# middle steps are second order (t+1 - t-1) / 2h
|
||||
middle_vel = (data_seq[:, 2:] - data_seq[:, 0:-2]) / (2 * h)
|
||||
# last step is backward diff (t - t-1) / h
|
||||
final_vel = (data_seq[:, -1:] - data_seq[:, -2:-1]) / h
|
||||
|
||||
vel_seq = torch.cat([init_vel, middle_vel, final_vel], dim=1)
|
||||
return vel_seq
|
||||
|
||||
|
||||
def estimate_angular_velocity(rot_seq, h):
|
||||
"""
|
||||
Given a batch of sequences of T rotation matrices, estimates angular velocity at T-2 steps.
|
||||
Input sequence should be of shape (B, T, ..., 3, 3)
|
||||
"""
|
||||
# see https://en.wikipedia.org/wiki/Angular_velocity#Calculation_from_the_orientation_matrix
|
||||
dRdt = estimate_linear_velocity(rot_seq, h)
|
||||
R = rot_seq
|
||||
RT = R.transpose(-1, -2)
|
||||
# compute skew-symmetric angular velocity tensor
|
||||
w_mat = torch.matmul(dRdt, RT)
|
||||
# pull out angular velocity vector
|
||||
# average symmetric entries
|
||||
w_x = (-w_mat[..., 1, 2] + w_mat[..., 2, 1]) / 2.0
|
||||
w_y = (w_mat[..., 0, 2] - w_mat[..., 2, 0]) / 2.0
|
||||
w_z = (-w_mat[..., 0, 1] + w_mat[..., 1, 0]) / 2.0
|
||||
w = torch.stack([w_x, w_y, w_z], axis=-1)
|
||||
return w
|
||||
+643
@@ -0,0 +1,643 @@
|
||||
import os
|
||||
import glob
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
|
||||
from body_model import OP_IGNORE_JOINTS
|
||||
from geometry.mesh import save_mesh_scenes, vertices_to_trimesh
|
||||
from util.logger import Logger, log_cur_stats
|
||||
from util.tensor import move_to, detach_all
|
||||
from vis.output import prep_result_vis, animate_scene
|
||||
|
||||
from .losses import RootLoss, SMPLLoss #, MotionLoss
|
||||
from .output import save_camera_json
|
||||
import time
|
||||
|
||||
LINE_SEARCH = "strong_wolfe"
|
||||
LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353)
|
||||
|
||||
"""
|
||||
Optimization happens in multiple stages
|
||||
Stage 1: fit root orients and trans of each frame indendpently
|
||||
Stage 2: fit SMPL poses and betas of each frame independently
|
||||
Stage 3: fit poses and roots in same global coordinate frame
|
||||
"""
|
||||
|
||||
|
||||
class StageOptimizer(object):
|
||||
def __init__(
|
||||
self,
|
||||
name,
|
||||
model,
|
||||
param_names,
|
||||
lr=1.0,
|
||||
lbfgs_max_iter=20,
|
||||
save_every=10,
|
||||
vis_every=-1,
|
||||
save_meshes=True,
|
||||
max_chunk_steps=10,
|
||||
**kwargs,
|
||||
):
|
||||
Logger.log(f"INITIALIZING OPTIMIZER {name} for {param_names}")
|
||||
self.name = name
|
||||
self.model = model
|
||||
|
||||
self.set_opt_vars(param_names)
|
||||
|
||||
self.optim = torch.optim.LBFGS(
|
||||
self.opt_params, max_iter=lbfgs_max_iter, lr=lr, line_search_fn=LINE_SEARCH
|
||||
)
|
||||
# LBFGS computes losses multiple times per iteration,
|
||||
# save a dict mapping iteration to list of stats_dicts
|
||||
self.loss_dicts = {}
|
||||
|
||||
self.save_every = save_every
|
||||
self.vis_every = vis_every
|
||||
self.save_meshes = save_meshes
|
||||
self.max_chunk_steps = max_chunk_steps
|
||||
|
||||
self.cur_step = 0
|
||||
|
||||
self.add_chunk = 0
|
||||
self.cur_loss = 0
|
||||
self.prev_loss = np.inf
|
||||
self.last_updated = 0
|
||||
self.reached_max = False
|
||||
self.reached_max_iter = -1
|
||||
|
||||
def set_opt_vars(self, param_names):
|
||||
Logger.log("Set param names:")
|
||||
Logger.log(param_names)
|
||||
|
||||
self.param_names = param_names
|
||||
self.model.params.set_require_grads(self.param_names)
|
||||
self.opt_params = [
|
||||
getattr(self.model.params, name) for name in self.param_names
|
||||
]
|
||||
|
||||
def forward_pass(self, obs_data):
|
||||
raise NotImplementedError
|
||||
|
||||
def load_checkpoint(self, out_dir, device=None):
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
|
||||
param_path = os.path.join(out_dir, f"{self.name}_params.pth")
|
||||
if os.path.isfile(param_path):
|
||||
param_dict = torch.load(param_path, map_location=device)
|
||||
self.model.params.load_dict(param_dict)
|
||||
Logger.log(f"Params loaded from {param_path}")
|
||||
|
||||
optim_path = os.path.join(out_dir, f"{self.name}_optim.pth")
|
||||
if os.path.isfile(optim_path):
|
||||
optim_dict = torch.load(optim_path)
|
||||
self.optim.load_state_dict(optim_dict["optim"])
|
||||
self.cur_step = optim_dict["cur_step"]
|
||||
Logger.log(f"Optimizer loaded from {optim_path} at iter {self.cur_step}")
|
||||
|
||||
def save_checkpoint(self, out_dir):
|
||||
param_path = os.path.join(out_dir, f"{self.name}_params.pth")
|
||||
param_dict = self.model.params.get_dict()
|
||||
if "world_scale" in param_dict:
|
||||
print("WORLD_SCALE", param_dict["world_scale"].detach().cpu())
|
||||
if "floor_plane" in param_dict:
|
||||
print("FLOOR PLANE", param_dict["floor_plane"].detach().cpu())
|
||||
if "cam_f" in param_dict:
|
||||
print("CAM_F", param_dict["cam_f"].detach().cpu())
|
||||
torch.save(param_dict, param_path)
|
||||
Logger.log(f"Model saved at {param_path}")
|
||||
|
||||
optim_path = os.path.join(out_dir, f"{self.name}_optim.pth")
|
||||
torch.save(
|
||||
{"optim": self.optim.state_dict(), "cur_step": self.cur_step},
|
||||
optim_path,
|
||||
)
|
||||
Logger.log(f"Optimizer saved at {optim_path}")
|
||||
|
||||
def save_results(self, out_dir, seq_name):
|
||||
"""
|
||||
pred dict will be a dictionary of trajectories.
|
||||
each trajectory will have params and lists of trimesh sequences
|
||||
"""
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
with torch.no_grad():
|
||||
print('go into get_optim_result from optimizer.py save_results')
|
||||
pred_dict = self.model.get_optim_result()
|
||||
pred_dict = move_to(detach_all(pred_dict), "cpu")
|
||||
|
||||
i = self.cur_step
|
||||
for name, results in pred_dict.items():
|
||||
print('save_results, name: ', name)
|
||||
# save parameters of trajectory
|
||||
out_path = f"{out_dir}/{seq_name}_{i:06d}_{name}_results.npz"
|
||||
Logger.log(f"saving params to {out_path}")
|
||||
np.savez(out_path, **results)
|
||||
print('cam_R in save_results and its .npz', results['cam_R'])
|
||||
|
||||
# also save the cameras
|
||||
# print('save the cameras in ptimizer.py save_results')
|
||||
# print(self.model.params)
|
||||
# with torch.no_grad():
|
||||
# cam_R, cam_t = self.model.params.get_extrinsics()
|
||||
# intrins = self.model.params.get_intrinsics()
|
||||
# print('cam_R in 0000.json:', cam_R)
|
||||
# save_camera_json(
|
||||
# f"{out_dir}/{seq_name}_cameras_{self.cur_step:06d}.json",
|
||||
# cam_R.detach().cpu(),
|
||||
# cam_t.detach().cpu(),
|
||||
# intrins.detach().cpu(),
|
||||
# )
|
||||
|
||||
# plot losses
|
||||
self.plot_losses(out_dir)
|
||||
|
||||
# def save_meshes_all(self, res_dir, obs_data, seq_name, num_steps=-1):
|
||||
|
||||
# # check which results are saved
|
||||
# seq_name = obs_data["seq_name"][0]
|
||||
# res_pre = f"{res_dir}/{seq_name}_opt_{self.cur_step:06d}"
|
||||
# with torch.no_grad():
|
||||
# print('go into get_optim_result from optimizer.py vis_result')
|
||||
# pred_dict = self.model.get_optim_result(num_steps=num_steps)
|
||||
|
||||
# res_dict = detach_all(pred_dict["world"])
|
||||
# scene_dict = move_to(
|
||||
# prep_result_vis(
|
||||
# res_dict,
|
||||
# obs_data["vis_mask"],
|
||||
# obs_data["track_id"],
|
||||
# self.model.body_model,
|
||||
# temporal_smooth=True
|
||||
# ),
|
||||
# "cpu",
|
||||
# )
|
||||
|
||||
# # save the meshes
|
||||
# os.makedirs(res_dir, exist_ok=True)
|
||||
# i = self.cur_step
|
||||
# scene_dir = f"{res_dir}/{seq_name}_{i:06d}_meshes"
|
||||
# os.makedirs(scene_dir, exist_ok=True)
|
||||
|
||||
# verts, joints, colors, l_faces, r_faces, is_right, bounds = scene_dict["geometry"]
|
||||
# T = len(verts)
|
||||
# print(f"{T} mesh frames")
|
||||
# times = list(range(0, T, 1))
|
||||
# for t in times:
|
||||
# if len(is_right[t]) > 1:
|
||||
# assert (is_right[t].cpu().numpy().tolist() == [0,1])
|
||||
# # l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
# # r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None])
|
||||
# # assert len(l_meshes) == 1
|
||||
# # assert len(r_meshes) == 1
|
||||
# tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
|
||||
# tmesh = vertices_to_trimesh(verts[t][1].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
|
||||
# else:
|
||||
# assert len(is_right[t]) == 1
|
||||
# if is_right[t] == 0:
|
||||
# # meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None])
|
||||
# tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj'))
|
||||
# elif is_right[t] == 1:
|
||||
# # meshes = make_batch_mesh(verts[t][0][None], r_faces[t], colors[t][0][None])
|
||||
# tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1)
|
||||
# tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj'))
|
||||
|
||||
def vis_result(self, res_dir, obs_data, vis=None, num_steps=-1):
|
||||
if vis is None or self.vis_every < 0:
|
||||
return
|
||||
|
||||
# check which results are saved
|
||||
seq_name = obs_data["seq_name"][0]
|
||||
res_pre = f"{res_dir}/{seq_name}_opt_{self.cur_step:06d}"
|
||||
with torch.no_grad():
|
||||
print('go into get_optim_result from optimizer.py vis_result')
|
||||
pred_dict = self.model.get_optim_result(num_steps=num_steps)
|
||||
|
||||
res_dict = detach_all(pred_dict["world"])
|
||||
scene_dict = move_to(
|
||||
prep_result_vis(
|
||||
res_dict,
|
||||
obs_data["vis_mask"],
|
||||
obs_data["track_id"],
|
||||
self.model.body_model,
|
||||
),
|
||||
"cpu",
|
||||
)
|
||||
animate_scene(vis, scene_dict, res_pre, render_views=["src_cam", "above"])
|
||||
|
||||
def log_losses(self, stats_dict):
|
||||
stats_dict = move_to(detach_all(stats_dict), "cpu")
|
||||
log_cur_stats(
|
||||
stats_dict,
|
||||
iter=self.cur_step,
|
||||
to_stdout=(self.cur_step % self.save_every == 0),
|
||||
)
|
||||
for loss_name, loss_val in stats_dict.items():
|
||||
loss_dict = self.loss_dicts.get(loss_name, {})
|
||||
loss_series = loss_dict.get(self.cur_step, [])
|
||||
loss_series.append(loss_val)
|
||||
loss_dict[self.cur_step] = loss_series
|
||||
self.loss_dicts[loss_name] = loss_dict
|
||||
|
||||
def record_current_losses(self, writer):
|
||||
"""
|
||||
record the mean of current step's loss values in tensorboard
|
||||
"""
|
||||
if len(self.loss_dicts) < 1:
|
||||
return
|
||||
|
||||
for loss_name, loss_dict in self.loss_dicts.items():
|
||||
loss_mean = np.mean(loss_dict[self.cur_step])
|
||||
writer.add_scalar(f"{self.name}/{loss_name}", loss_mean, self.cur_step)
|
||||
|
||||
def plot_losses(self, res_dir):
|
||||
"""
|
||||
plot a box plot for each BFGS iteration
|
||||
"""
|
||||
if len(self.loss_dicts) < 1:
|
||||
return
|
||||
for loss_name, loss_dict in self.loss_dicts.items():
|
||||
# times (list len T)
|
||||
# loss vals (list len T of loss value lists)
|
||||
times, loss_vals = zip(*loss_dict.items())
|
||||
plt.figure()
|
||||
plt.boxplot(loss_vals, labels=times, showfliers=False)
|
||||
plt.savefig(f"{res_dir}/{loss_name}.png")
|
||||
|
||||
def run(self, obs_data, num_iters, out_dir, vis=None, writer=None):
|
||||
self.cur_step = 0
|
||||
self.loss.cur_step = 0
|
||||
res_dir = os.path.join(out_dir, self.name)
|
||||
os.makedirs(res_dir, exist_ok=True)
|
||||
seq_name = obs_data["seq_name"][0]
|
||||
print("SEQ NAME", seq_name, "num_iters", num_iters)
|
||||
|
||||
# try to load from checkpoint if exists
|
||||
device = obs_data["joints2d"].device
|
||||
self.load_checkpoint(out_dir, device=device)
|
||||
|
||||
if self.cur_step >= num_iters:
|
||||
Logger.log(f"Current optimizer {self}")
|
||||
Logger.log(f"Checkpoint at {self.cur_step} >= {num_iters}, skipping")
|
||||
return
|
||||
|
||||
Logger.log(f"OPTIMIZING {self.name} FOR {num_iters} ITERATIONS")
|
||||
# time.sleep(5)
|
||||
|
||||
# save initial results and vis
|
||||
print('go into save_results in optimizers.py run()')
|
||||
self.save_results(res_dir, seq_name)
|
||||
|
||||
for i in range(self.cur_step, num_iters):
|
||||
Logger.log("ITER: %d" % (i))
|
||||
|
||||
if (i + 1) % self.save_every == 0: # save before
|
||||
self.save_checkpoint(out_dir)
|
||||
self.save_results(res_dir, seq_name)
|
||||
else:
|
||||
self.save_checkpoint(out_dir)
|
||||
|
||||
if (i + 1) % self.vis_every == 0: # render
|
||||
self.vis_result(res_dir, obs_data, vis)
|
||||
|
||||
self.cur_step = i
|
||||
self.loss.cur_step = i
|
||||
|
||||
#print("ITER: %d" % (i))
|
||||
self.optim_step(obs_data, i, writer)
|
||||
|
||||
# early termination in case of nans
|
||||
if np.isnan(self.cur_loss):
|
||||
# we need to backtrack
|
||||
self.load_checkpoint(out_dir, device=device)
|
||||
print(self.cur_loss, self.kkk, self.ppp)
|
||||
raise ValueError
|
||||
|
||||
# termination for the last chunk
|
||||
if self.reached_max and self.reached_max_iter < 0:
|
||||
self.reached_max_iter = i - 1
|
||||
if self.reached_max and i - self.reached_max_iter >= self.max_chunk_steps:
|
||||
break
|
||||
|
||||
# termination for middle chunks
|
||||
loss_change = self.prev_loss - self.cur_loss
|
||||
if self.last_updated == i - 1 and loss_change == 0:
|
||||
break
|
||||
if (
|
||||
(self.cur_loss < 0 and loss_change < 100)
|
||||
or (i - self.last_updated >= self.max_chunk_steps)
|
||||
or (loss_change < 20 and i - self.last_updated > 5)
|
||||
):
|
||||
self.add_chunk = self.add_chunk + 1
|
||||
self.last_updated = i
|
||||
self.prev_loss = self.cur_loss
|
||||
|
||||
# final save and vis step
|
||||
self.cur_step = num_iters
|
||||
self.save_checkpoint(out_dir)
|
||||
self.save_results(res_dir, seq_name)
|
||||
self.vis_result(res_dir, obs_data, vis)
|
||||
# self.save_meshes_all(res_dir, obs_data, seq_name)
|
||||
|
||||
def optim_step(self, obs_data, i, writer=None):
|
||||
def closure():
|
||||
self.optim.zero_grad()
|
||||
#print('******************************', i)
|
||||
loss, stats_dict, preds = self.forward_pass(obs_data)
|
||||
#print('******************************', i)
|
||||
stats_dict["total"] = loss
|
||||
self.log_losses(move_to(detach_all(stats_dict), "cpu"))
|
||||
self.cur_loss = stats_dict["total"].detach().cpu().item()
|
||||
self.kkk = stats_dict
|
||||
self.ppp = preds
|
||||
loss_keys = stats_dict.keys()
|
||||
if "motion_prior" in loss_keys and "pose_prior" in loss_keys:
|
||||
self.cur_loss = (
|
||||
self.cur_loss
|
||||
- 0.04 * stats_dict["pose_prior"].detach().cpu().item()
|
||||
)
|
||||
print("optimizer print", loss, loss.shape, stats_dict, loss.requires_grad)
|
||||
loss.backward()
|
||||
# print('end of loss backward')
|
||||
return loss
|
||||
|
||||
#print('!!!!!!!!!!!!!!!!!!!!!!!!!!', i)
|
||||
self.optim.step(closure)
|
||||
#print('!!!!!!!!!!!!!!!!!!!!!!!!!!', i)
|
||||
if writer is not None:
|
||||
self.record_current_losses(writer)
|
||||
|
||||
|
||||
class RootOptimizer(StageOptimizer):
|
||||
name = "root_fit"
|
||||
stage = 0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
all_loss_weights,
|
||||
use_chamfer=False,
|
||||
robust_loss_type="none",
|
||||
robust_tuning_const=4.6851,
|
||||
joints2d_sigma=100,
|
||||
**kwargs,
|
||||
):
|
||||
param_names = ["trans", "root_orient"]
|
||||
super().__init__(self.name, model, param_names, **kwargs)
|
||||
|
||||
self.loss = RootLoss(
|
||||
all_loss_weights[self.stage],
|
||||
ignore_op_joints=OP_IGNORE_JOINTS,
|
||||
joints2d_sigma=joints2d_sigma,
|
||||
use_chamfer=use_chamfer,
|
||||
robust_loss=robust_loss_type,
|
||||
robust_tuning_const=robust_tuning_const,
|
||||
faces=model.body_model.faces_tensor
|
||||
)
|
||||
|
||||
def forward_pass(self, obs_data):
|
||||
"""
|
||||
Takes in observed data, predicts the smpl parameters and returns loss
|
||||
"""
|
||||
# raise ValueError
|
||||
# Use current params to go through SMPL and get joints3d, verts3d, points3d
|
||||
pred_data = self.model.pred_params_mano(obs_data["is_right"])
|
||||
pred_data["cameras"] = self.model.params.get_cameras()
|
||||
|
||||
# compute data losses only
|
||||
vis_mask = obs_data["vis_mask"] >= 0
|
||||
# print('entering loss from RootOptimizer')
|
||||
loss, stats_dict = self.loss(obs_data, pred_data, vis_mask)
|
||||
return loss, stats_dict, pred_data
|
||||
|
||||
|
||||
class SmoothOptimizer(StageOptimizer):
|
||||
name = "smooth_fit"
|
||||
stage = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
all_loss_weights,
|
||||
use_chamfer=False,
|
||||
robust_loss_type="none",
|
||||
robust_tuning_const=4.6851,
|
||||
joints2d_sigma=100,
|
||||
**kwargs,
|
||||
):
|
||||
param_names = ["trans", "root_orient", "betas", "latent_pose"]
|
||||
if model.opt_scale:
|
||||
param_names += ["world_scale"]
|
||||
if model.opt_cams:
|
||||
param_names += ["cam_f", "delta_cam_R"]
|
||||
|
||||
super().__init__(self.name, model, param_names, **kwargs)
|
||||
|
||||
self.loss = SMPLLoss(
|
||||
all_loss_weights[self.stage],
|
||||
ignore_op_joints=OP_IGNORE_JOINTS,
|
||||
joints2d_sigma=joints2d_sigma,
|
||||
use_chamfer=use_chamfer,
|
||||
robust_loss=robust_loss_type,
|
||||
robust_tuning_const=robust_tuning_const,
|
||||
faces=model.body_model.faces_tensor
|
||||
)
|
||||
|
||||
def forward_pass(self, obs_data):
|
||||
"""
|
||||
Takes in observed data, predicts the smpl parameters and returns loss
|
||||
"""
|
||||
# raise ValueError
|
||||
# Use current params to go through SMPL and get joints3d, verts3d, points3d
|
||||
pred_data = self.model.pred_params_mano(obs_data["is_right"])
|
||||
pred_data["cameras"] = self.model.params.get_cameras()
|
||||
pred_data.update(self.model.params.get_vars())
|
||||
|
||||
# camera predictions
|
||||
# pts_ij, target, weight = self.model.params.get_reprojected_points()
|
||||
# pred_data["bg2d_err"] = (
|
||||
# (weight * (pts_ij - target) ** 2).sum(dim=(-1, -2)).mean()
|
||||
# )
|
||||
pred_data["cam_R"], pred_data["cam_t"] = self.model.params.get_extrinsics()
|
||||
|
||||
# compute data losses only
|
||||
vis_mask = obs_data["vis_mask"] >= 0
|
||||
loss, stats_dict = self.loss(obs_data, pred_data, self.model.seq_len, vis_mask)
|
||||
return loss, stats_dict, pred_data
|
||||
|
||||
|
||||
# class MotionOptimizer(StageOptimizer):
|
||||
# name = "motion_fit"
|
||||
# stage = 2
|
||||
|
||||
# def __init__(self, model, all_loss_weights, opt_cams=False, **kwargs):
|
||||
# self.opt_cams = opt_cams and model.opt_cams
|
||||
# param_names = [
|
||||
# "trans",
|
||||
# "root_orient",
|
||||
# "latent_pose",
|
||||
# "trans_vel",
|
||||
# "root_orient_vel",
|
||||
# "joints_vel",
|
||||
# "betas",
|
||||
# "latent_motion",
|
||||
# "floor_plane",
|
||||
# ]
|
||||
# if model.opt_scale:
|
||||
# param_names += ["world_scale"]
|
||||
|
||||
# if self.opt_cams:
|
||||
# Logger.log(f"{self.name} OPTIMIZING CAMERAS")
|
||||
# param_names += ["delta_cam_R", "cam_f"]
|
||||
|
||||
# super().__init__(self.name, model, param_names, **kwargs)
|
||||
# self.set_loss(model, all_loss_weights[self.stage], **kwargs)
|
||||
|
||||
# def set_loss(
|
||||
# self,
|
||||
# model,
|
||||
# loss_weights,
|
||||
# use_chamfer=False,
|
||||
# robust_loss_type="none",
|
||||
# robust_tuning_const=4.6851,
|
||||
# joints2d_sigma=100,
|
||||
# **kwargs,
|
||||
# ):
|
||||
# self.loss = MotionLoss(
|
||||
# loss_weights,
|
||||
# init_motion_prior=model.init_motion_prior,
|
||||
# ignore_op_joints=OP_IGNORE_JOINTS,
|
||||
# joints2d_sigma=joints2d_sigma,
|
||||
# use_chamfer=use_chamfer,
|
||||
# robust_loss=robust_loss_type,
|
||||
# robust_tuning_const=robust_tuning_const,
|
||||
# )
|
||||
|
||||
# def get_motion_scale(self):
|
||||
# return 1.0
|
||||
|
||||
# def forward_pass(self, obs_data, num_steps=-1):
|
||||
# p = self.model.params
|
||||
# param_names = [
|
||||
# "betas",
|
||||
# "joints_vel",
|
||||
# "trans_vel",
|
||||
# "root_orient_vel",
|
||||
# ]
|
||||
# param_dict = p.get_vars(param_names)
|
||||
# # param_dict["floor_plane"] = p.floor_plane[p.floor_idcs]
|
||||
|
||||
# preds, world_preds = self.model.rollout_mano_steps(num_steps)
|
||||
# preds.update(param_dict)
|
||||
# world_preds.update(param_dict)
|
||||
|
||||
# # track mask is the length of the sequence that contains num_steps of each track
|
||||
# track_mask = world_preds.get("track_mask", None)
|
||||
# T = track_mask.shape[1] if track_mask is not None else num_steps
|
||||
# world_preds["cameras"] = p.get_cameras(np.arange(T))
|
||||
# if self.opt_cams:
|
||||
# cam_R, cam_t = p.get_extrinsics()
|
||||
# world_preds["cam_R"], world_preds["cam_t"] = cam_R[:T], cam_t[:T]
|
||||
|
||||
# obs_data = slice_dict(obs_data, 0, T)
|
||||
# motion_scale = self.get_motion_scale()
|
||||
# loss, stats_dict = self.loss(
|
||||
# obs_data,
|
||||
# preds,
|
||||
# world_preds,
|
||||
# self.model.seq_len,
|
||||
# valid_mask=track_mask,
|
||||
# init_motion_scale=motion_scale,
|
||||
# )
|
||||
# return loss, stats_dict, (preds, world_preds)
|
||||
|
||||
|
||||
# def slice_dict(d, start, end):
|
||||
# out = d.copy()
|
||||
# for k, v in d.items():
|
||||
# if not isinstance(v, torch.Tensor) or v.ndim < 3 or v.shape[1] < end:
|
||||
# continue
|
||||
# out[k] = v[:, start:end]
|
||||
# return out
|
||||
|
||||
|
||||
# class MotionOptimizerChunks(MotionOptimizer):
|
||||
# """
|
||||
# Incrementally optimize sequence in chunks (with all variables)
|
||||
# :param chunk_size (int) number of frames per chunk
|
||||
# :param init_steps (int) number of optimization steps to spend on first chunk
|
||||
# :param chunk_steps (int) number of opt steps to spend on subsequent chunks
|
||||
# """
|
||||
|
||||
# name = "motion_chunks"
|
||||
# stage = 2
|
||||
|
||||
# def __init__(self, *args, chunk_size=20, init_steps=20, chunk_steps=10, **kwargs):
|
||||
# self.chunk_size = chunk_size
|
||||
# self.init_steps = init_steps
|
||||
# self.chunk_steps = chunk_steps
|
||||
# super().__init__(*args, **kwargs)
|
||||
|
||||
# @property
|
||||
# def num_iters(self):
|
||||
# num_chunks = int(np.ceil(self.model.seq_len / self.chunk_size)) + 1
|
||||
# return self.init_steps + self.chunk_steps * num_chunks
|
||||
|
||||
# @property
|
||||
# def end_idx(self):
|
||||
# if self.cur_step < self.init_steps and self.add_chunk == 0:
|
||||
# return self.chunk_size
|
||||
# chunk_idx = max(
|
||||
# (self.cur_step - self.init_steps) // self.chunk_steps + 1, self.add_chunk
|
||||
# )
|
||||
# if int(np.ceil(self.model.seq_len / self.chunk_size)) == self.add_chunk:
|
||||
# self.reached_max = True
|
||||
# return min(self.chunk_size * (chunk_idx + 1), self.model.seq_len)
|
||||
|
||||
# @property
|
||||
# def start_idx(self):
|
||||
# return 0
|
||||
|
||||
# def get_motion_scale(self):
|
||||
# num_frames = self.end_idx - self.start_idx
|
||||
# return max(1.0, float(self.model.seq_len) / num_frames)
|
||||
|
||||
# def forward_pass(self, obs_data):
|
||||
# return super().forward_pass(obs_data, num_steps=self.end_idx)
|
||||
|
||||
# def vis_result(self, res_dir, obs_data, vis=None, **kwargs):
|
||||
# print("start, end", self.start_idx, self.end_idx)
|
||||
# return super().vis_result(res_dir, obs_data, vis=vis, num_steps=self.end_idx)
|
||||
|
||||
|
||||
# class MotionOptimizerFreeze(MotionOptimizer):
|
||||
# """
|
||||
# Only optimizing a subset of parameters
|
||||
# """
|
||||
|
||||
# name = "motion_freeze"
|
||||
# stage = 2
|
||||
|
||||
# def __init__(self, model, all_loss_weights, no_contacts=True, **kwargs):
|
||||
# loss_weights = all_loss_weights[self.stage].copy()
|
||||
|
||||
# if no_contacts:
|
||||
# loss_weights["contact_height"] = 0.0
|
||||
# loss_weights["contact_vel"] = 0.0
|
||||
|
||||
# param_names = ["latent_motion", "betas", "trans", "floor_plane"]
|
||||
# if model.optim_scale:
|
||||
# param_names += ["world_scale"]
|
||||
|
||||
# StageOptimizer.__init__(self, self.name, model, param_names, **kwargs)
|
||||
# self.set_loss(model, loss_weights, **kwargs)
|
||||
+188
@@ -0,0 +1,188 @@
|
||||
import os
|
||||
import imageio
|
||||
import glob
|
||||
import json
|
||||
import subprocess
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from util.tensor import move_to, detach_all, to_torch
|
||||
from util.logger import Logger
|
||||
|
||||
|
||||
def get_results_paths(res_dir):
|
||||
"""
|
||||
get the iterations of all saved results in res_dir
|
||||
:param res_dir (str) result dir
|
||||
returns a dict of iter to result path
|
||||
"""
|
||||
res_files = sorted(glob.glob(f"{res_dir}/*_results.npz"))
|
||||
print(f"found {len(res_files)} results in {res_dir}")
|
||||
|
||||
path_dict = {}
|
||||
for res_file in res_files:
|
||||
it, name, _ = os.path.basename(res_file).split("_")[-3:]
|
||||
assert name in ["world", "prior"]
|
||||
if it not in path_dict:
|
||||
path_dict[it] = {}
|
||||
path_dict[it][name] = res_file
|
||||
return path_dict
|
||||
|
||||
|
||||
def load_result(res_path_dict):
|
||||
"""
|
||||
load all saved results for a given iteration
|
||||
:param res_path_dict (dict) paths to relevant results
|
||||
returns dict of results
|
||||
"""
|
||||
res_dict = {}
|
||||
for name, path in res_path_dict.items():
|
||||
res = np.load(path)
|
||||
res_dict[name] = to_torch({k: res[k] for k in res.files})
|
||||
return res_dict
|
||||
|
||||
|
||||
def save_initial_predictions(model, out_dir, seq_name):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
with torch.no_grad():
|
||||
print('go into get_optim_result from output.py save_initial_predictions')
|
||||
pred_dict = model.get_optim_result()
|
||||
pred_dict = move_to(detach_all(pred_dict), "cpu")
|
||||
|
||||
for name, results in pred_dict.items():
|
||||
out_path = f"{out_dir}/{seq_name}_{0:06d}_init_{name}_results.npz"
|
||||
Logger.log(f"saving initial results to {out_path}")
|
||||
np.savez(out_path, **results)
|
||||
# print('cam_R in save_initial_predictions: ', results['cam_R'].shape, results['cam_R'])
|
||||
# print((results['cam_R'][0] == results['cam_R'][1]).all())
|
||||
# print((results['cam_R'][1] == results['cam_R'][2]).all())
|
||||
# print((results['cam_R'][2] == results['cam_R'][3]).all())
|
||||
# print((results['cam_R'][3] == results['cam_R'][4]).all())
|
||||
# print((results['cam_R'][4] == results['cam_R'][5]).all())
|
||||
# print((results['cam_R'][5] == results['cam_R'][6]).all())
|
||||
|
||||
|
||||
def save_input_poses(dataset, out_dir, seq_name, name="phalp"):
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
dataset.load_data(interp_input=False)
|
||||
d = dataset.data_dict
|
||||
res = {
|
||||
"pose_body": np.stack(d["init_body_pose"], axis=0),
|
||||
"trans": np.stack(d["init_trans"], axis=0),
|
||||
"root_orient": np.stack(d["init_root_orient"], axis=0),
|
||||
}
|
||||
print({k: v.shape for k, v in res.items()})
|
||||
out_path = f"{out_dir}/{seq_name}_{0:06d}_{name}_world_results.npz"
|
||||
Logger.log(f"saving inputs to {out_path}")
|
||||
np.savez(out_path, **res)
|
||||
|
||||
|
||||
def save_input_frames_ffmpeg(dataset, out_dir, name="input", fps=30, overwrite=False):
|
||||
vid_path = f"{out_dir}/{name}.mp4"
|
||||
if not overwrite and os.path.isfile(vid_path):
|
||||
return
|
||||
|
||||
list_path = f"{out_dir}/src_paths.txt"
|
||||
with open(list_path, "w") as f:
|
||||
f.write("\n".join([f"file '{p}'" for p in dataset.sel_img_paths]))
|
||||
|
||||
filter_str = f"-framerate {fps} -vf 'format=yuv420p'"
|
||||
cmd = (
|
||||
f"ffmpeg -loglevel error -f concat -safe 0 "
|
||||
f"-i {list_path} -c:v libx264 {filter_str} {vid_path} -y"
|
||||
)
|
||||
print(cmd)
|
||||
subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
||||
print(f"SAVED {len(dataset.sel_img_paths)} INPUT FRAMES TO {vid_path}")
|
||||
return vid_path
|
||||
|
||||
|
||||
def save_input_frames(dataset, vid_path, fps=30, overwrite=False):
|
||||
if not overwrite and os.path.isfile(vid_path):
|
||||
return
|
||||
|
||||
writer = imageio.get_writer(vid_path, fps=fps)
|
||||
for path in dataset.sel_img_paths:
|
||||
writer.append_data(imageio.imread(path))
|
||||
writer.close()
|
||||
print(f"SAVED {len(dataset.sel_img_paths)} INPUT FRAMES TO {vid_path}")
|
||||
return vid_path
|
||||
|
||||
|
||||
def load_track_info(path):
|
||||
assert os.path.isfile(path)
|
||||
|
||||
with open(path, "r") as f:
|
||||
track_info = json.load(f)
|
||||
|
||||
meta = track_info["meta"]
|
||||
tracks = track_info["tracks"]
|
||||
track_ids, idcs, vis_masks = map(
|
||||
torch.as_tensor,
|
||||
zip(
|
||||
*[
|
||||
(int(tid), info["index"], info["vis_mask"])
|
||||
for tid, info in tracks.items()
|
||||
]
|
||||
),
|
||||
)
|
||||
track_ids = track_ids[idcs]
|
||||
vis_masks = vis_masks[idcs]
|
||||
data_interval = meta["data_interval"]
|
||||
seq_interval = meta["seq_interval"]
|
||||
return track_ids, vis_masks, data_interval, seq_interval
|
||||
|
||||
|
||||
def save_track_info(dataset, out_dir):
|
||||
# track indices are relative to the entire sequence
|
||||
track_ids = dataset.track_ids
|
||||
track_dict = {}
|
||||
for i, (tid, mask) in enumerate(zip(dataset.track_ids, dataset.track_vis_masks)):
|
||||
track_dict[int(tid)] = {"index": i, "vis_mask": mask.tolist()}
|
||||
|
||||
out_dict = {
|
||||
"tracks": track_dict,
|
||||
"meta": {
|
||||
"seq_interval": (int(dataset.start_idx), int(dataset.end_idx)),
|
||||
"data_interval": (int(dataset.data_start), int(dataset.data_end)),
|
||||
},
|
||||
}
|
||||
|
||||
with open(f"{out_dir}/track_info.json", "w") as f:
|
||||
json.dump(out_dict, f)
|
||||
print("SAVED TRACK INFO")
|
||||
|
||||
|
||||
def load_camera_json(path):
|
||||
assert os.path.isfile(path)
|
||||
|
||||
with open(path, "r") as f:
|
||||
cam_data = json.load(f)
|
||||
R = torch.as_tensor(cam_data["rotation"]).reshape(-1, 3, 3)
|
||||
t = torch.as_tensor(cam_data["translation"])
|
||||
intrins = torch.as_tensor(cam_data["intrinsics"])
|
||||
return R.float(), t.float(), intrins.float()
|
||||
|
||||
|
||||
def save_camera_json(path, cam_R, cam_t, intrins):
|
||||
"""
|
||||
:param path
|
||||
:param cam_R (N, 3, 3)
|
||||
:param cam_t (N, 3)
|
||||
:param intrins (N, 4)
|
||||
"""
|
||||
N = len(cam_R)
|
||||
T = torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=torch.float32)
|
||||
cam_R = torch.einsum("ij,bjk->bik", T, cam_R)
|
||||
cam_t = torch.einsum("ij,bj->bi", T, cam_t)
|
||||
with open(path, "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"rotation": cam_R.reshape(N, 9).tolist(),
|
||||
"translation": cam_t.tolist(),
|
||||
"intrinsics": intrins.tolist(),
|
||||
},
|
||||
f,
|
||||
indent=1,
|
||||
)
|
||||
+174
@@ -0,0 +1,174 @@
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# from body_model import MANO_JOINTS
|
||||
from geometry.camera import invert_camera, compose_cameras
|
||||
from geometry.rotation import batch_rodrigues
|
||||
from util.logger import Logger
|
||||
from util.tensor import detach_all
|
||||
|
||||
# J_BODY = len(MANO_JOINTS) - 1 # no root
|
||||
# assert J_BODY == 15
|
||||
|
||||
class Params(nn.Module):
|
||||
def __init__(self, batch_size):
|
||||
super().__init__()
|
||||
self.batch_size = batch_size
|
||||
self.param_names = set()
|
||||
|
||||
def set_param(self, name, val, requires_grad=False):
|
||||
print("SETTING PARAM", name, val.shape)
|
||||
with torch.no_grad():
|
||||
setattr(
|
||||
self, name, nn.Parameter(val.contiguous(), requires_grad=requires_grad)
|
||||
)
|
||||
self.param_names.add(name)
|
||||
|
||||
def get_param(self, name):
|
||||
if name not in self.param_names:
|
||||
raise ValueError(f"{name} not stored as opt param")
|
||||
return getattr(self, name)
|
||||
|
||||
def load_dict(self, param_dict):
|
||||
for name, val in param_dict.items():
|
||||
self.set_param(name, val, requires_grad=False)
|
||||
|
||||
def get_dict(self):
|
||||
return {name: self.get_param_item(name) for name in self.param_names}
|
||||
|
||||
def get_vars(self, names=None):
|
||||
if names is None:
|
||||
names = self.param_names
|
||||
return {name: self.get_param(name) for name in names}
|
||||
|
||||
def get_param_item(self, name):
|
||||
with torch.no_grad():
|
||||
param = self.get_param(name)
|
||||
return param.detach()
|
||||
|
||||
def _set_param_grad(self, name, val: bool):
|
||||
if name not in self.param_names:
|
||||
raise ValueError(f"{name} not stored as param")
|
||||
param = getattr(self, name)
|
||||
assert isinstance(param, torch.Tensor)
|
||||
param.requires_grad = val
|
||||
|
||||
def set_require_grads(self, names):
|
||||
"""
|
||||
set parameters in names to True, set all others to False
|
||||
"""
|
||||
for name in self.param_names:
|
||||
self._set_param_grad(name, False)
|
||||
|
||||
for name in names:
|
||||
self._set_param_grad(name, True)
|
||||
|
||||
Logger.log("Set parameter grads:")
|
||||
Logger.log(
|
||||
{name: getattr(self, name).requires_grad for name in self.param_names}
|
||||
)
|
||||
|
||||
|
||||
class CameraParams(Params):
|
||||
"""
|
||||
Parameter container with cameras
|
||||
"""
|
||||
|
||||
def set_cameras(
|
||||
self, cam_data, opt_scale=True, opt_cams=False, opt_focal=False, **kwargs
|
||||
):
|
||||
self.opt_scale = opt_scale
|
||||
self.opt_cams = opt_cams
|
||||
self.opt_focal = opt_focal
|
||||
|
||||
# (T, 3, 3), (T, 3)
|
||||
cam_R, cam_t = cam_data["cam_R"], cam_data["cam_t"]
|
||||
intrins = cam_data["intrins"] # (T, 4)
|
||||
|
||||
T = cam_R.shape[0]
|
||||
device = cam_R.device
|
||||
|
||||
# assume focal length and center are same for all timesteps
|
||||
# print(intrins.shape, cam_R.shape, cam_t.shape)
|
||||
# torch.Size([128, 4]) torch.Size([128, 3, 3]) torch.Size([128, 3])
|
||||
|
||||
self.cam_center = intrins[:, 2:] # (T, 2)
|
||||
cam_f = intrins[:, :2] # (T, 2)
|
||||
if self.opt_focal:
|
||||
self.set_param("cam_f", cam_f)
|
||||
else:
|
||||
self.cam_f = cam_f
|
||||
|
||||
self._cam_R = cam_R # (T, 3, 3)
|
||||
self._cam_t = cam_t # (T, 3)
|
||||
|
||||
world_scale = torch.ones(1, 1, device=device)
|
||||
if self.opt_scale:
|
||||
if "world_scale" in kwargs:
|
||||
world_scale = kwargs["world_scale"]
|
||||
self.set_param("world_scale", world_scale)
|
||||
else:
|
||||
self.world_scale = world_scale
|
||||
|
||||
if self.opt_cams:
|
||||
init_delta_R = torch.zeros(T, 3, device=device)
|
||||
init_delta_t = torch.zeros(T, 3, device=device)
|
||||
if "delta_cam_R" in kwargs:
|
||||
print("setting delta_cam_R from kwargs")
|
||||
init_delta_R = kwargs["delta_cam_R"]
|
||||
if "delta_cam_t" in kwargs:
|
||||
print("setting delta_cam_t from kwargs")
|
||||
init_delta_t = kwargs["delta_cam_t"]
|
||||
|
||||
self.set_param("delta_cam_R", init_delta_R)
|
||||
self.set_param("delta_cam_t", init_delta_t)
|
||||
|
||||
@property
|
||||
def intrins(self): # (4,)
|
||||
return torch.cat([self.cam_f[0], self.cam_center[0]], dim=-1).detach().cpu()
|
||||
|
||||
def get_extrinsics(self):
|
||||
"""
|
||||
returns (T, 3, 3), (T, 3)
|
||||
"""
|
||||
cam_R, cam_t = self._cam_R, self._cam_t
|
||||
print('get_extrinsics')
|
||||
if self.opt_cams:
|
||||
raise ValueError
|
||||
print('self.opt_cams:', self.opt_cams)
|
||||
# raise ValueError
|
||||
dR = batch_rodrigues(self.delta_cam_R)
|
||||
cam_R = torch.matmul(cam_R, dR)
|
||||
cam_t = cam_t + self.delta_cam_t
|
||||
return cam_R, cam_t * self.world_scale
|
||||
|
||||
def get_intrinsics(self):
|
||||
"""
|
||||
returns (T, 4)
|
||||
"""
|
||||
return torch.cat([self.cam_f, self.cam_center], axis=-1)
|
||||
|
||||
def get_cameras(self, idcs=None):
|
||||
"""
|
||||
returns cam_R (B, T, 3, 3) cam_t (B, T, 3), cam_f (T, 2), cam_center (T, 2)
|
||||
"""
|
||||
if idcs is None:
|
||||
idcs = np.arange(self._cam_R.shape[0])
|
||||
|
||||
# print('idcs:', idcs, self._cam_R.shape[0], self._cam_R.shape)
|
||||
cam_R, cam_t = self.get_extrinsics()
|
||||
cam_R, cam_t = cam_R[None, idcs], cam_t[None, idcs]
|
||||
cam_f, cam_center = self.cam_f[idcs], self.cam_center[idcs]
|
||||
|
||||
B = self.batch_size
|
||||
cam_R = cam_R.repeat(B, 1, 1, 1)
|
||||
cam_t = cam_t.repeat(B, 1, 1)
|
||||
# print(cam_R.shape, cam_t.shape, cam_f.shape, cam_center.shape)
|
||||
# print(B, self.batch_size)
|
||||
# raise ValueError
|
||||
# torch.Size([2, 128, 3, 3]) torch.Size([2, 128, 3]) torch.Size([128, 2]) torch.Size([128, 2])
|
||||
# 2 2
|
||||
return cam_R, cam_t, cam_f, cam_center
|
||||
Reference in New Issue
Block a user