1732 lines
70 KiB
Python
1732 lines
70 KiB
Python
import os
|
|
import cv2
|
|
import copy
|
|
import glob
|
|
import time
|
|
import torch
|
|
# import ffmpeg
|
|
import joblib
|
|
import random
|
|
import shutil
|
|
import subprocess
|
|
import numpy as np
|
|
import torch.nn as nn
|
|
from loguru import logger
|
|
from argparse import Namespace
|
|
import matplotlib.pyplot as plt
|
|
import torch.nn.functional as F
|
|
from torch.autograd import Variable
|
|
|
|
# import scenepic_viz
|
|
# from utils import slerp
|
|
# import open3d_viz_overlay
|
|
# from datasets.amass import *
|
|
|
|
from arguments import Arguments
|
|
from argparse import ArgumentParser
|
|
from nemf.generative import Architecture
|
|
from nemf.fk import ForwardKinematicsLayer
|
|
from human_body_prior.tools.omni_tools import copy2cpu as c2c
|
|
from nemf.losses import GeodesicLoss, pos_smooth_loss, rot_smooth_loss
|
|
from rotations import (axis_angle_to_matrix, matrix_to_axis_angle, matrix_to_quaternion, matrix_to_rotation_6d,
|
|
quaternion_to_matrix, rotation_6d_to_matrix, axis_angle_to_quaternion, quaternion_to_axis_angle, quat_to_aa)
|
|
from fitting_utils import (process_gt, gmof, perspective_projection, get_joints2d, run_pymafx, run_metro, process_pymafx_mano, compute_seq_intervals,
|
|
save_quantitative_evaluation, get_seqname_ho3d_v3, get_seqname_arctic_data, get_seqname_in_the_wild, get_seqname_dexycb,
|
|
RIGHT_WRIST_BASE_LOC, joints2d_loss, # map_openpose_joints_to_mano, map_mano_joints_to_openpose,
|
|
export_pymafx_json, blend_keypoints, BMCLoss)
|
|
from utils import estimate_angular_velocity, estimate_linear_velocity, render_openpose
|
|
bmc = BMCLoss(lambda_bl=1, lambda_rb=1, lambda_a=1)
|
|
|
|
HAND_JOINT_NUM = 16
|
|
# mano2openpose = map_mano_joints_to_openpose()
|
|
# openpose2mano = map_openpose_joints_to_mano()
|
|
|
|
IGNORE_KEYS = ['cam_f', 'cam_center', 'img_height', 'img_width', 'img_dir',
|
|
'save_path', 'frame_id', 'config_type', 'rh_verts', 'handedness']
|
|
POSSIBLE_KEYP_SOURCES = ["mmpose", "mediapipe", "mediapipe_std", "mediapipe_multiview",
|
|
"pymafx", "pymafx_std", "gt", "metro", "blend",
|
|
"blend_mediapipe_std_mmpose",
|
|
"blend_std", "blend_smooth"]
|
|
|
|
MANO_JOINTS = {
|
|
'wrist': 0,
|
|
'index1': 1,
|
|
'index2': 2,
|
|
'index3': 3,
|
|
'middle1': 4,
|
|
'middle2': 5,
|
|
'middle3': 6,
|
|
'pinky1': 7,
|
|
'pinky2': 8,
|
|
'pinky3': 9,
|
|
'ring1': 10,
|
|
'ring2': 11,
|
|
'ring3': 12,
|
|
'thumb1': 13,
|
|
'thumb2': 14,
|
|
'thumb3': 15,
|
|
}
|
|
def run_mano(body_model, trans, root_orient, body_pose, is_right, betas=None, only_right=False):
|
|
"""
|
|
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 : (optional) B x D
|
|
"""
|
|
B, T, _ = trans.shape
|
|
bm_batch_size = body_model.batch_size
|
|
assert bm_batch_size % B == 0
|
|
# assert (is_right[0]==is_right[0][0]).all(), f'{is_right}'
|
|
# assert (is_right[1]==is_right[1][0]).all(), f'{is_right}'
|
|
|
|
seq_len = bm_batch_size // B
|
|
bm_num_betas = body_model.num_betas
|
|
J_BODY = len(MANO_JOINTS) - 1 # all joints except root
|
|
# print('utils.py, seq_len, T: ', seq_len, T, bm_batch_size, B) # 1, 120, 1, 1
|
|
# print('trans, root_orient, body_pose:', trans.shape, root_orient.shape, body_pose.shape)
|
|
if T == 1:
|
|
raise ValueError
|
|
# must expand to use with body model
|
|
trans = trans.expand(B, seq_len, 3)
|
|
root_orient = root_orient.expand(B, seq_len, 3)
|
|
body_pose = body_pose.expand(B, seq_len, J_BODY * 3)
|
|
elif T != seq_len:
|
|
# print(trans.shape, root_orient.shape, body_pose.shape)
|
|
trans, root_orient, body_pose = zero_pad_tensors(
|
|
[trans, root_orient, body_pose], seq_len - T
|
|
)
|
|
if betas is None:
|
|
betas = torch.zeros(B, bm_num_betas, device=trans.device)
|
|
betas = betas.reshape((B, 1, bm_num_betas)).expand((B, seq_len, bm_num_betas))
|
|
# print('body_pose: ', body_pose.reshape((B * seq_len, -1)).shape)
|
|
|
|
mano_output = body_model(
|
|
hand_pose=body_pose.reshape((B * seq_len, -1)),
|
|
betas=betas.reshape((B * seq_len, -1)),
|
|
global_orient=root_orient.reshape((B * seq_len, -1)),
|
|
transl=trans.reshape((B * seq_len, -1)),
|
|
)
|
|
joints = mano_output.joints
|
|
verts = mano_output.vertices
|
|
|
|
joints = joints.reshape(B, seq_len, -1, 3)[:, :T]
|
|
verts = verts.reshape(B, seq_len, -1, 3)[:, :T]
|
|
is_right = is_right.unsqueeze(-1)
|
|
|
|
if not only_right:
|
|
joints[:, :, :, 0] = (2*is_right-1)*joints[:, :, :, 0]
|
|
verts[:, :, :, 0] = (2*is_right-1)*verts[:, :, :, 0]
|
|
|
|
return {
|
|
"joints": joints,
|
|
"vertices": verts,
|
|
"l_faces": body_model.faces_tensor[:,[0,2,1]],
|
|
"r_faces": body_model.faces_tensor,
|
|
"is_right": is_right.squeeze(-1),
|
|
'body_pose': body_pose.clone()
|
|
}
|
|
|
|
def zero_pad_tensors(pad_list, pad_size):
|
|
"""
|
|
Assumes tensors in pad_list are B x T x D and pad temporal dimension
|
|
"""
|
|
B = pad_list[0].size(0)
|
|
new_pad_list = []
|
|
for pad_idx, pad_tensor in enumerate(pad_list):
|
|
padding = torch.zeros((B, pad_size, pad_tensor.size(2))).to(pad_tensor)
|
|
new_pad_list.append(torch.cat([pad_tensor, padding], dim=1))
|
|
return new_pad_list
|
|
|
|
def pred_mano(hand_model, trans, root_orient, body_pose, is_right, betas, only_right=False):
|
|
"""
|
|
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(hand_model, trans, root_orient, body_pose, is_right, betas=betas, only_right=only_right)
|
|
joints3d, points3d = mano_out["joints"], mano_out["vertices"]
|
|
verts3d = points3d
|
|
|
|
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 reproject(points3d, cam_R, cam_t, cam_f, cam_center):
|
|
"""
|
|
reproject points3d into the scene cameras
|
|
:param points3d (B, T, N, 3)
|
|
:param cam_R (B, T, 3, 3)
|
|
:param cam_t (B, T, 3)
|
|
:param cam_f (T, 2)
|
|
:param cam_center (T, 2)
|
|
"""
|
|
B, T, N, _ = points3d.shape
|
|
points3d = torch.einsum("btij,btnj->btni", cam_R, points3d)
|
|
points3d = points3d + cam_t[..., None, :] # (B, T, N, 3)
|
|
points2d = points3d[..., :2] / points3d[..., 2:3]
|
|
points2d = cam_f[None, :, None] * points2d + cam_center[None, :, None]
|
|
return points2d
|
|
|
|
def get_stage2_res(base_dir, device, npz_init_dict, hand_model):
|
|
args = Arguments(base_dir, os.path.dirname(__file__), filename='amass.yaml')
|
|
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
fk = ForwardKinematicsLayer(args)
|
|
to_th = lambda x: torch.from_numpy(x).to(device)
|
|
|
|
cdata = npz_init_dict
|
|
|
|
gender = 'neutral'
|
|
if gender.startswith('b'):
|
|
gender = str(cdata['gender'], encoding='utf-8')
|
|
|
|
N = len(cdata['poses'])
|
|
print(f'Sequence has {N} frames')
|
|
|
|
# no matter what the keypoint source is, we need to read the bbox from mmpose
|
|
# joint2d_data = to_th(np.stack(keyp_frames, axis=0)) # T x J x 3 (x,y,conf)
|
|
joint2d_data = npz_init_dict['keyp2d']
|
|
is_right = npz_init_dict['is_right']
|
|
vis_mask = npz_init_dict['vis_mask']
|
|
|
|
assert (len(np.where(is_right==1)[0])==0) or (len(np.where(is_right==0)[0])==0), f'{len(np.where(is_right==0)[0]), len(np.where(is_right==1)[0])}'
|
|
|
|
# ONLY FOR RIGHT HAND
|
|
root_orient_aa = cdata['root_orient']
|
|
pose_body = cdata['poses']
|
|
data_poses = np.concatenate((root_orient_aa, pose_body), axis=1)
|
|
|
|
pose = torch.from_numpy(np.asarray(data_poses, np.float32)).to(device)
|
|
pose = pose.view(-1, 15 + 1, 3) # axis-angle (T, J, 3)
|
|
assert len(pose) == 128
|
|
|
|
trans = torch.from_numpy(np.asarray(cdata['trans'], np.float32)).to(device) # global translation (T, 3)
|
|
|
|
# Compute necessary data for model training.
|
|
rotmat = axis_angle_to_matrix(pose) # rotation matrix (T, J, 3, 3)
|
|
root_orient = rotmat[:, 0].clone()
|
|
root_orient = matrix_to_rotation_6d(root_orient) # root orientation (T, 6)
|
|
|
|
# print(pose.shape, trans.shape, rotmat.shape, root_orient.shape)
|
|
# raise ValueError
|
|
# defined in amass.yaml. Set as True
|
|
if args.unified_orientation:
|
|
identity = torch.eye(3).cuda()
|
|
identity = identity.view(1, 3, 3).repeat(rotmat.shape[0], 1, 1)
|
|
rotmat[:, 0] = identity
|
|
rot6d = matrix_to_rotation_6d(rotmat) # 6D rotation representation (T, J, 6)
|
|
|
|
rot_seq = rotmat.clone()
|
|
angular = estimate_angular_velocity(rot_seq.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0) # angular velocity of all the joints (T, J, 3)
|
|
|
|
pos, global_xform = fk(rot6d) # local joint positions (T, J, 3), global transformation matrix for each joint (T, J, 4, 4)
|
|
|
|
pos = pos.contiguous()
|
|
global_xform = global_xform.contiguous()
|
|
velocity = estimate_linear_velocity(pos.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0) # linear velocity of all the joints (T, J, 3)
|
|
|
|
# defined in amass.yaml. Set as True
|
|
if args.unified_orientation:
|
|
root_rotation = rotation_6d_to_matrix(root_orient) # (T, 3, 3)
|
|
root_rotation = root_rotation.unsqueeze(1).repeat(1, args.smpl.joint_num, 1, 1) # (T, J, 3, 3)
|
|
global_pos = torch.matmul(root_rotation, pos.unsqueeze(-1)).squeeze(-1)
|
|
height = global_pos + trans.unsqueeze(1)
|
|
else:
|
|
height = pos + trans.unsqueeze(1)
|
|
height = height[..., 'xyz'.index(args.data.up)] # (T, J)
|
|
|
|
root_vel = estimate_linear_velocity(trans.unsqueeze(0), dt=1.0 / args.data.fps).squeeze(0) # linear velocity of the root joint (T, 3)
|
|
|
|
global_xform = global_xform[:, :, :3, :3] # (T, J, 3, 3)
|
|
global_xform = matrix_to_rotation_6d(global_xform) # (T, J, 6)
|
|
|
|
betas = to_th(cdata["betas"])
|
|
out_mano = pred_mano(hand_model, \
|
|
trans[None], torch.tensor(root_orient_aa).to(device)[None], \
|
|
torch.tensor(pose_body).to(device)[None], torch.tensor(is_right).to(device)[None], \
|
|
torch.tensor(cdata["betas"]).to(device)[None], only_right=True)
|
|
|
|
assert len(out_mano['joints3d']) == 1
|
|
joints3d, vertices3d = out_mano['joints3d'][0], out_mano['verts3d'][0]
|
|
# print(npz_init_dict.keys())
|
|
# for i in npz_init_dict.keys():
|
|
# print(i, npz_init_dict[i].shape)
|
|
|
|
data = {'rotmat': rotmat,
|
|
'pos': pos,
|
|
'trans': trans,
|
|
'root_vel': root_vel,
|
|
'height': height,
|
|
'rot6d': rot6d,
|
|
'angular': angular,
|
|
'betas': betas,
|
|
'global_xform': global_xform,
|
|
'velocity': velocity,
|
|
'root_orient': root_orient,
|
|
'joints2d': torch.tensor(joint2d_data).to(device),
|
|
'joints3d': joints3d,
|
|
'vertices': vertices3d,
|
|
'cam_t': torch.tensor(npz_init_dict['cam_t']).to(device),
|
|
"cam_R": torch.tensor(npz_init_dict['cam_R']).to(device),
|
|
"cam_center": torch.tensor(npz_init_dict['cam_center']).to(device),
|
|
"cam_f": torch.tensor(npz_init_dict['cam_f']).to(device),
|
|
'is_right': torch.tensor(is_right).to(device),
|
|
'vis_mask': torch.tensor(vis_mask).to(device),
|
|
}
|
|
# for i in data.keys():
|
|
# print(i, data[i].shape)
|
|
|
|
return data
|
|
|
|
def forward_mano(output):
|
|
rotmat = output['rotmat'] # (B, T, J, 3, 3)
|
|
B, T, J, _, _ = rotmat.size()
|
|
|
|
# b_size, _, n_joints = rotmat.shape[:3]
|
|
local_rotmat = fk.global_to_local(rotmat.view(-1, J, 3, 3)) # (B x T, J, 3, 3)
|
|
local_rotmat = local_rotmat.view(B, -1, J, 3, 3)
|
|
|
|
root_orient = rotation_6d_to_matrix(output['root_orient']) # (B, T, 3, 3)
|
|
|
|
local_rotmat[:, :, 0] = root_orient
|
|
|
|
poses = matrix_to_axis_angle(local_rotmat) # (T, J, 3)
|
|
poses = poses.view(-1, J * 3)
|
|
|
|
# no_shift is the flag for not shifting the wrist location to the origin
|
|
mano_out = args.hand_model(input_dict={"betas":output['betas'].view(-1, 10),
|
|
"global_orient":poses[..., :3].view(-1, 3),
|
|
"hand_pose":poses[..., 3:].view(-1, 45),
|
|
"no_shift":True,
|
|
"return_finger_tips": True,
|
|
"transl":output['trans'].view(-1, 3)})
|
|
|
|
return mano_out
|
|
|
|
|
|
def L_pose_prior(output):
|
|
|
|
hposer.eval()
|
|
|
|
rotmat = output['rotmat'] # (B, T, J, 3, 3)
|
|
B, T, J, _, _ = rotmat.size()
|
|
|
|
# b_size, _, n_joints = rotmat.shape[:3]
|
|
local_rotmat = fk.global_to_local(rotmat.view(-1, J, 3, 3)) # (B x T, J, 3, 3)
|
|
local_rotmat = local_rotmat.view(B, -1, J, 3, 3)
|
|
|
|
local_rotmat_casted = local_rotmat[:, :, 1:, ...].view(B*T, -1)
|
|
|
|
pose_latent_code = hposer.encode(local_rotmat_casted)
|
|
pose_prior_mean_squared = (pose_latent_code.mean ** 2).mean(-1)
|
|
|
|
loss = pose_prior_mean_squared.mean()
|
|
|
|
return loss
|
|
|
|
|
|
def L_rot(pred, gt, T, conf=None):
|
|
"""
|
|
Args:
|
|
source, target: rotation matrices in the shape B x T x J x 3 x 3.
|
|
T: temporal masks.
|
|
|
|
Returns:
|
|
reconstruction loss evaluated on the rotation matrices.
|
|
"""
|
|
if conf is not None:
|
|
criterion_rec = nn.L1Loss(reduction='none') if args.l1_loss else nn.MSELoss(reduction='none')
|
|
else:
|
|
criterion_rec = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
|
criterion_geo = GeodesicLoss()
|
|
|
|
B, seqlen, J, _, _ = pred.shape
|
|
|
|
if args.geodesic_loss:
|
|
if conf is not None:
|
|
loss = (conf.squeeze(-1) ** 2) * criterion_geo(pred[:, T].view(-1, 3, 3), gt[:, T].view(-1, 3, 3), reduction='none').reshape(B, seqlen, J)
|
|
loss = loss.mean()
|
|
else:
|
|
loss = criterion_geo(pred[:, T].view(-1, 3, 3), gt[:, T].view(-1, 3, 3))
|
|
else:
|
|
if conf is not None:
|
|
loss = (conf.unsqueeze(-1) ** 2) * criterion_rec(pred[:, T], gt[:, T])
|
|
loss = loss.mean()
|
|
else:
|
|
loss = criterion_rec(pred[:, T], gt[:, T])
|
|
|
|
return loss
|
|
|
|
def pose_prior_loss(latent_pose_pred, mask=None):
|
|
"""
|
|
:param latent_pose_pred (B, T, D)
|
|
:param mask (optional) (B, T)
|
|
"""
|
|
# prior is isotropic gaussian so take L2 distance from 0
|
|
loss = latent_pose_pred**2
|
|
if mask is not None:
|
|
loss = loss[mask.bool()]
|
|
loss = torch.sum(loss)
|
|
return loss
|
|
|
|
|
|
def L_PCA(pose):
|
|
|
|
bs = pose.shape[0]
|
|
|
|
# convert aa to rotmat
|
|
pose_rotmat_global = axis_angle_to_matrix(pose)
|
|
|
|
# convert global to local
|
|
pose_rotmat_local = fk.global_to_local(pose_rotmat_global.reshape(-1, 16, 3, 3))
|
|
|
|
# convert back to aa
|
|
pose_aa_local = matrix_to_axis_angle(pose_rotmat_local)
|
|
|
|
pose_reshaped = pose_aa_local[:, 1:, :].reshape(-1, args.data.clip_length * 15 * 3)
|
|
|
|
pose_reshaped_centered = pose_reshaped - pca_mean
|
|
|
|
pose_projected = pose_reshaped_centered @ pca_pc.T
|
|
|
|
normalized_pose_loss = abs(pose_projected) / pca_sv
|
|
|
|
mp_loss = normalized_pose_loss.mean(0).mean()
|
|
|
|
return mp_loss
|
|
|
|
def L_GMM(pose):
|
|
|
|
bs = pose.shape[0]
|
|
|
|
# convert aa to rotmat
|
|
pose_rotmat_global = axis_angle_to_matrix(pose)
|
|
|
|
# convert global to local
|
|
pose_rotmat_local = fk.global_to_local(pose_rotmat_global.reshape(-1, 16, 3, 3))
|
|
|
|
# convert back to aa
|
|
pose_aa_local = matrix_to_axis_angle(pose_rotmat_local).reshape(bs, args.data.clip_length, -1, 3)
|
|
|
|
mp_loss_list = []
|
|
loss_tot = 0
|
|
|
|
for i in range(bs):
|
|
loss_i = gmm_aa.log_likelihood(pose_aa_local[i, :, 1:].reshape(1, -1).cpu())
|
|
loss_tot += loss_i
|
|
|
|
mp_loss = loss_tot / bs
|
|
return mp_loss.to("cuda")
|
|
|
|
def L_orient(source, target, T, bbox_conf=None):
|
|
"""
|
|
Args:
|
|
source: predicted root orientation in the shape B x T x 6.
|
|
target: root orientation in the shape of B x T x 6.
|
|
T: temporal masks.
|
|
|
|
Returns:
|
|
reconstruction loss evaluated on the root orientation.
|
|
"""
|
|
criterion_rec = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
|
criterion_geo = GeodesicLoss()
|
|
|
|
source = rotation_6d_to_matrix(source) # (B, T, 3, 3)
|
|
target = rotation_6d_to_matrix(target) # (B, T, 3, 3)
|
|
|
|
if args.geodesic_loss:
|
|
|
|
if bbox_conf is not None:
|
|
|
|
loss = criterion_geo(source[:, T].view(-1, 3, 3), target[:, T].view(-1, 3, 3), reduction='none')
|
|
bbox_conf_coef = bbox_conf.reshape(-1)
|
|
loss = ((bbox_conf_coef ** 2) * loss).mean()
|
|
|
|
else:
|
|
loss = criterion_geo(source[:, T].view(-1, 3, 3), target[:, T].view(-1, 3, 3))
|
|
|
|
else:
|
|
loss = criterion_rec(source[:, T], target[:, T])
|
|
|
|
return loss
|
|
|
|
|
|
def L_trans(source, target, T, bbox_conf=None):
|
|
"""
|
|
Args:
|
|
source: predict global translation in the shape B x T x 3 (the origin is (0, 0, height)).
|
|
target: global translation of the root joint in the shape B x T x 3.
|
|
T: temporal masks.
|
|
|
|
Returns:
|
|
reconstruction loss evaluated on the global translation.
|
|
"""
|
|
|
|
trans = source
|
|
trans_gt = target
|
|
|
|
# dont make reduction and weight by bbox_conf
|
|
if bbox_conf is not None:
|
|
criterion_pred = nn.L1Loss(reduction='none') if args.l1_loss else nn.MSELoss(reduction='none')
|
|
|
|
# reshape to (T * N, 3)
|
|
loss = criterion_pred(trans[:, T].reshape(-1, 3), trans_gt[:, T].reshape(-1, 3)).mean(1)
|
|
bbox_conf_coef = bbox_conf.reshape(-1)
|
|
loss = ((bbox_conf_coef **2) * loss).mean()
|
|
|
|
else:
|
|
criterion_pred = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
|
loss = criterion_pred(trans[:, T], trans_gt[:, T])
|
|
|
|
return loss
|
|
|
|
|
|
def L_pos(source, target, T):
|
|
"""
|
|
Args:
|
|
source, target: joint local positions in the shape B x T x J x 3.
|
|
T: temporal masks.
|
|
|
|
Returns:
|
|
reconstruction loss evaluated on the joint local positions.
|
|
"""
|
|
criterion_rec = nn.L1Loss() if args.l1_loss else nn.MSELoss()
|
|
loss = criterion_rec(source[:, T], target[:, T])
|
|
|
|
return loss
|
|
|
|
|
|
def contact_vel_loss(contacts_conf, joints3d):
|
|
'''
|
|
Velocity should be zero at predicted contacts
|
|
'''
|
|
delta_pos = (joints3d[:,1:] - joints3d[:,:-1])**2
|
|
cur_loss = delta_pos.sum(dim=-1) * contacts_conf[:,1:]
|
|
cur_loss = 0.5 * torch.mean(cur_loss)
|
|
|
|
return cur_loss
|
|
|
|
def motion_prior_loss(latent_motion_pred):
|
|
# assume standard normal
|
|
loss = latent_motion_pred**2
|
|
loss = torch.mean(loss)
|
|
|
|
return loss
|
|
|
|
|
|
def motion_reconstruction(hand_model, target, output_dir, steps, T=None, idx=0):
|
|
|
|
z_l, _, _ = model.encode_local()
|
|
z_g, _, _ = model.encode_global()
|
|
|
|
# optimize pose directly
|
|
pose_aa = matrix_to_axis_angle(model.input_data["rotmat"].detach().clone()) if motion_prior_type in ['gmm', 'pca'] else None
|
|
|
|
z_l, z_g, opt_betas, opt_root_orient, opt_trans, cam_R, cam_t, cam_f, cam_center, opt_pose_aa = latent_optimization(hand_model, target, T=T, z_l=z_l, z_g=z_g, pose=pose_aa)
|
|
|
|
for step in steps:
|
|
|
|
fps = int(args.data.fps / step)
|
|
|
|
with torch.no_grad():
|
|
B, seqlen, _ = opt_trans.shape
|
|
|
|
if motion_prior_type in ["pca", "gmm"]:
|
|
output = {"rotmat": axis_angle_to_matrix(opt_pose_aa)}
|
|
else:
|
|
output = model.decode(z_l, z_g=z_g, length=args.data.clip_length, step=step)
|
|
|
|
rotmat = output['rotmat'] # (B, T, J, 3, 3)
|
|
rotmat_gt = target['rotmat'] # (B, T, J, 3, 3)
|
|
b_size, _, n_joints = rotmat.shape[:3]
|
|
|
|
local_rotmat = fk.global_to_local(rotmat.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
|
local_rotmat = local_rotmat.view(b_size, -1, n_joints, 3, 3)
|
|
# local_rotmat_gt = fk.global_to_local(rotmat_gt.view(-1, n_joints, 3, 3)) # (B x T, J, 3, 3)
|
|
# local_rotmat_gt = local_rotmat_gt.view(b_size, -1, n_joints, 3, 3)
|
|
|
|
root_orient = rotation_6d_to_matrix(opt_root_orient) # (B, T, 3, 3)
|
|
# root_orient_gt = rotation_6d_to_matrix(target['root_orient']) # (B, T, 3, 3)
|
|
|
|
global_trans = opt_trans.clone()
|
|
B = global_trans.shape[0]
|
|
|
|
output['betas'] = opt_betas[:, None, None, :].repeat(1, B, seqlen, 1)
|
|
output['trans'] = opt_trans.clone().reshape(B*args.nsubject, seqlen, 3)
|
|
output['root_orient'] = opt_root_orient
|
|
|
|
####################
|
|
body_pose = matrix_to_axis_angle(local_rotmat)[:, :, 1:] # matrix_to_axis_angle(local_rotmat)
|
|
is_right = target['is_right'].clone()
|
|
####################
|
|
|
|
R = matrix_to_axis_angle(rotation_6d_to_matrix(opt_root_orient))
|
|
T = output['trans']
|
|
P = body_pose.reshape(1, 128, 45)
|
|
Be = opt_betas
|
|
|
|
DR = matrix_to_axis_angle(local_rotmat[:, :, 0]).clone()
|
|
if args.data.root_transform:
|
|
local_rotmat[:, :, 0] = root_orient
|
|
# local_rotmat_gt[:, :, 0] = root_orient_gt
|
|
|
|
trans = global_trans
|
|
trans_gt = target['trans'] # (B, T, 3)
|
|
|
|
if trans.shape[0] > 1:
|
|
raise ValueError
|
|
|
|
assert local_rotmat.shape[0] == 1
|
|
return R, T, P, Be, DR
|
|
# for i in range(local_rotmat.shape[0]): # (1, T. J, 3, 3)
|
|
|
|
# poses = c2c(matrix_to_axis_angle(local_rotmat[i])) # (T, J, 3)
|
|
# poses_gt = c2c(matrix_to_axis_angle(local_rotmat_gt[i])) # (T, J, 3)
|
|
|
|
# poses = poses.reshape((poses.shape[0], -1)) # (T, 48)
|
|
# poses_gt = poses_gt.reshape((poses_gt.shape[0], -1)) # (T, 66)
|
|
|
|
# limit = poses.shape[0]
|
|
|
|
# pred_save_path = os.path.join(output_dir, f'recon_{offset + i:03d}_{fps}fps.npz')
|
|
# gt_save_path = pred_save_path if args.raw_config else os.path.join(output_dir, f'recon_{offset + i:03d}_gt.npz')
|
|
|
|
# save opt results if not _pymafx_raw or _metro_raw
|
|
# save_keys = ['root_orient', 'trans', 'latent_pose', 'is_right', 'init_body_pose', \
|
|
# 'world_scale', 'betas', 'pose_body', 'cam_R', 'cam_t', 'intrins']
|
|
|
|
# if not args.raw_config:
|
|
# np.savez(pred_save_path,
|
|
# poses=poses[:limit].detach().cpu().numpy(),
|
|
# trans=c2c(trans[i][:limit]),
|
|
# betas=opt_betas.detach().cpu().numpy()[:limit],
|
|
# gender=args.data.gender,
|
|
# mocap_framerate=fps,
|
|
# cam_R=c2c(target['cam_R'][i][0]),
|
|
# cam_t=c2c(target["cam_t"][i][0]),
|
|
# cam_f=c2c(target['cam_f'].unsqueeze(0)),
|
|
# cam_center=c2c(target['cam_center'].unsqueeze(0)),
|
|
# joints_2d=c2c(joints2d_pred[i][:limit]),
|
|
# keypoints_2d=c2c(target['joints2d'][i][:limit]), # detected keypoints with confidence
|
|
# vertices=c2c(vertices_pred[i][:limit]),
|
|
# save_path=pred_save_path,
|
|
# handedness=target["handedness"].detach().cpu().numpy(),
|
|
# # img_dir=str(target['img_dir']),
|
|
# # frame_id=target['frame_id'],
|
|
# # img_height=target['img_height'],
|
|
# # img_width=target['img_width'],
|
|
# # config_type=target['config_type'],
|
|
# joints_3d=c2c(joints3d_pred)[i][:limit])
|
|
|
|
# if ("_encode_decode" in gt_save_path) or args.raw_config:
|
|
# np.savez(gt_save_path,
|
|
# poses=poses_gt[:limit].detach().cpu().numpy(),
|
|
# trans=c2c(trans_gt[i][:limit]),
|
|
# betas=target["betas"][0].detach().cpu().numpy()[:limit],
|
|
# vertices=target["rh_verts"].detach().cpu().numpy()[:limit],
|
|
# config_type=target['config_type'] if args.raw_config else None,
|
|
# gender=args.data.gender,
|
|
# save_path=gt_save_path,
|
|
# mocap_framerate=args.data.fps,
|
|
# # frame_id=target['frame_id'],
|
|
# # img_height=target['img_height'],
|
|
# # img_width=target['img_width'],
|
|
# cam_R=c2c(target['cam_R'][i][0]),
|
|
# cam_t=c2c(target['cam_t'][i][0]),
|
|
# cam_f=c2c(target['cam_f'].unsqueeze(0)),
|
|
# cam_center=c2c(target['cam_center'].unsqueeze(0)),
|
|
# joints2d=c2c(target['joints2d'][i, ...,])[:limit],
|
|
# joints_2d=c2c(target['joints2d'][i, ...,])[:limit],
|
|
# keypoints_2d=c2c(target['joints2d'][i, ...,])[:limit], # duplicated, because alignment reads that key value
|
|
# joints_3d=c2c(target['joints3d'][i])[:limit])
|
|
|
|
def get_gt_path(expname):
|
|
if args.dataname == "DexYCB":
|
|
gt_path = args.vid_path
|
|
subjectname = args.vid_path.split("/")[4]
|
|
elif args.dataname == "HO3D_v3":
|
|
gt_path = os.path.join(os.path.dirname(args.vid_path), "meta")
|
|
subjectname = args.vid_path.split("/")[-2]
|
|
elif args.dataname == "arctic_data":
|
|
subjectname = args.vid_path.split("/")[-3]
|
|
split_path = os.path.join(f"./external/arctic/data/arctic_data/data/splits/p1_{expname}.npy")
|
|
meta_path = os.path.join(f"./external/arctic/data/arctic_data/data/meta/misc.json")
|
|
gt_path = {"gt_path": args.vid_path, "split": split_path, "meta_path": meta_path}
|
|
else:
|
|
args.dataname = "in_the_wild"
|
|
subjectname = args.vid_path.split("/")[-2]
|
|
gt_path = os.path.join(os.path.dirname(args.vid_path), "meta")
|
|
|
|
return gt_path, subjectname
|
|
|
|
|
|
def multi_stage_opt(opt, device, obs_data, res_dict, hand_model, config_f, exp_setup_name, init_method_name):
|
|
|
|
logger.info(f"Running reconstruction with prior")
|
|
|
|
config_type = config_f.split('/')[-1].split('.')[:-1]
|
|
config_type = '.'.join(config_type)
|
|
args.raw_config = config_type in ["_pymafx_raw", "_metro_raw"]
|
|
|
|
keypoint_blend_weight = 1.0
|
|
vid_path = args.vid_path
|
|
abs_video_path = os.path.join(os.getcwd(), args.vid_path)
|
|
|
|
args.dataname = args.vid_path.split("/")[3]
|
|
args.N_frames = len(glob.glob(os.path.join(abs_video_path, "*.jpg")))
|
|
|
|
if args.N_frames == 0:
|
|
args.N_frames = len(glob.glob(os.path.join(abs_video_path, "*.png")))
|
|
|
|
args.hand_model = hand_model
|
|
gt_path, subjectname = get_gt_path(exp_setup_name)
|
|
init_method_out_path = os.path.join(os.path.dirname(vid_path), f"{init_method_name}_out")
|
|
|
|
print('loading data...')
|
|
assert len(res_dict) == 1
|
|
res_dict = res_dict[0]
|
|
print('obs_data: ', obs_data.keys())
|
|
print('res_dict: ', res_dict.keys())
|
|
for i in res_dict.keys():
|
|
res_dict[i] = res_dict[i].cpu().detach().numpy()
|
|
print(i, res_dict[i].shape)
|
|
|
|
for i in obs_data.keys():
|
|
try:
|
|
obs_data[i] = obs_data[i].cpu().detach().numpy()
|
|
except:
|
|
pass
|
|
|
|
R_list = []
|
|
T_list = []
|
|
P_list = []
|
|
Be_list = []
|
|
DR_list = []
|
|
for idx in range(len(res_dict['pose_body'])):
|
|
|
|
assert (res_dict['is_right'][idx] == (obs_data['is_right'][idx])).all()
|
|
rhand_orient_padded, rhand_betas_padded, rhand_trans_padded, rhand_pose_padded, is_right = \
|
|
res_dict['root_orient'][idx], res_dict['betas'][idx], res_dict['trans'][idx], res_dict['pose_body'][idx], res_dict['is_right'][idx]
|
|
|
|
cam_center = torch.tensor(res_dict['intrins'][2:][None]).repeat(128, 1) # (T, 2)
|
|
cam_f = torch.tensor(res_dict['intrins'][:2][None]).repeat(128, 1) # (T, 2)
|
|
|
|
init_dict = {
|
|
"keyp2d": obs_data['joints2d'][idx],
|
|
"betas": rhand_betas_padded,
|
|
"trans": rhand_trans_padded,
|
|
"root_orient": rhand_orient_padded,
|
|
"poses": rhand_pose_padded.reshape(-1, 45),
|
|
"cam_R": res_dict['cam_R'][idx],
|
|
"cam_t": res_dict['cam_t'][idx],
|
|
"img_dir":abs_video_path,
|
|
"is_right": res_dict['is_right'][idx],
|
|
"cam_f": cam_f,
|
|
"cam_center": cam_center,
|
|
'vis_mask': obs_data['vis_mask'][idx]
|
|
}
|
|
|
|
# save the results, expand the length by padding so that it matches the number of frames in the video
|
|
args.pkl_output_dir = os.path.join(args.save_path, "pkls")
|
|
os.makedirs(args.pkl_output_dir, exist_ok=True)
|
|
|
|
# run stage3 optimization
|
|
os.makedirs(args.save_path, exist_ok=True)
|
|
data = get_stage2_res(opt.paths.base_dir, device, init_dict, hand_model)
|
|
|
|
data['save_path'] = os.path.join(args.save_path, 'pymaf_output.npz')
|
|
data['config_type'] = config_type
|
|
args.orig_seq_len = data['trans'].shape[0]
|
|
|
|
shutil.copy2(config_f, args.save_path)
|
|
|
|
data['betas'] = data['betas'][None].repeat(128, 1)
|
|
|
|
for k, v in data.items():
|
|
if k in IGNORE_KEYS:
|
|
continue
|
|
else:
|
|
if k == 'betas':
|
|
print(k, v.shape)
|
|
if v.shape[0] > 128:
|
|
|
|
# in case of batch optimization, we need to split the data into chunks of 128 frames
|
|
if args.overlap_len > 0:
|
|
# compute start and end indices
|
|
seq_intervals = compute_seq_intervals(v.shape[0], 128, args.overlap_len)
|
|
data_split = []
|
|
for seq_s, seq_e in seq_intervals:
|
|
data_split.append(v[seq_s:seq_e])
|
|
else:
|
|
data_split = list(torch.split(v, 128))
|
|
|
|
if data_split[-1].shape[0] == 128:
|
|
data[k] = torch.stack(data_split, dim=0)
|
|
else:
|
|
pad_repeat = 128 - data_split[-1].shape[0]
|
|
last_el = data_split[-1]
|
|
last_el = torch.cat([last_el, last_el[-1:].repeat_interleave(pad_repeat, 0)])
|
|
data_split[-1] = last_el
|
|
data[k] = torch.stack(data_split, dim=0)
|
|
else:
|
|
pad_repeat = 128 - v.shape[0]
|
|
data[k] = torch.cat([v, v[-1:].repeat_interleave(pad_repeat, 0)]).unsqueeze(0) # BxTxJxD
|
|
|
|
args.data.clip_length = data['pos'].shape[1]
|
|
model.set_input(data)
|
|
|
|
target = dict()
|
|
target['pos'] = data['pos'].to(model.device)
|
|
target['rotmat'] = rotation_6d_to_matrix(data['global_xform'].to(model.device))
|
|
target['trans'] = data['trans'].to(model.device)
|
|
target['root_orient'] = data['root_orient'].to(model.device)
|
|
target['cam_R'] = data['cam_R'].to(model.device)
|
|
target['cam_t'] = data['cam_t'].to(model.device)
|
|
target['cam_f'] = data['cam_f'].to(model.device).squeeze(0)
|
|
target['cam_center'] = data['cam_center'].to(model.device).squeeze(0)
|
|
target['joints2d'] = data['joints2d'].to(model.device)
|
|
target['joints3d'] = data['joints3d'].to(model.device)
|
|
# target['verts3d'] = data['verts3d'].to(model.device)
|
|
target['betas'] = data['betas'].to(model.device)
|
|
target['save_path'] = data['save_path']
|
|
target['config_type'] = data['config_type']
|
|
target['handedness'] = data['is_right']
|
|
target['vis_mask'] = data['vis_mask']
|
|
target['is_right'] = data['is_right']
|
|
|
|
####################
|
|
# vis for debugging
|
|
# rhand_trans_padded = target['trans']
|
|
# rhand_orient_padded = matrix_to_axis_angle(rotation_6d_to_matrix(target['root_orient']))
|
|
# rhand_betas_padded = target['betas'][:, 0, ]
|
|
# is_right = target['is_right']
|
|
# rhand_pose_padded = fk.global_to_local(target['rotmat'].view(-1, HAND_JOINT_NUM, 3, 3)) # (B x T, J, 3, 3)
|
|
# rhand_pose_padded = rhand_pose_padded.view(1*args.nsubject, -1, HAND_JOINT_NUM, 3, 3) # (B x T, J, 3, 3)
|
|
# rhand_pose_padded = matrix_to_axis_angle(rhand_pose_padded)[:, :, 1:]
|
|
# print(rhand_pose_padded.shape)
|
|
|
|
# # print(rhand_trans_padded.shape, rhand_orient_padded.shape, rhand_pose_padded.reshape(1, 128, 45).shape, is_right.shape, rhand_betas_padded.shape)
|
|
# rh_mano_out = pred_mano(hand_model, rhand_trans_padded, rhand_orient_padded, rhand_pose_padded.reshape(1, 128, 45), is_right, rhand_betas_padded, only_right=False)
|
|
|
|
# joints3d_pred = rh_mano_out['joints3d'].view(1, 128, -1, 3)
|
|
# vertices_pred = rh_mano_out['verts3d'].view(1, 128, -1, 3)
|
|
# print(joints3d_pred.shape, vertices_pred.shape)
|
|
|
|
# cam_R = torch.tensor(res_dict['cam_R'][idx]).cuda()[None]
|
|
# cam_t = torch.tensor(res_dict['cam_t'][idx]).cuda()[None]
|
|
# cam_center = torch.tensor(res_dict['intrins'][2:][None]).repeat(128, 1).cuda() # (T, 2)
|
|
# cam_f = torch.tensor(res_dict['intrins'][:2][None]).repeat(128, 1).cuda() # (T, 2)
|
|
|
|
# print(joints3d_pred.shape, cam_R.shape, cam_t.shape, cam_f.shape, cam_center.shape)
|
|
# joints2d_pred = reproject(
|
|
# joints3d_pred, cam_R, cam_t, cam_f, cam_center
|
|
# )
|
|
|
|
# # 2d keypoints debugging.
|
|
# openpose_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
|
|
# gt_indices = openpose_indices
|
|
# num_hand = 1
|
|
# num_frames = joints2d_pred.shape[1]
|
|
# for x in range(num_frames):
|
|
|
|
# vit_img = cv2.imread(f'/data/home/x/code/hand/Dyn-HaMR/test/images/dance_old/' + str(x+1).zfill(6) + '.jpg')
|
|
# print(f'/data/home/x/code/hand/Dyn-HaMR/test/images/dance_old/' + str(x+1).zfill(6) + '.jpg')
|
|
# all_vit_2d = joints2d_pred[:, x].detach().cpu().numpy() # observed_data["joints2d"][:, x].cpu().numpy()
|
|
# v = np.ones((len(all_vit_2d), 21, 1))
|
|
# all_vit_2d = np.concatenate((all_vit_2d, v), axis=-1)
|
|
# for i in range(num_hand):
|
|
# body_keypoints_2d = all_vit_2d[i, :21].copy()
|
|
# for op, gt in zip(openpose_indices, gt_indices):
|
|
# if all_vit_2d[i, gt, -1] > body_keypoints_2d[op, -1]:
|
|
# body_keypoints_2d[op] = all_vit_2d[i, gt]
|
|
# vit_img = render_openpose(vit_img, body_keypoints_2d)
|
|
# cv2.imwrite(f'/data/home/x/code/hand/Dyn-HaMR/' + str(x+1).zfill(6) + '.jpg', vit_img)
|
|
|
|
R, T, P, Be, DR = motion_reconstruction(hand_model, target, args.save_path, steps=[1.0], idx=idx)
|
|
|
|
# vis
|
|
is_right = target['is_right'].clone()
|
|
rh_mano_out = pred_mano(hand_model, T, R, P, is_right, Be, only_right=False)
|
|
joints3d_pred = rh_mano_out['joints3d'].view(1, 128, -1, 3)
|
|
vertices_pred = rh_mano_out['verts3d'].view(1, 128, -1, 3)
|
|
joints2d_pred = reproject(
|
|
joints3d_pred, target['cam_R'], target['cam_t'], target['cam_f'], target['cam_center']
|
|
)
|
|
|
|
R_list.append(R.detach().cpu().numpy())
|
|
T_list.append(T.detach().cpu().numpy())
|
|
P_list.append(P.detach().cpu().numpy())
|
|
Be_list.append(Be.detach().cpu().numpy())
|
|
DR_list.append(DR.detach().cpu().numpy())
|
|
|
|
save_keys = ['root_orient', 'trans', 'latent_pose', 'is_right', 'init_body_pose', 'world_scale', 'betas', 'pose_body', 'cam_R', 'cam_t', 'intrins']
|
|
R_list = np.vstack(R_list)
|
|
T_list = np.vstack(T_list)
|
|
P_list = np.vstack(P_list)
|
|
Be_list = np.vstack(Be_list)
|
|
DR_list = np.vstack(DR_list)
|
|
res_dict['root_orient'] = R_list
|
|
res_dict['trans'] = T_list
|
|
res_dict['latent_pose'] = P_list
|
|
num_hands = len(P_list)
|
|
res_dict['pose_body'] = P_list.reshape(num_hands, 128, 15, 3)
|
|
res_dict['betas'] = Be_list
|
|
res_dict['decode_root'] = DR_list
|
|
pred_save_path = os.path.join(args.save_path, os.path.basename(args.vid_path).split('.')[0] + f'_000000_world_results.npz')
|
|
|
|
for i in res_dict.keys():
|
|
print(i, res_dict[i].shape)
|
|
np.savez(pred_save_path, **res_dict)
|
|
|
|
def run_quantitative_evaluation(pred_npz_path, gt_dict, pymafx_npz_path, viz_flag, misc={}):
|
|
|
|
from eval import alignment
|
|
|
|
out_dir = "/".join(pred_npz_path.split("/")[:-1])
|
|
evaluator_object = alignment.Evaluator(out_dir)
|
|
|
|
# load predictions
|
|
pred_dict = dict(np.load(pred_npz_path))
|
|
|
|
# change keynames in gt_dict for quant evaluation
|
|
gt_dict["joints_3d"] = torch.tensor(np.array(gt_dict["joints_3d"]))
|
|
gt_dict["joints_2d"] = torch.tensor(np.array(gt_dict["joints_2d"]))
|
|
|
|
|
|
pred_dict, gt_dict = process_gt(pred_dict, gt_dict, args.dataname)
|
|
|
|
if pymafx_npz_path is not None:
|
|
pymafx_dict = dict(np.load(pymafx_npz_path, allow_pickle=True))
|
|
pymafx_dict, gt_dict = process_gt(pymafx_dict, gt_dict, args.dataname)
|
|
else:
|
|
pymafx_dict = None
|
|
|
|
|
|
# MPJPE
|
|
no_align = alignment.PointError(alignment_object=alignment.NoAlignment(), return_aligned=True)
|
|
procrustes_align = alignment.PointError(alignment_object=alignment.ProcrustesAlignment(), return_aligned=True)
|
|
root_align = alignment.PointError(alignment_object=alignment.RootAlignment(), return_aligned=True)
|
|
|
|
# ACCELERATION
|
|
no_align_accel = alignment.AccelError(alignment_object=alignment.NoAlignment(), return_aligned=True)
|
|
root_align_accel = alignment.AccelError(alignment_object=alignment.RootAlignment(), return_aligned=True)
|
|
procrustes_align_accel = alignment.AccelError(alignment_object=alignment.ProcrustesAlignment(), return_aligned=True)
|
|
|
|
# F-SCORE
|
|
f_thresholds = np.array([5/1000, 15/1000])
|
|
root_align_f = alignment.FScores(thresholds=f_thresholds, alignment_object=alignment.RootAlignment(), return_aligned=True)
|
|
procrustes_align_f = alignment.FScores(thresholds=f_thresholds, alignment_object=alignment.ProcrustesAlignment(), return_aligned=True)
|
|
|
|
# scale alignment and procrustes alignment are the same. Only input shapes are different. SO, there is no need to use.
|
|
align_dict_3d = {"no_align": no_align, "root_align": root_align, "procrustes_align": procrustes_align}
|
|
align_dict_fscore = {"root_align": root_align_f, "procrustes_align": procrustes_align_f}
|
|
align_dict_2d = {"no_align": no_align}
|
|
|
|
align_dict_accel_score = {"no_align": no_align_accel, "root_align": root_align_accel, "procrustes_align": procrustes_align_accel}
|
|
|
|
metrics = {"mpjpe_3d": align_dict_3d, "mpjpe_2d": align_dict_2d, "acc_err": align_dict_accel_score, "f_score": align_dict_fscore}
|
|
|
|
save_quantitative_evaluation(evaluator_object, pred_dict, pymafx_dict=pymafx_dict, gt_dict=gt_dict, metrics=metrics, viz_flag=viz_flag, misc=misc)
|
|
|
|
return
|
|
|
|
|
|
def latent_optimization(hand_model, target, T=None, z_l=None, z_g=None, pose=None):
|
|
|
|
if T is None:
|
|
T = torch.arange(args.data.clip_length)
|
|
|
|
cam_R = target['cam_R'].clone()
|
|
cam_t = target['cam_t'].clone()
|
|
cam_f = target['cam_f'].clone()
|
|
cam_center = target['cam_center'].clone()
|
|
is_right = target['is_right'].clone()
|
|
|
|
optim_trans = target['trans'].clone()
|
|
optim_root_orient = target["root_orient"].clone()
|
|
|
|
# mp_bbox_conf = target["mediapipe_bbox_conf"]
|
|
vis_mask = torch.tensor(target["vis_mask"])
|
|
|
|
B, seqlen, _ = optim_trans.shape
|
|
optim_trans.requires_grad = True
|
|
optim_root_orient.requires_grad = True
|
|
|
|
if not pose is None:
|
|
optim_pose = pose.clone()
|
|
optim_pose.requires_grad = False
|
|
else:
|
|
optim_pose = pose
|
|
|
|
z_global = torch.zeros_like(z_g).to(z_g)
|
|
z_global.requires_grad = False
|
|
|
|
init_z_l = z_l.clone().detach()
|
|
init_z_l.requires_grad = False
|
|
|
|
full_cam_R = cam_R
|
|
full_cam_t = cam_t
|
|
|
|
# take pymafx mean as starting point
|
|
mean_betas = target["betas"].mean(dim=1).mean(dim=0)
|
|
|
|
if args.opt_betas:
|
|
betas = Variable(mean_betas.clone().unsqueeze(0).repeat_interleave(args.nsubject, 0), requires_grad=True)
|
|
logger.info(f'Optimizing betas: {betas}')
|
|
else:
|
|
betas = mean_betas.clone().unsqueeze(0).repeat_interleave(args.nsubject, 0)
|
|
betas.requires_grad = False
|
|
|
|
# print(target["betas"].shape)
|
|
# print('args.opt_betas: ', args.opt_betas, betas.shape, target["betas"].shape, mean_betas.shape)
|
|
stg_configs = [args.stg1]
|
|
|
|
if hasattr(args, 'stg2'):
|
|
stg_configs.append(args.stg2)
|
|
|
|
if hasattr(args, 'stg3'):
|
|
stg_configs.append(args.stg3)
|
|
|
|
if hasattr(args, 'stg4'):
|
|
stg_configs.append(args.stg4)
|
|
|
|
if hasattr(args, 'stg5'):
|
|
stg_configs.append(args.stg5)
|
|
|
|
stg_int_results = (full_cam_R, optim_trans, optim_root_orient, z_l, z_global, betas,
|
|
target, B, seqlen, mean_betas, T, full_cam_R, full_cam_t)
|
|
|
|
joblib.dump(stg_int_results, f'{args.pkl_output_dir}/stg_0.pkl')
|
|
logger.info(f'Saved intermediate results to {args.pkl_output_dir}/stg_0.pkl')
|
|
|
|
is_nan_loss = False
|
|
iter = 0
|
|
|
|
# iterate over different optimization steps
|
|
while iter < len(stg_configs):
|
|
stg_conf = stg_configs[iter]
|
|
logger.info(f'Stage {iter+1}: Learning rate: {stg_conf.lr}')
|
|
stg_id = iter
|
|
|
|
# break is better than continue here
|
|
if stg_conf.niters == 0 and iter!=0:
|
|
break
|
|
# this corresponds to encode-decode stage, we need to calculate joints2d, joints3d etc.
|
|
elif stg_conf.niters == 0 and iter == 0:
|
|
logger.info('Encode-Decode case')
|
|
# cannot plot loss this case
|
|
args.plot_loss = False
|
|
break
|
|
|
|
if is_nan_loss:
|
|
# will give error here
|
|
prev_stg_results = joblib.load(f'{args.pkl_output_dir}/stg_{stg_id}.pkl')
|
|
_, optim_trans, optim_root_orient, z_l, z_global, betas, target, B, seqlen, mean_betas, T, \
|
|
full_cam_R, full_cam_t = prev_stg_results
|
|
|
|
stg_results = optim_step(hand_model, stg_conf, stg_id, z_l, z_global, betas, target,
|
|
B, seqlen, optim_trans, optim_root_orient, init_z_l, mean_betas,
|
|
T, full_cam_R, full_cam_t, cam_f, cam_center, is_right, vis_mask=vis_mask.to("cuda"), pose=optim_pose)
|
|
|
|
if isinstance(stg_results, int):
|
|
is_nan_loss = True
|
|
logger.error(f'[Stage {stg_id+1}] NaN loss detected, restarting stage {stg_id+1}')
|
|
logger.warning(f'Decreasing learning rate by 0.5 for the current stage')
|
|
stg_configs[stg_id].lr *= 0.5
|
|
|
|
else:
|
|
z_l, z_global, optim_cam_R, optim_cam_t, optim_trans, optim_root_orient, betas, optim_pose = stg_results
|
|
|
|
full_cam_R = optim_cam_R # matrix_to_rotation_6d(optim_cam_R.detach())
|
|
full_cam_t = optim_cam_t.detach()
|
|
|
|
stg_int_results = (full_cam_R, optim_trans, optim_root_orient,
|
|
z_l, z_global, betas, target, B, seqlen,
|
|
mean_betas, T, full_cam_R, full_cam_t)
|
|
|
|
joblib.dump(stg_int_results, f'{args.pkl_output_dir}/stg_{stg_id+1}.pkl')
|
|
logger.info(f'Saved intermediate results to {args.pkl_output_dir}/stg_{stg_id+1}.pkl')
|
|
|
|
iter += 1
|
|
|
|
if args.plot_loss:
|
|
plot_list = []
|
|
|
|
for num in range(iter):
|
|
loss_i = joblib.load(f'{args.pkl_output_dir}/stage_{num}_loss.pkl')
|
|
plt.figure()
|
|
|
|
for k, v in loss_i.items():
|
|
if not v == []:
|
|
plt.plot(v, label=k)
|
|
plt.legend()
|
|
plt.title(f'Stage {num}')
|
|
plt.savefig(f'{args.pkl_output_dir}/stage_{num}_loss.jpg')
|
|
|
|
# concatenate all the losses
|
|
plot_list.append(cv2.imread(f'{args.pkl_output_dir}/stage_0_loss.jpg'))
|
|
|
|
plt_concat = np.concatenate(plot_list, axis=0)
|
|
cv2.imwrite(f'{args.pkl_output_dir}/all_stages_loss.jpg', plt_concat)
|
|
|
|
return z_l, z_global, betas, optim_root_orient, optim_trans, full_cam_R, full_cam_t, cam_f, cam_center, optim_pose
|
|
|
|
|
|
def optim_step_new(hand_model, stg_conf, stg_id, z_l, z_g, betas, target, B,
|
|
seqlen, trans, root_orient, init_z_l, mean_betas, T, cam_R, cam_t, cam_f, cam_center, is_right, vis_mask=None, pose=None):
|
|
def closure():
|
|
optimizer.zero_grad()
|
|
|
|
if motion_prior_type == "hmp":
|
|
output = model.decode(z_l, z_g, length=args.data.clip_length, step=1)
|
|
|
|
for k, v in output.items():
|
|
if torch.isnan(v).any():
|
|
logger.warning(f'{k} in output is NaN, skipping this stage...')
|
|
return 0
|
|
|
|
# instead of latent code, work with pose, it is in global coordinates
|
|
else:
|
|
raise ValueError
|
|
output = {"rotmat": axis_angle_to_matrix(pose)}
|
|
|
|
output['betas'] = betas[:, None, None, :].repeat(1, B, seqlen, 1)
|
|
global_trans = trans.clone().reshape(args.nsubject, B, seqlen, 3)
|
|
global_trans = global_trans.reshape(B*args.nsubject, seqlen, 3)
|
|
|
|
output['trans'] = global_trans
|
|
output['root_orient'] = root_orient
|
|
|
|
local_rotmat = fk.global_to_local(output['rotmat'].view(-1, HAND_JOINT_NUM, 3, 3)) # (B x T, J, 3, 3)
|
|
local_rotmat = local_rotmat.view(B*args.nsubject, -1, HAND_JOINT_NUM, 3, 3) # (B x T, J, 3, 3)
|
|
body_pose = matrix_to_axis_angle(local_rotmat)[:, :, 1:]
|
|
|
|
rh_mano_out = pred_mano(hand_model, output['trans'], matrix_to_axis_angle(rotation_6d_to_matrix(root_orient)), body_pose.reshape(1, 128, 45), is_right, betas, only_right=False)
|
|
|
|
joints3d = rh_mano_out['joints3d'].view(B, seqlen, -1, 3)
|
|
vertices3d = rh_mano_out['verts3d'].view(B, seqlen, -1, 3)
|
|
|
|
# print(joints3d.shape, cam_R.shape)
|
|
joints2d_pred = reproject(
|
|
joints3d, cam_R, cam_t, cam_f, cam_center
|
|
)
|
|
output['joints2d'] = joints2d_pred
|
|
output['joints3d'] = joints3d
|
|
|
|
############################
|
|
# 2d keypoints debugging.
|
|
# openpose_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
|
|
# gt_indices = openpose_indices
|
|
# num_hand = 1
|
|
# num_frames = joints2d_pred.shape[1]
|
|
# print(target['joints2d'].shape)
|
|
# for x in range(num_frames):
|
|
|
|
# raise ValueError
|
|
###########################
|
|
|
|
_bbox_conf_ = None
|
|
# _bbox_conf_[_bbox_conf_<0.6] = 0.0
|
|
|
|
local_rotmat_gt = target['rotmat']
|
|
loss_dict = {}
|
|
|
|
if stg_conf.lambda_rot > 0:
|
|
mano_joint_conf = torch.zeros_like(target['rotmat'][..., :1, 0])
|
|
|
|
for si in range(16):
|
|
op_conf = [target['joints2d'][:, :, si, 2]]
|
|
max_conf = torch.stack(op_conf, dim=0).max(0).values
|
|
mano_joint_conf[:, :, si] = max_conf.unsqueeze(-1)
|
|
|
|
# use bbox conf if that is the case
|
|
if _bbox_conf_ is not None:
|
|
bbox_coef = torch.repeat_interleave(_bbox_conf_[..., None], dim=2, repeats=15)
|
|
|
|
rot_loss = L_rot(local_rotmat[:, :, 1:],
|
|
local_rotmat_gt[:, :, 1:],
|
|
T, conf=bbox_coef)
|
|
else:
|
|
rot_loss = L_rot(local_rotmat[:, :, 1:],
|
|
local_rotmat_gt[:, :, 1:],
|
|
T, conf=mano_joint_conf[:, :, 1:])
|
|
|
|
loss_dict['rot'] = stg_conf.lambda_rot * rot_loss
|
|
|
|
if stg_conf.lambda_bio > 0:
|
|
assert len(output['joints3d']) == 1
|
|
joints = output['joints3d'][0] # (128,21,3)
|
|
loss_total, _ = bmc.compute_loss(joints)
|
|
# print("loss_total=", loss_total)
|
|
# print("loss_dict=", loss_dict)
|
|
loss_dict['bio'] = stg_conf.lambda_bio * loss_total
|
|
|
|
if stg_conf.lambda_consistency > 0:
|
|
cur_loss = pose_prior_loss(body_pose)
|
|
loss_dict['pose_prior'] = stg_conf.lambda_reproj * cur_loss
|
|
|
|
if stg_conf.lambda_reproj > 0:
|
|
reproj_loss = joints2d_loss(joints2d_obs=target['joints2d'], joints2d_pred=joints2d_pred, bbox_conf=_bbox_conf_)
|
|
loss_dict['reproj'] = stg_conf.lambda_reproj * reproj_loss
|
|
|
|
if stg_conf.lambda_orient > 0:
|
|
orient_loss = L_orient(output['root_orient'], target['root_orient'], T, bbox_conf=_bbox_conf_)
|
|
loss_dict['orient'] = stg_conf.lambda_orient * orient_loss
|
|
|
|
if stg_conf.lambda_trans > 0:
|
|
trans_loss = L_trans(output['trans'], target['trans'], T, bbox_conf=_bbox_conf_)
|
|
loss_dict['trans'] = stg_conf.lambda_trans * trans_loss
|
|
|
|
if stg_conf.lambda_rot_smooth > 0:
|
|
rot_smooth_l = rot_smooth_loss(local_rotmat)
|
|
loss_dict['rot_sm'] = stg_conf.lambda_rot_smooth * rot_smooth_l
|
|
|
|
if stg_conf.lambda_orient_smooth > 0:
|
|
matrot_root_orient = rotation_6d_to_matrix(root_orient)
|
|
orient_smooth_l = rot_smooth_loss(matrot_root_orient)
|
|
loss_dict['orient_sm'] = stg_conf.lambda_orient_smooth * orient_smooth_l
|
|
|
|
# Smoothness objectives
|
|
if stg_conf.lambda_j3d_smooth > 0:
|
|
joints3d = output['joints3d']
|
|
|
|
j3d_smooth_l = pos_smooth_loss(joints3d)
|
|
loss_dict['j3d_sm'] = stg_conf.lambda_j3d_smooth * j3d_smooth_l
|
|
|
|
if stg_conf.lambda_trans_smooth > 0:
|
|
# tr = mask_data(output['trans'], mask)
|
|
tr = output['trans']
|
|
tr = tr.reshape(args.nsubject, B, seqlen, 3)
|
|
trans_smooth_l = 0
|
|
for sid in range(args.nsubject):
|
|
trans_smooth_l += pos_smooth_loss(tr[sid])
|
|
loss_dict['trans_sm'] = stg_conf.lambda_trans_smooth * trans_smooth_l
|
|
|
|
if stg_conf.lambda_motion_prior > 0:
|
|
|
|
if motion_prior_type == "pca":
|
|
mp_local_loss = L_PCA(pose)
|
|
elif motion_prior_type == "gmm":
|
|
mp_local_loss = L_GMM(pose)
|
|
else:
|
|
mp_local_loss = motion_prior_loss(z_l)
|
|
loss_dict['mot_prior'] = stg_conf.lambda_motion_prior * mp_local_loss
|
|
|
|
|
|
if stg_conf.lambda_init_z_prior > 0:
|
|
zl_init_prior_l = F.mse_loss(z_l, init_z_l)
|
|
loss_dict['init_z_prior'] = stg_conf.lambda_init_z_prior * (zl_init_prior_l)
|
|
|
|
if stg_conf.lambda_pose_prior > 0 and opt.HMP.use_hposer:
|
|
raise NotImplementedError
|
|
# loss_dict['pose_prior'] = stg_conf.lambda_pose_prior * L_pose_prior(output)
|
|
|
|
if hasattr(stg_conf, 'lambda_batch_cs'):
|
|
if stg_conf.lambda_batch_cs > 0:
|
|
if args.overlap_len == 0:
|
|
logger.warning('Batch consistency won\'t be effective since overlap_len is 0')
|
|
if B > 1:
|
|
# joints3d = mask_data(output['joints3d'], mask)
|
|
joints3d = joints3d.reshape(args.nsubject, B, seqlen, -1, 3)
|
|
batch_cs_l = 0
|
|
for sid in range(args.nsubject):
|
|
batch_cs_l += L_pos(joints3d[sid, :-1, -args.overlap_len:], joints3d[sid, 1:, :args.overlap_len], T)
|
|
loss_dict['batch_cs'] = stg_conf.lambda_batch_cs * batch_cs_l
|
|
else:
|
|
if i < 5:
|
|
logger.warning('Batch consistency won\'t be effective since batch size is 1')
|
|
|
|
if hasattr(stg_conf, 'betas_prior'):
|
|
if stg_conf.betas_prior > 0:
|
|
if betas is None:
|
|
logger.error('Cannot compute betas prior since args.opt_betas is False')
|
|
betas_prior_l = torch.pow(betas - mean_betas, 2).mean()
|
|
loss_dict['betas_prior'] = stg_conf.betas_prior * betas_prior_l
|
|
|
|
|
|
loss = sum(loss_dict.values())
|
|
loss_dict['loss'] = loss
|
|
|
|
# copy loss values to loss_dict_by_step
|
|
for k, v in loss_dict.items():
|
|
loss_dict_by_step[k].append(v.detach().item())
|
|
|
|
if not torch.isnan(loss):
|
|
loss.backward()
|
|
return loss
|
|
else:
|
|
logger.warning('Loss is NaN, skipping this stage')
|
|
raise ValueError
|
|
|
|
print(stg_id, i, stg_conf.niters)
|
|
loss_log_str = f'Stage {stg_id+1} [{i:03d}/{stg_conf.niters}]'
|
|
for k, v in loss_dict.items():
|
|
loss_dict[k] = v.item()
|
|
loss_log_str += f'{k}: {v.item():.3f}\t'
|
|
logger.info(loss_log_str)
|
|
|
|
logger.info(f'Running optimization stage {stg_id+1} ...')
|
|
opt_params = []
|
|
for param in stg_conf.opt_params:
|
|
if param == 'root_orient':
|
|
opt_params.append(root_orient)
|
|
elif param == 'trans':
|
|
opt_params.append(trans)
|
|
elif param == 'z_l':
|
|
opt_params.append(z_l)
|
|
elif param == 'pose':
|
|
opt_params.append(pose)
|
|
elif param == 'betas':
|
|
if not args.opt_betas:
|
|
continue
|
|
if betas is None:
|
|
logger.error('Cannot optimize betas if args.opt_betas is False')
|
|
opt_params.append(betas)
|
|
else:
|
|
raise ValueError(f'Unknown parameter {param}')
|
|
|
|
for param in opt_params:
|
|
param.requires_grad = True
|
|
|
|
# optimizer = torch.optim.Adam(opt_params, lr=stg_conf.lr)
|
|
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.scheduler.step_size, args.scheduler.gamma)
|
|
LINE_SEARCH = "strong_wolfe"
|
|
lr = 1.0
|
|
lbfgs_max_iter = 20
|
|
optimizer = torch.optim.LBFGS(
|
|
opt_params, max_iter=lbfgs_max_iter, lr=lr, line_search_fn=LINE_SEARCH
|
|
)
|
|
|
|
def mask_data(data, mask):
|
|
ml = len(mask.shape)
|
|
dl = len(data.shape)
|
|
for _ in range(dl-ml):
|
|
mask = mask[..., None]
|
|
return data * mask
|
|
|
|
loss_dict_by_step = {"rot": [], "reproj": [], "rot_sm": [], "orient_sm": [], 'betas_prior': [], "j3d_sm": [], "pose_prior": [],
|
|
"trans_sm": [], "mot_prior": [], "init_z_prior": [], "orient": [], "trans": [], "loss": [], "bio": []}
|
|
|
|
# optimize the z_l and root_orient, pos, trans, 2d kp objectives
|
|
start_time = time.time()
|
|
print('start latent optimization...')
|
|
for i in range(stg_conf.niters):
|
|
optimizer.step(closure)
|
|
|
|
end_time = time.time()
|
|
|
|
# save the loss dict.
|
|
joblib.dump(loss_dict_by_step, open(os.path.join(args.pkl_output_dir, f'stage_{stg_id}_loss.pkl'), 'wb'))
|
|
|
|
print(f'Stage {stg_id+1} finished in {time.strftime("%H:%M:%S", time.gmtime(end_time - start_time))}')
|
|
|
|
if not betas is None:
|
|
logger.info(f'mean_betas: {mean_betas.detach().cpu().numpy()}')
|
|
logger.info(f'betas: {betas.detach().cpu().numpy()}')
|
|
|
|
return z_l, z_g, cam_R, cam_t, trans, root_orient, betas, pose
|
|
|
|
def optim_step(hand_model, stg_conf, stg_id, z_l, z_g, betas, target, B,
|
|
seqlen, trans, root_orient, init_z_l, mean_betas, T, cam_R, cam_t, cam_f, cam_center, is_right, vis_mask=None, pose=None):
|
|
|
|
logger.info(f'Running optimization stage {stg_id+1} ...')
|
|
opt_params = []
|
|
for param in stg_conf.opt_params:
|
|
if param == 'root_orient':
|
|
opt_params.append(root_orient)
|
|
elif param == 'trans':
|
|
opt_params.append(trans)
|
|
elif param == 'z_l':
|
|
opt_params.append(z_l)
|
|
elif param == 'pose':
|
|
opt_params.append(pose)
|
|
elif param == 'betas':
|
|
if not args.opt_betas:
|
|
continue
|
|
if betas is None:
|
|
logger.error('Cannot optimize betas if args.opt_betas is False')
|
|
opt_params.append(betas)
|
|
else:
|
|
raise ValueError(f'Unknown parameter {param}')
|
|
|
|
for param in opt_params:
|
|
param.requires_grad = True
|
|
|
|
optimizer = torch.optim.Adam(opt_params, lr=stg_conf.lr)
|
|
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.scheduler.step_size, args.scheduler.gamma)
|
|
|
|
def mask_data(data, mask):
|
|
ml = len(mask.shape)
|
|
dl = len(data.shape)
|
|
for _ in range(dl-ml):
|
|
mask = mask[..., None]
|
|
return data * mask
|
|
|
|
loss_dict_by_step = {"rot": [], "reproj": [], "rot_sm": [], "orient_sm": [], 'betas_prior': [], "j3d_sm": [], "pose_prior": [],
|
|
"trans_sm": [], "mot_prior": [], "init_z_prior": [], "orient": [], "trans": [], "loss": [], "bio": []}
|
|
|
|
# optimize the z_l and root_orient, pos, trans, 2d kp objectives
|
|
start_time = time.time()
|
|
print(f'start latent optimization... {stg_conf.niters} iters')
|
|
for i in range(stg_conf.niters):
|
|
optimizer.zero_grad()
|
|
|
|
if motion_prior_type == "hmp":
|
|
output = model.decode(z_l, z_g, length=args.data.clip_length, step=1)
|
|
|
|
for k, v in output.items():
|
|
if torch.isnan(v).any():
|
|
logger.warning(f'{k} in output is NaN, skipping this stage...')
|
|
return 0
|
|
|
|
# instead of latent code, work with pose, it is in global coordinates
|
|
else:
|
|
raise ValueError
|
|
output = {"rotmat": axis_angle_to_matrix(pose)}
|
|
|
|
output['betas'] = betas[:, None, None, :].repeat(1, B, seqlen, 1)
|
|
global_trans = trans.clone().reshape(args.nsubject, B, seqlen, 3)
|
|
global_trans = global_trans.reshape(B*args.nsubject, seqlen, 3)
|
|
|
|
output['trans'] = global_trans
|
|
output['root_orient'] = root_orient
|
|
|
|
local_rotmat = fk.global_to_local(output['rotmat'].view(-1, HAND_JOINT_NUM, 3, 3)) # (B x T, J, 3, 3)
|
|
local_rotmat = local_rotmat.view(B*args.nsubject, -1, HAND_JOINT_NUM, 3, 3) # (B x T, J, 3, 3)
|
|
body_pose = matrix_to_axis_angle(local_rotmat)[:, :, 1:]
|
|
|
|
rh_mano_out = pred_mano(hand_model, output['trans'], matrix_to_axis_angle(rotation_6d_to_matrix(root_orient)), body_pose.reshape(1, 128, 45), is_right, betas, only_right=False)
|
|
|
|
joints3d = rh_mano_out['joints3d'].view(B, seqlen, -1, 3)
|
|
vertices3d = rh_mano_out['verts3d'].view(B, seqlen, -1, 3)
|
|
|
|
# print(joints3d.shape, cam_R.shape)
|
|
joints2d_pred = reproject(
|
|
joints3d, cam_R, cam_t, cam_f, cam_center
|
|
)
|
|
output['joints2d'] = joints2d_pred
|
|
output['joints3d'] = joints3d
|
|
|
|
############################
|
|
# 2d keypoints debugging.
|
|
# openpose_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
|
|
# gt_indices = openpose_indices
|
|
# num_hand = 1
|
|
# num_frames = joints2d_pred.shape[1]
|
|
# print(target['joints2d'].shape)
|
|
# for x in range(num_frames):
|
|
|
|
# raise ValueError
|
|
###########################
|
|
|
|
_bbox_conf_ = None
|
|
# _bbox_conf_[_bbox_conf_<0.6] = 0.0
|
|
|
|
local_rotmat_gt = target['rotmat']
|
|
loss_dict = {}
|
|
|
|
if stg_conf.lambda_rot > 0:
|
|
mano_joint_conf = torch.zeros_like(target['rotmat'][..., :1, 0])
|
|
|
|
for si in range(16):
|
|
op_conf = [target['joints2d'][:, :, si, 2]]
|
|
max_conf = torch.stack(op_conf, dim=0).max(0).values
|
|
mano_joint_conf[:, :, si] = max_conf.unsqueeze(-1)
|
|
|
|
# use bbox conf if that is the case
|
|
if _bbox_conf_ is not None:
|
|
bbox_coef = torch.repeat_interleave(_bbox_conf_[..., None], dim=2, repeats=15)
|
|
|
|
rot_loss = L_rot(local_rotmat[:, :, 1:],
|
|
local_rotmat_gt[:, :, 1:],
|
|
T, conf=bbox_coef)
|
|
else:
|
|
rot_loss = L_rot(local_rotmat[:, :, 1:],
|
|
local_rotmat_gt[:, :, 1:],
|
|
T, conf=mano_joint_conf[:, :, 1:])
|
|
|
|
loss_dict['rot'] = stg_conf.lambda_rot * rot_loss
|
|
|
|
if stg_conf.lambda_bio > 0:
|
|
assert len(output['joints3d']) == 1
|
|
joints = output['joints3d'][0] # (128,21,3)
|
|
loss_total, _ = bmc.compute_loss(joints)
|
|
# print("loss_total=", loss_total)
|
|
# print("loss_dict=", loss_dict)
|
|
loss_dict['bio'] = stg_conf.lambda_bio * loss_total
|
|
|
|
if stg_conf.lambda_consistency > 0:
|
|
cur_loss = pose_prior_loss(body_pose)
|
|
loss_dict['pose_prior'] = stg_conf.lambda_reproj * cur_loss
|
|
|
|
if stg_conf.lambda_reproj > 0:
|
|
reproj_loss = joints2d_loss(joints2d_obs=target['joints2d'], joints2d_pred=joints2d_pred, bbox_conf=_bbox_conf_)
|
|
loss_dict['reproj'] = stg_conf.lambda_reproj * reproj_loss
|
|
|
|
if stg_conf.lambda_orient > 0:
|
|
orient_loss = L_orient(output['root_orient'], target['root_orient'], T, bbox_conf=_bbox_conf_)
|
|
loss_dict['orient'] = stg_conf.lambda_orient * orient_loss
|
|
|
|
if stg_conf.lambda_trans > 0:
|
|
trans_loss = L_trans(output['trans'], target['trans'], T, bbox_conf=_bbox_conf_)
|
|
loss_dict['trans'] = stg_conf.lambda_trans * trans_loss
|
|
|
|
if stg_conf.lambda_rot_smooth > 0:
|
|
rot_smooth_l = rot_smooth_loss(local_rotmat)
|
|
loss_dict['rot_sm'] = stg_conf.lambda_rot_smooth * rot_smooth_l
|
|
|
|
if stg_conf.lambda_orient_smooth > 0:
|
|
matrot_root_orient = rotation_6d_to_matrix(root_orient)
|
|
orient_smooth_l = rot_smooth_loss(matrot_root_orient)
|
|
loss_dict['orient_sm'] = stg_conf.lambda_orient_smooth * orient_smooth_l
|
|
|
|
# Smoothness objectives
|
|
if stg_conf.lambda_j3d_smooth > 0:
|
|
joints3d = output['joints3d']
|
|
|
|
j3d_smooth_l = pos_smooth_loss(joints3d)
|
|
loss_dict['j3d_sm'] = stg_conf.lambda_j3d_smooth * j3d_smooth_l
|
|
|
|
if stg_conf.lambda_trans_smooth > 0:
|
|
# tr = mask_data(output['trans'], mask)
|
|
tr = output['trans']
|
|
tr = tr.reshape(args.nsubject, B, seqlen, 3)
|
|
trans_smooth_l = 0
|
|
for sid in range(args.nsubject):
|
|
trans_smooth_l += pos_smooth_loss(tr[sid])
|
|
loss_dict['trans_sm'] = stg_conf.lambda_trans_smooth * trans_smooth_l
|
|
|
|
if stg_conf.lambda_motion_prior > 0:
|
|
|
|
if motion_prior_type == "pca":
|
|
mp_local_loss = L_PCA(pose)
|
|
elif motion_prior_type == "gmm":
|
|
mp_local_loss = L_GMM(pose)
|
|
else:
|
|
mp_local_loss = motion_prior_loss(z_l)
|
|
loss_dict['mot_prior'] = stg_conf.lambda_motion_prior * mp_local_loss
|
|
|
|
|
|
if stg_conf.lambda_init_z_prior > 0:
|
|
zl_init_prior_l = F.mse_loss(z_l, init_z_l)
|
|
loss_dict['init_z_prior'] = stg_conf.lambda_init_z_prior * (zl_init_prior_l)
|
|
|
|
if stg_conf.lambda_pose_prior > 0 and opt.HMP.use_hposer:
|
|
raise NotImplementedError
|
|
# loss_dict['pose_prior'] = stg_conf.lambda_pose_prior * L_pose_prior(output)
|
|
|
|
if hasattr(stg_conf, 'lambda_batch_cs'):
|
|
if stg_conf.lambda_batch_cs > 0:
|
|
if args.overlap_len == 0:
|
|
logger.warning('Batch consistency won\'t be effective since overlap_len is 0')
|
|
if B > 1:
|
|
# joints3d = mask_data(output['joints3d'], mask)
|
|
joints3d = joints3d.reshape(args.nsubject, B, seqlen, -1, 3)
|
|
batch_cs_l = 0
|
|
for sid in range(args.nsubject):
|
|
batch_cs_l += L_pos(joints3d[sid, :-1, -args.overlap_len:], joints3d[sid, 1:, :args.overlap_len], T)
|
|
loss_dict['batch_cs'] = stg_conf.lambda_batch_cs * batch_cs_l
|
|
else:
|
|
if i < 5:
|
|
logger.warning('Batch consistency won\'t be effective since batch size is 1')
|
|
|
|
if hasattr(stg_conf, 'betas_prior'):
|
|
if stg_conf.betas_prior > 0:
|
|
if betas is None:
|
|
logger.error('Cannot compute betas prior since args.opt_betas is False')
|
|
betas_prior_l = torch.pow(betas - mean_betas, 2).mean()
|
|
loss_dict['betas_prior'] = stg_conf.betas_prior * betas_prior_l
|
|
|
|
|
|
loss = sum(loss_dict.values())
|
|
loss_dict['loss'] = loss
|
|
|
|
# copy loss values to loss_dict_by_step
|
|
for k, v in loss_dict.items():
|
|
loss_dict_by_step[k].append(v.detach().item())
|
|
|
|
if not torch.isnan(loss):
|
|
loss.backward()
|
|
torch.nn.utils.clip_grad_norm_(opt_params, 5.0)
|
|
optimizer.step()
|
|
else:
|
|
logger.warning('Loss is NaN, skipping this stage')
|
|
return 0
|
|
|
|
scheduler.step()
|
|
loss_log_str = f'Stage {stg_id+1} [{i:03d}/{stg_conf.niters}]'
|
|
for k, v in loss_dict.items():
|
|
loss_dict[k] = v.item()
|
|
loss_log_str += f'{k}: {v.item():.3f}\t'
|
|
logger.info(loss_log_str)
|
|
|
|
end_time = time.time()
|
|
|
|
# save the loss dict.
|
|
joblib.dump(loss_dict_by_step, open(os.path.join(args.pkl_output_dir, f'stage_{stg_id}_loss.pkl'), 'wb'))
|
|
|
|
print(f'Stage {stg_id+1} finished in {time.strftime("%H:%M:%S", time.gmtime(end_time - start_time))}')
|
|
|
|
if not betas is None:
|
|
logger.info(f'mean_betas: {mean_betas.detach().cpu().numpy()}')
|
|
logger.info(f'betas: {betas.detach().cpu().numpy()}')
|
|
|
|
return z_l, z_g, cam_R, cam_t, trans, root_orient, betas, pose
|
|
|
|
def fitting_prior(obs_data, res_dict, hand_model, opt, data_args, out_dir, device):
|
|
global model, fk, ngpu, hposer, motion_prior_type, pca_aa, gmm_aa, args
|
|
args = Arguments(opt.paths.base_dir, os.path.dirname(__file__), filename=opt.HMP.config)
|
|
|
|
opt.HMP.exp_name = data_args.seq
|
|
|
|
cfg_name = opt.HMP.config.split(".")[0]
|
|
opt.HMP.use_hposer = False
|
|
|
|
args.save_path = out_dir
|
|
args.plot_loss = True
|
|
|
|
args.vid_path = opt.HMP.vid_path
|
|
|
|
args.root = opt.paths.base_dir
|
|
args.dataset_dir = os.path.join(opt.paths.base_dir, '_DATA/hmp_model')
|
|
args.save_dir = os.path.join(opt.paths.base_dir, '_DATA/hmp_model')
|
|
|
|
init_method = args.init_method if hasattr(args, 'init_method') else "pymafx"
|
|
assert init_method in ["metro", "pymafx"]
|
|
|
|
if hasattr(args, 'motion_prior_type'):
|
|
raise ValueError
|
|
else:
|
|
motion_prior_type = "hmp"
|
|
|
|
assert motion_prior_type in ["hmp"]
|
|
|
|
# load hposer
|
|
if opt.HMP.use_hposer:
|
|
raise ValueError
|
|
# hposer, _ = load_hposer()
|
|
# hposer.to("cuda")
|
|
# hposer.eval()
|
|
else:
|
|
hposer = None
|
|
|
|
torch.manual_seed(0)
|
|
np.random.seed(0)
|
|
random.seed(0)
|
|
|
|
ngpu = 1
|
|
|
|
model = Architecture(args, ngpu)
|
|
model.load(optimal=True)
|
|
model.eval()
|
|
|
|
fk = ForwardKinematicsLayer(args)
|
|
multi_stage_opt(opt, device, obs_data, res_dict, hand_model, os.path.join(os.path.dirname(__file__), opt.HMP.config), opt.HMP.exp_name, init_method)
|
|
|
|
|
|
def run_prior(
|
|
cfg,
|
|
dataset,
|
|
out_dir,
|
|
device,
|
|
phases,
|
|
obs_data, hand_model, opt, data_args, prior_out,
|
|
save_dir=None
|
|
):
|
|
|
|
def load_result(res_path_dict):
|
|
def to_torch(obj):
|
|
if isinstance(obj, np.ndarray):
|
|
return torch.from_numpy(obj).float()
|
|
if isinstance(obj, dict):
|
|
return {k: to_torch(v) for k, v in obj.items()}
|
|
if isinstance(obj, (list, tuple)):
|
|
return [to_torch(x) for x in obj]
|
|
return obj
|
|
"""
|
|
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 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
|
|
|
|
save_dir = out_dir if save_dir is None else save_dir
|
|
print("OUT_DIR", out_dir)
|
|
print("SAVE_DIR", save_dir)
|
|
print("VISUALIZING PHASES", phases)
|
|
print("PRIOR OUT", prior_out)
|
|
|
|
phase_results = {}
|
|
phase_max_iters = {}
|
|
for phase in phases:
|
|
res_dir = os.path.join(out_dir, phase)
|
|
if phase == "input":
|
|
res = get_input_dict(dataset)
|
|
it = f"{0:06d}"
|
|
|
|
elif os.path.isdir(res_dir):
|
|
res_path_dict = get_results_paths(res_dir)
|
|
it = sorted(res_path_dict.keys())[-1]
|
|
res = load_result(res_path_dict[it])["world"]
|
|
|
|
else:
|
|
print(f"{res_dir} does not exist, skipping")
|
|
continue
|
|
|
|
out_name = f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}"
|
|
phase_max_iters[phase] = it
|
|
|
|
# out_paths = [f"{out_name}_{view}{out_ext}" for view in render_views]
|
|
# if not overwrite and all(os.path.exists(p) for p in out_paths):
|
|
# print("FOUND OUT PATHS", out_paths)
|
|
# continue
|
|
|
|
phase_results[phase] = out_name, res
|
|
|
|
if len(phase_results) > 0:
|
|
out_names, res_dicts = zip(*phase_results.values())
|
|
|
|
assert len(res_dicts) == 1
|
|
assert len(res_dicts) == len(out_names)
|
|
if res_dicts[0]['trans'].shape[1] != 128:
|
|
from HMP.windowed import fit_windowed
|
|
fit_windowed(obs_data, res_dicts, opt, data_args, prior_out, device, fitting_prior)
|
|
else:
|
|
fitting_prior(obs_data, res_dicts, hand_model, opt, data_args, prior_out, device)
|