1119 lines
42 KiB
Python
1119 lines
42 KiB
Python
import os
|
|
import cv2
|
|
import glob
|
|
import json
|
|
import torch
|
|
import joblib
|
|
import importlib
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
# import scenepic as sp
|
|
from smplx import MANO
|
|
from loguru import logger
|
|
import matplotlib.pyplot as plt
|
|
from PIL import Image, ImageDraw
|
|
from rotations import axis_angle_to_matrix, matrix_to_axis_angle
|
|
import torch.nn.functional as torch_f
|
|
from nemf.losses import pos_smooth_loss
|
|
|
|
AUGMENTED_MANO_CHAIN = np.array([-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11, 0, 13, 14, 15, 3, 6, 12, 9])
|
|
openpose_skeleton = [-1, 0, 1, 2, 3, 0, 5, 6, 7, 0, 9, 10, 11, 0, 13, 14, 15, 0, 17, 18, 19]
|
|
|
|
RIGHT_WRIST_BASE_LOC = torch.tensor([[0.0957, 0.0064, 0.0062]])
|
|
LEFT_WRIST_BASE_LOC = torch.tensor([[-0.0957, 0.0064, 0.0062]])
|
|
MANO_RH_DIR = "./data/body_models/mano/MANO_RIGHT.pkl"
|
|
|
|
|
|
def convert_pred_to_full_img_cam(pare_cam, bbox_height, bbox_center,
|
|
img_w, img_h, focal_length, crop_res=224):
|
|
# Converts weak perspective camera estimated by PARE in
|
|
# bbox coords to perspective camera in full image coordinates
|
|
# from https://arxiv.org/pdf/2009.06549.pdf
|
|
|
|
s, tx, ty = pare_cam[:, 0], pare_cam[:, 1], pare_cam[:, 2]
|
|
|
|
r = bbox_height / crop_res
|
|
|
|
|
|
tz = 2 * focal_length / (r * crop_res * s)
|
|
|
|
cx = 2 * (bbox_center[:, 0] - (img_w / 2.)) / (s * bbox_height)
|
|
cy = 2 * (bbox_center[:, 1] - (img_h / 2.)) / (s * bbox_height)
|
|
|
|
kk = 1.0
|
|
|
|
# fantezi
|
|
tx_final = (tx + cx) / kk
|
|
ty_final = (ty + cy) / kk
|
|
tz_final = tz / kk
|
|
|
|
cam_t = torch.stack([tx_final, ty_final, tz_final], dim=-1)
|
|
|
|
return cam_t
|
|
|
|
# reverse
|
|
# def map_openpose_joints_to_mano():
|
|
# return np.array([0, 5, 6, 7, 9, 10, 11, 17, 18, 19, 13, 14, 15, 1, 2, 3, 4, 8, 12, 16, 20], dtype=np.int32)
|
|
|
|
# def map_mano_joints_to_openpose():
|
|
# return np.array([0, 13, 14, 15, 16, 1, 2, 3, 17, 4, 5, 6, 18, 10, 11, 12, 19, 7, 8, 9, 20], dtype=np.int32)
|
|
|
|
def forward_mano_pymafx(output):
|
|
root_orient = output['root_orient'].to("cuda") # (T, 3)
|
|
poses = output['poses'].to("cuda") # (T, J, 3)
|
|
|
|
T, J, _ = poses.size()
|
|
|
|
hand_model = MANO(model_path=MANO_RH_DIR, flat_hand_mean=True, use_pca=False, batch_size=poses.shape[0]).to('cuda')
|
|
|
|
mano_out = hand_model(betas=output['betas'].view(-1, 10).to("cuda"),
|
|
global_orient=root_orient.to("cuda"),
|
|
hand_pose=poses.view(-1, 45).to("cuda"),
|
|
return_tips=True,
|
|
transl=output['trans'].view(-1, 3).to("cuda"))
|
|
|
|
# mano_out.vertices -= RIGHT_WRIST_BASE_LOC.to(mano_out.vertices)
|
|
# mano_out.joints -= RIGHT_WRIST_BASE_LOC.to(mano_out.joints)
|
|
|
|
return mano_out
|
|
|
|
|
|
def joints2d_loss(joints2d_obs, joints2d_pred, bbox_conf=None):
|
|
'''
|
|
Cam extrinsics are assumed the same for entire sequence
|
|
- cam_t : (B, 3)
|
|
- cam_R : (B, 3, 3)
|
|
'''
|
|
B, T, _, _ = joints2d_obs.size()
|
|
|
|
# compared to observations
|
|
|
|
if bbox_conf is None:
|
|
joints2d_obs_conf = joints2d_obs[:,:,:,2:3]
|
|
else:
|
|
bbox_conf = torch.repeat_interleave(bbox_conf.unsqueeze(-1), 21, dim=-1)
|
|
joints2d_obs_conf = torch.repeat_interleave(bbox_conf.unsqueeze(-1), 2, dim=-1).detach()
|
|
|
|
# weight errors by detection confidence
|
|
robust_sqr_dist = gmof(joints2d_pred[:, :, :, :] - joints2d_obs[:,:,:,:2], sigma=40)
|
|
reproj_err = (joints2d_obs_conf**2) * robust_sqr_dist
|
|
loss = torch.mean(reproj_err)
|
|
return loss
|
|
|
|
def get_seqname_ho3d_v3(path):
|
|
return path.split("/")[-2]
|
|
|
|
def get_seqname_dexycb(path):
|
|
# return SUBJECTNAME/SEQNAME/SUBSEQNAME'
|
|
return "/".join(path.split("/")[-4:-1])
|
|
|
|
def get_seqname_arctic_data(path):
|
|
# return SUBJECTNAME/SEQNAME/VIEVNUM'
|
|
return "/".join(path.split("/")[-4:-1])
|
|
|
|
def get_seqname_in_the_wild(path):
|
|
return path.split("/")[-2].split(".")[0]
|
|
|
|
def compute_seq_intervals(seq_len, split_len=128, overlap_len=16):
|
|
# for seq_len = 256 [0, 128], [112, 240], [224, 256]
|
|
# for seq_len = 320 [0, 128], [112, 240], [224, 320]
|
|
# for seq_len = 565 [0, 128], [112, 240], [224, 352], [336, 464], [448, 565]
|
|
|
|
intervals = []
|
|
for i in range(0, seq_len, split_len - overlap_len):
|
|
start = i
|
|
end = min(i + split_len, seq_len)
|
|
|
|
intervals.append((start, end))
|
|
if (end - start) < split_len:
|
|
break
|
|
return intervals
|
|
|
|
# optimize on pymafx. Optimize translation only.
|
|
def optimize_on_pymafx(pymafx_dict, pymafx_keypoints2d_path):
|
|
|
|
os.makedirs(pymafx_keypoints2d_path, exist_ok=True)
|
|
|
|
seqlen = 128
|
|
|
|
# pad values
|
|
for k, v in pymafx_dict.items():
|
|
if k in ['cam_f', 'cam_center']:
|
|
continue
|
|
else:
|
|
data_split = list(torch.split(v, 128))
|
|
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
|
|
pymafx_dict[k] = torch.stack(data_split, dim=0)
|
|
|
|
rh_opt_trans = torch.tensor(pymafx_dict["trans"].clone(), dtype=torch.float32)
|
|
B = rh_opt_trans.shape[0]
|
|
|
|
cam_R = torch.zeros((seqlen * B, 3, 3), dtype=torch.float32).to(rh_opt_trans)
|
|
cam_R[:, 0, 0] = 1.0
|
|
cam_R[:, 1, 1] = 1.0
|
|
cam_R[:, 2, 2] = 1.0
|
|
cam_t = torch.zeros((seqlen * B, 3), dtype=torch.float32).to(rh_opt_trans)
|
|
|
|
rh_opt_trans.requires_grad = True
|
|
cam_R.requires_grad = False
|
|
cam_t.requires_grad = False
|
|
|
|
# optimize based on 2d keypoints
|
|
opt_params = [rh_opt_trans]
|
|
optimizer = torch.optim.Adam(opt_params, lr=0.3)
|
|
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=200, gamma=0.95, verbose=False)
|
|
|
|
K_iter = 3000
|
|
|
|
# project to image and get 2d keypoints
|
|
# for i in tqdm(range(500)):
|
|
for i in range(K_iter):
|
|
|
|
loss_dict = {}
|
|
optimizer.zero_grad()
|
|
|
|
# values change in every iteration
|
|
output = {"poses": pymafx_dict["poses"].view(-1, 15, 3),
|
|
"root_orient": pymafx_dict["root_orient"].view(-1, 3),
|
|
"betas": pymafx_dict["betas"].view(-1, 10),
|
|
"trans": rh_opt_trans}
|
|
|
|
# rotmat ==> (B, T, J, 3, 3), root_orient ==> (B, T, 3, 3), trans ==> (B, T, 3)
|
|
rh_mano_out = forward_mano_pymafx(output)
|
|
joints3d = rh_mano_out.joints
|
|
|
|
joints2d_pred = get_joints2d(joints3d_pred=joints3d.reshape(B, seqlen, -1, 3),
|
|
cam_t=cam_t.unsqueeze(0).repeat_interleave(1, 0),
|
|
cam_R=cam_R.unsqueeze(0).repeat_interleave(1, 0),
|
|
cam_f=torch.tensor([5000., 5000.]),
|
|
cam_center=pymafx_dict['cam_center'][0])
|
|
|
|
reproj_loss = joints2d_loss(joints2d_obs=pymafx_dict['keypoints2d'].reshape(B, seqlen, -1, 3), joints2d_pred=joints2d_pred)
|
|
|
|
loss_dict['reproj'] = 0.01 * reproj_loss
|
|
loss_dict['trans_sm'] = 400 * pos_smooth_loss(rh_opt_trans)
|
|
loss_dict['j3d_sm'] = 100 * pos_smooth_loss(joints3d.reshape(B, seqlen, -1, 3))
|
|
|
|
loss = sum(loss_dict.values())
|
|
loss.backward()
|
|
optimizer.step()
|
|
scheduler.step()
|
|
|
|
loss_log_str = f'[{i:03d}/{K_iter}]'
|
|
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)
|
|
|
|
pymafx_dict["trans"] = rh_opt_trans.detach()
|
|
|
|
return pymafx_dict
|
|
|
|
# process ground truth and predictions. Discard frames with no or invalid ground truth
|
|
def process_gt(pred_dict, gt_dict, dataset_name):
|
|
|
|
# DEXYCB
|
|
if dataset_name == "DexYCB":
|
|
|
|
jts3d = - torch.ones((1, 21, 3))
|
|
jts2d = - torch.ones((1, 21, 2))
|
|
|
|
invalid_indices3d = torch.where(torch.all(torch.all(torch.eq(gt_dict["joints_3d"], jts3d), dim=1), dim=1))[0]
|
|
invalid_indices2d = torch.where(torch.all(torch.all(torch.eq(gt_dict["joints_2d"], jts2d), dim=1), dim=1))[0]
|
|
|
|
assert torch.all(torch.eq(invalid_indices3d, invalid_indices2d)), "Inconsistency in ground truth data"
|
|
|
|
valid_indices = np.setdiff1d(np.arange(gt_dict["joints_3d"].shape[0]), invalid_indices3d)
|
|
|
|
# discard those frames from frame_id
|
|
gt_dict["frame_id"] = list(set(gt_dict["frame_id"]).difference(set(np.array(invalid_indices3d))))
|
|
|
|
# find indices of frames with value -1. This indicates the ground truth is not available
|
|
gt_dict["joints_3d"] = gt_dict["joints_3d"][valid_indices]
|
|
gt_dict["vertices_3d"] = gt_dict["vertices_3d"][valid_indices]
|
|
gt_dict["joints_2d"] = gt_dict["joints_2d"][valid_indices]
|
|
gt_dict["poses"] = np.array(gt_dict["poses"])[valid_indices]
|
|
|
|
# cast to seq length
|
|
pred_dict["joints_3d"] = pred_dict["joints_3d"][:gt_dict['num_frames']]
|
|
pred_dict["joints_2d"] = pred_dict["joints_2d"][:gt_dict['num_frames']]
|
|
|
|
|
|
return pred_dict, gt_dict
|
|
|
|
|
|
def gmof(res, sigma):
|
|
"""
|
|
Geman-McClure error function
|
|
- residual
|
|
- sigma scaling factor
|
|
"""
|
|
x_squared = res ** 2
|
|
sigma_squared = sigma ** 2
|
|
return (sigma_squared * x_squared) / (sigma_squared + x_squared)
|
|
|
|
|
|
def get_joints2d(joints3d_pred, cam_t, cam_R, cam_f, cam_center):
|
|
B, T, _, _ = joints3d_pred.size()
|
|
|
|
# project points to 2D
|
|
joints2d_pred = perspective_projection(joints3d_pred.reshape((B*T, -1, 3)),
|
|
cam_R.reshape((B*T, 3, 3)),
|
|
cam_t.reshape((B*T, 3)),
|
|
cam_f.unsqueeze(0).unsqueeze(0).repeat(B, T, 1).reshape((B*T, 2)),
|
|
cam_center.unsqueeze(0).unsqueeze(0).repeat(B, T, 1).reshape((B*T, 2))).reshape(B, T, -1, 2)
|
|
return joints2d_pred
|
|
|
|
|
|
def perspective_projection(points, rotation, translation,
|
|
focal_length, camera_center):
|
|
"""
|
|
Adapted from https://github.com/mkocabas/VIBE/blob/master/lib/models/spin.py
|
|
This function computes the perspective projection of a set of points.
|
|
Input:
|
|
points (bs, N, 3): 3D points
|
|
rotation (bs, 3, 3): Camera rotation
|
|
translation (bs, 3): Camera translation
|
|
focal_length (bs, 2): Focal length
|
|
camera_center (bs, 2): Camera center
|
|
"""
|
|
batch_size = points.shape[0]
|
|
K = torch.zeros([batch_size, 3, 3], device=points.device)
|
|
K[:,0,0] = focal_length[:,0]
|
|
K[:,1,1] = focal_length[:,1]
|
|
K[:,2,2] = 1.
|
|
K[:,:-1, -1] = camera_center
|
|
|
|
# Transform points
|
|
points = torch.einsum('bij,bkj->bki', rotation, points)
|
|
points = points + translation.unsqueeze(1)
|
|
|
|
# Apply perspective distortion
|
|
projected_points = points / points[:,:,-1].unsqueeze(-1)
|
|
|
|
# Apply camera intrinsics
|
|
projected_points = torch.einsum('bij,bkj->bki', K, projected_points)
|
|
|
|
return projected_points[:, :, :-1]
|
|
|
|
def run_metro(video_p, out_path, joints2d_image_path=None, jts_vid=None, gt_bbox=None):
|
|
|
|
import open3d_viz_overlay
|
|
|
|
pickle_path = os. path.join(out_path, "output.pkl")
|
|
cur_path = os.getcwd()
|
|
|
|
if not os.path.exists(pickle_path):
|
|
|
|
os.makedirs(out_path, exist_ok=True)
|
|
|
|
rel_vid_p = os.path.join("../..", video_p)
|
|
rel_out_p = os.path.join("../..", out_path)
|
|
|
|
command_list = [
|
|
"/home/eduran2/miniconda3/envs/nemf/bin/python",
|
|
"./metro/tools/end2end_inference_handmesh.py",
|
|
"--resume_checkpoint",
|
|
"./models/metro_hand_state_dict.bin",
|
|
"--image_file_or_path",
|
|
rel_vid_p
|
|
]
|
|
|
|
command_str = " ".join(command_list)
|
|
os.chdir("./external/MeshTransformer")
|
|
os.system(command_str)
|
|
|
|
os.chdir(cur_path)
|
|
# visualize with our open3d code
|
|
# cropped image dir
|
|
# open3d_viz_overlay.vis_opt_results(pred_file_path=os.path.join(out_path, "output.pkl"),
|
|
# gt_file_path=None,
|
|
# img_dir=os.path.join(os.path.dirname(video_p), "metro_out/rgb_crop")
|
|
|
|
|
|
open3d_viz_overlay.vis_opt_results(pred_file_path=os.path.join(out_path, "output.pkl"),
|
|
gt_file_path=None,
|
|
img_dir=video_p)
|
|
|
|
# read the output pickle file, and get the hand pose and shape parameters if not exist
|
|
if joblib.load(pickle_path).get("orient") is None:
|
|
|
|
import external.v2a_code.scripts.fit_mano_to_metro
|
|
external.v2a_code.scripts.fit_mano_to_metro.main(video_p)
|
|
|
|
out_obj = joblib.load(pickle_path)
|
|
out_obj["video_path"] = video_p
|
|
|
|
# full image dir
|
|
# yp = os.path.join(os.path.dirname(video_p), "metro_out/open3d_img/*.jpg")
|
|
# make video out of the images in the directory video_p
|
|
# os.system(f"/usr/bin/ffmpeg -framerate 30 -y -pattern_type glob -i {yp} {os.path.join(out_path, 'video.mp4')}")
|
|
|
|
# create keypoints from it.
|
|
|
|
return out_obj
|
|
|
|
|
|
def run_pymafx(video_p, pymaf_out_path, joints2d_image_path=None, jts_vid=None, gt_bbox=None):
|
|
|
|
pickle_path = os. path.join(pymaf_out_path, "output.pkl")
|
|
|
|
# run PyMAF-X if not already done
|
|
if not os.path.exists(pickle_path):
|
|
os.makedirs(pymaf_out_path, exist_ok=True)
|
|
|
|
# first save the gt bbox so that pymafx read it
|
|
if gt_bbox is not None:
|
|
np.savez(os.path.join(pymaf_out_path, "gt_bbox.npz"), **gt_bbox)
|
|
|
|
cur_path = os.getcwd()
|
|
os.chdir("./external/PyMAF-X")
|
|
|
|
rel_vid_p = os.path.join("../..", video_p)
|
|
rel_out_p = os.path.join("../..", pymaf_out_path)
|
|
|
|
|
|
command_list = ["python", "-m", "apps.demo_mano", "--image_folder", rel_vid_p, "--detection_threshold", "0.3",
|
|
"--misc", "MODEL.PyMAF.OPT_HEAD", "False", "TRAIN.BHF_MODE", "hand_only", "MODEL.MESH_MODEL", "mano",
|
|
"--output_folder", rel_out_p, "--pretrained_model", "./data/pretrained_model/PyMAF-X_model_checkpoint_v1.1.pt",
|
|
"--out_vid_path", jts_vid]
|
|
|
|
if joints2d_image_path is not None:
|
|
joints2d_image_path = os.path.join("../..", joints2d_image_path)
|
|
command_list.extend(["--joints2d_image_folder", joints2d_image_path])
|
|
|
|
# uncomment this line if you want to have the direct open3d rendering of the PyMAF-X output
|
|
# commant_list.extend([", --render"])
|
|
|
|
command_str = " ".join(command_list)
|
|
|
|
os.system(command_str)
|
|
os.chdir(cur_path)
|
|
|
|
out_obj = joblib.load(pickle_path)
|
|
out_obj["video_path"] = video_p
|
|
out_obj["cam_f"] = np.array([5000, 5000])
|
|
return out_obj
|
|
|
|
def process_pymafx_mano(pickle_obj):
|
|
|
|
if not pickle_obj["smplx_params"]:
|
|
return None, None, None, None, None
|
|
|
|
|
|
batch_count = len(pickle_obj["smplx_params"])
|
|
frame_count = len(pickle_obj["frame_ids"])
|
|
|
|
right_hand_pose, rh_orient, rh_betas = [], [], []
|
|
|
|
for i in range(batch_count):
|
|
bs = pickle_obj["smplx_params"][i]["body_pose"].shape[0]
|
|
|
|
right_hand_pose.append(matrix_to_axis_angle(pickle_obj["smplx_params"][i]["body_pose"].reshape(-1, 3, 3)).reshape(bs, -1, 3))
|
|
rh_betas.append(pickle_obj["smplx_params"][i]["shape"])
|
|
rh_orient.append(matrix_to_axis_angle(pickle_obj["smplx_params"][i]["root_orient"].reshape(-1, 3, 3)).reshape(bs, 3))
|
|
|
|
right_hand_pose = torch.vstack(right_hand_pose)
|
|
rh_betas = torch.vstack(rh_betas)
|
|
rh_orient = torch.vstack(rh_orient)
|
|
|
|
transl = torch.tensor(pickle_obj["orig_cam_t"], dtype=torch.float32)
|
|
|
|
export_pymafx_json(fpath=pickle_obj["video_path"], jts_2d=pickle_obj["pymafx_joints2d"])
|
|
rh_verts = pickle_obj["verts"]
|
|
|
|
return rh_orient, transl, right_hand_pose, rh_betas, rh_verts
|
|
|
|
|
|
def export_pymafx_json(fpath, jts_2d):
|
|
|
|
if ("HO3D_v3" in fpath) or ("DexYCB" in fpath) or ("in_the_wild" in fpath) or ("arctic_data"):
|
|
out_keypoints_path = os.path.join(os.path.dirname(fpath), "pymafx_keypoints2d")
|
|
else:
|
|
raise NotImplementedError("PYMAFX export not implemented for this dataset")
|
|
|
|
# slerp(time smoothing) for 3d case is done in fitting_app
|
|
os.makedirs(out_keypoints_path, exist_ok=True)
|
|
|
|
right_hand_tmp = jts_2d # confidences come from mmpose detection confidence values
|
|
rgb_image_list = sorted([os.path.join(fpath, x) for x in os.listdir(fpath) if x.endswith('.png') or x.endswith('.jpg')])
|
|
|
|
for im_k, im_path_k in enumerate(rgb_image_list):
|
|
json_pathname = os.path.basename(im_path_k).split(".")[0] + "_keypoints.json"
|
|
json_pathname = os.path.join(out_keypoints_path, json_pathname)
|
|
|
|
|
|
dic = {}
|
|
dic['version'] = '1.5'
|
|
dic["people"] = []
|
|
person_dic = {}
|
|
|
|
person_dic["person_id"] = [-1]
|
|
person_dic["pose_keypoints_2d"] = []
|
|
person_dic["face_keypoints_2d"] = []
|
|
person_dic["hand_left_keypoints_2d"] = []
|
|
person_dic["hand_right_keypoints_2d"] = list(right_hand_tmp[im_k].reshape(-1))
|
|
|
|
dic["people"].append(person_dic)
|
|
|
|
with open(json_pathname, 'w') as fp:
|
|
json.dump(dic, fp)
|
|
|
|
return
|
|
|
|
def blend_keypoints(source_path1, source_path2, target_path, gt_frames=None, raw_img_path=None,
|
|
blend_ratio=0.5, blend_vid_out_path=None, smooth_flag=True,
|
|
render=False, circle_rad=4, line_width=2):
|
|
""" source_path1: path to mediapipe keypoints1
|
|
source_path2: path to pymafx keypoints2
|
|
"""
|
|
raise NotImplementedError
|
|
# green for pymafx, red for mediapipe
|
|
COLORS = {"pymafx": (0, 255, 0), "together": (0, 0, 255)}
|
|
|
|
from datasets.data_utils import read_keypoints
|
|
|
|
smoother = None
|
|
|
|
if blend_vid_out_path is not None and render:
|
|
os.makedirs(blend_vid_out_path, exist_ok=True)
|
|
|
|
if raw_img_path is not None:
|
|
raw_img_list = sorted([os.path.join(raw_img_path, x) for x in os.listdir(raw_img_path) if x.endswith('.png') or x.endswith('.jpg')])
|
|
|
|
os.makedirs(target_path, exist_ok=True)
|
|
|
|
assert os.path.exists(source_path1), "Source path 1 does not exist"
|
|
assert os.path.exists(source_path2), "Source path 1 does not exist"
|
|
|
|
kypts1 = sorted(glob.glob(os.path.join(source_path1, "*.json")))
|
|
kypts2 = sorted(glob.glob(os.path.join(source_path2, "*.json")))
|
|
|
|
assert len(kypts1) == len(kypts2), "Number of keypoints do not match"
|
|
|
|
|
|
kypts1 = [read_keypoints(f) for f in kypts1]
|
|
kypts2 = [read_keypoints(f_) for f_ in kypts2]
|
|
|
|
blend_res_list, smooth_list, valid_ind, keyp_source_list = [], [], [], []
|
|
track_id = 0
|
|
|
|
for i, (kypt1, kypt2) in enumerate(tqdm((zip(kypts1, kypts2)))):
|
|
|
|
# check if the mediapipe confidence are all zero ( 0,0,0 sometimes is violated (0, 480, 0) input comes )
|
|
if np.array_equiv(kypt1[:, -1], np.zeros_like(kypt1[:, -1])):
|
|
blend_rat = 0.0
|
|
keyp_source_list.append("pymafx")
|
|
else:
|
|
blend_rat = blend_ratio
|
|
keyp_source_list.append("together")
|
|
|
|
blend_res_i = np.array(kypt1) * blend_rat + np.array(kypt2) * (1 - blend_rat)
|
|
|
|
# to see if the gt bbox is discontinued
|
|
if not np.array_equiv(kypt2, np.zeros_like(kypt2)):
|
|
# print(kypt2.mean(), kypt1.mean(), blend_res_i.mean())
|
|
|
|
if smoother:
|
|
blend_res_i[:, :2] = smoother.smooth([{"keypoints": blend_res_i[:, :2],
|
|
"track_id": track_id}])[0]["keypoints"]
|
|
valid_ind.append(i)
|
|
else:
|
|
track_id += 1
|
|
|
|
blend_res_list.append(blend_res_i)
|
|
target_json_pathname = os.path.join(target_path, f"{i:04d}_keypoints.json")
|
|
|
|
if sum(blend_res_i[:, -1]) == 0.0:
|
|
assert i not in gt_frames, "GT frame is not valid"
|
|
|
|
dic = {}
|
|
dic['version'] = '1.5'
|
|
dic["people"] = []
|
|
person_dic = {}
|
|
|
|
# openpose format
|
|
person_dic["person_id"] = [-1]
|
|
person_dic["pose_keypoints_2d"] = []
|
|
person_dic["face_keypoints_2d"] = []
|
|
person_dic["hand_left_keypoints_2d"] = []
|
|
person_dic["hand_right_keypoints_2d"] = blend_res_i.reshape(-1, 1).tolist()
|
|
|
|
# if blend_res_i[:, -1].mean() == 0.0:
|
|
# import ipdb; ipdb.set_trace()
|
|
|
|
dic["people"].append(person_dic)
|
|
|
|
with open(target_json_pathname, 'w') as fp:
|
|
json.dump(dic, fp)
|
|
|
|
if render:
|
|
img_rgb = raw_img_list[i]
|
|
color = COLORS[keyp_source_list[i]]
|
|
|
|
_img = Image.open(img_rgb)
|
|
draw = ImageDraw.Draw(_img)
|
|
|
|
# check if there is a corresponding gt for this frame
|
|
joints2d = blend_res_i[:, :2]
|
|
|
|
for k in range(21):
|
|
kps_parent = joints2d[openpose_skeleton[k]]
|
|
kps_child = joints2d[k]
|
|
|
|
|
|
if (kps_parent == [0.0, 0.0]).all() or (kps_child == [0.0, 0.0]).all():
|
|
continue
|
|
|
|
if openpose_skeleton[k] != -1:
|
|
draw.line([(kps_child[0], kps_child[1]), (kps_parent[0], kps_parent[1])], fill=color, width=line_width)
|
|
|
|
if k in [3, 4]:
|
|
jts_color = (255, 0, 0)
|
|
elif k in [7, 8]:
|
|
jts_color = (139,69,19)
|
|
elif k in [19, 20]:
|
|
jts_color = (0, 0, 0)
|
|
else:
|
|
jts_color = (100, 100, 100)
|
|
|
|
|
|
draw.ellipse((joints2d[k][0]-circle_rad, joints2d[k][1]-circle_rad, joints2d[k][0]+circle_rad, joints2d[k][1]+circle_rad), fill=jts_color)
|
|
|
|
os.makedirs(blend_vid_out_path, exist_ok=True)
|
|
_img.save(blend_vid_out_path + f"/{i:04d}.jpg")
|
|
|
|
|
|
if not blend_vid_out_path is None and render:
|
|
os.system(f"/usr/bin/ffmpeg -y -framerate 30 -i {blend_vid_out_path}/%04d.jpg -vcodec libx264 -pix_fmt yuv420p {blend_vid_out_path}.mp4")
|
|
|
|
return
|
|
|
|
|
|
def save_quantitative_evaluation(evaluator_object, pred_dict, pymafx_dict, gt_dict, metrics, viz_flag=True, misc={}):
|
|
|
|
# form camera intrinsics
|
|
img_H, img_W = pred_dict["img_height"], pred_dict["img_width"]
|
|
pred_dict["source"] = "hmp"
|
|
|
|
hmp_quant_metric_vals = evaluator_object.compute_metric(model_output=pred_dict,
|
|
targets=gt_dict,
|
|
metrics=metrics)
|
|
|
|
if viz_flag:
|
|
for alignment_type in ['root', 'procrustes']:
|
|
html_filepath = os.path.join(os.path.dirname(pred_dict["save_path"].item()), f"{alignment_type}_aligned_hmp.html")
|
|
vid_filepath = os.path.join(os.path.dirname(pred_dict["save_path"].item()), f"{alignment_type}_aligned_hmp.mp4")
|
|
frame_dir = os.path.join(os.path.dirname(pred_dict["save_path"].item()), f"{alignment_type}_aligned_hmp_open3d_img")
|
|
|
|
evaluator_object.create_canvas_open3d(img_width=img_H, img_height=img_W, cam_intrinsics=gt_dict["cam_intrinsics"],
|
|
img_dir=pred_dict["img_dir"].item(), video_path=vid_filepath, alignment_type=alignment_type,
|
|
valid_indices=gt_dict["frame_id"], frame_dir=frame_dir)
|
|
|
|
|
|
# optim/CONFIGNAME/performance_sequence.pkl
|
|
cfg_name = pred_dict["config_type"].item()
|
|
pymafx_quant_metric_vals = None
|
|
|
|
if cfg_name in ["_slerp_encode_decode", "_encode_decode"]:
|
|
pymafx_dict["source"] = "pymafx"
|
|
pymafx_quant_metric_vals = evaluator_object.compute_metric(model_output=pymafx_dict,
|
|
targets=gt_dict,
|
|
metrics=metrics)
|
|
|
|
# if viz_flag:
|
|
# for alignment_type in ['root', 'procrustes']:
|
|
# html_filepath = os.path.join(os.path.dirname(pred_dict["save_path"].item()), f"{alignment_type}_aligned_pymafx.html")
|
|
# vid_filepath = os.path.join(os.path.dirname(pred_dict["save_path"].item()), f"{alignment_type}_aligned_pymafx.mp4")
|
|
# frame_dir = os.path.join(os.path.dirname(pred_dict["save_path"].item()), f"{alignment_type}_aligned_pymafx_open3d_img")
|
|
|
|
# # evaluator_object.create_canvas_scenepic(img_width=img_H , img_height=img_W, cam_intrinsics=gt_dict["cam_intrinsics"],
|
|
# # img_dir=pred_dict["img_dir"], save_dir=html_filepath, alignment_type=alignment_type,
|
|
# # valid_indices=gt_dict["frame_id"])
|
|
# evaluator_object.create_canvas_open3d(img_width=img_H, img_height=img_W, cam_intrinsics=gt_dict["cam_intrinsics"],
|
|
# img_dir=pred_dict["img_dir"].item(), video_path=vid_filepath, alignment_type=alignment_type,
|
|
# valid_indices=gt_dict["frame_id"], frame_dir=frame_dir)
|
|
|
|
dname = gt_dict["dataset_name"]
|
|
exp_setup_name = gt_dict["exp_setup_name"]
|
|
|
|
|
|
# performance_sequence = f"./optim/{cfg_name}/{dname}/performance_sequence_{exp_setup_name}.pkl"
|
|
# performance_per_frame = f"./optim/{cfg_name}/{dname}/performance_per_frame_{exp_setup_name}.pkl"
|
|
|
|
name = gt_dict["path"].replace("/", "_").replace(".", "")
|
|
|
|
os.makedirs(f"./optim/{cfg_name}/{dname}/performance_sequence", exist_ok=True)
|
|
os.makedirs(f"./optim/{cfg_name}/{dname}/performance_per_frame", exist_ok=True)
|
|
|
|
performance_sequence = f"./optim/{cfg_name}/{dname}/performance_sequence/{name}.pkl"
|
|
performance_per_frame = f"./optim/{cfg_name}/{dname}/performance_per_frame/{name}.pkl"
|
|
|
|
|
|
if not os.path.isfile(performance_sequence):
|
|
print("Creating new performance file")
|
|
joblib.dump({cfg_name: {}, "pymafx": {}}, performance_sequence)
|
|
|
|
if not os.path.isfile(performance_per_frame):
|
|
print("Creating new performance file")
|
|
joblib.dump({cfg_name: {}, "pymafx": {}}, performance_per_frame)
|
|
|
|
print("===================HMP=====================" )
|
|
for metric_name, metric_val in hmp_quant_metric_vals.items():
|
|
if metric_name != "valid_frames":
|
|
# take mean over joints.
|
|
hmp_quant_metric_vals[metric_name] = metric_val.mean(axis=1)
|
|
print(metric_name, "(mm): ", f"{metric_val.mean():.3f}")
|
|
|
|
if not pymafx_quant_metric_vals is None:
|
|
print("==================PYMAFX===================" )
|
|
for metric_name, pymafx_metric_val in pymafx_quant_metric_vals.items():
|
|
if metric_name != "valid_frames":
|
|
pymafx_quant_metric_vals[metric_name] = pymafx_metric_val.mean(axis=1)
|
|
print(metric_name, "(mm): ", f"{pymafx_metric_val.mean():.3f}")
|
|
|
|
print("===========================================" )
|
|
|
|
with open(performance_per_frame, "rb") as f:
|
|
perframe_pkl_file = joblib.load(f)
|
|
|
|
if cfg_name not in perframe_pkl_file.keys():
|
|
perframe_pkl_file[cfg_name] = {}
|
|
|
|
# save per frame values
|
|
perframe_pkl_file[cfg_name][gt_dict["path"]] = hmp_quant_metric_vals
|
|
perframe_pkl_file["pymafx"][gt_dict["path"]] = pymafx_quant_metric_vals
|
|
|
|
joblib.dump(perframe_pkl_file, performance_per_frame)
|
|
|
|
with open(performance_sequence, "rb") as f:
|
|
seq_pkl_file = joblib.load(f)
|
|
|
|
if cfg_name not in seq_pkl_file.keys():
|
|
seq_pkl_file[cfg_name] = {}
|
|
|
|
# take mean over time.
|
|
seq_pkl_file[cfg_name][gt_dict["path"]] = {k:v.mean(axis=0) if k!= "valid_frames" else v for k, v in hmp_quant_metric_vals.items()}
|
|
if pymafx_quant_metric_vals is not None:
|
|
seq_pkl_file["pymafx"][gt_dict["path"]] = {k:v.mean(axis=0) if k!= "valid_frames" else v for k, v in pymafx_quant_metric_vals.items()}
|
|
joblib.dump(seq_pkl_file, performance_sequence)
|
|
|
|
return
|
|
|
|
def vis_mano_skeleton_matplotlib(gt_dict, pred_dict, idx=None):
|
|
|
|
if idx is None:
|
|
idx = np.random.randint(0, len(gt_dict["joints_3d"]))
|
|
|
|
fig = plt.figure()
|
|
ax = fig.add_subplot(projection='3d')
|
|
ax.scatter(pred_dict["joints_3d"][idx, :, 0], pred_dict["joints_3d"][idx, :, 1], pred_dict["joints_3d"][idx, :, 2], color='red')
|
|
ax.scatter(gt_dict["joints_3d"][idx, :, 0], gt_dict["joints_3d"][idx, :, 1], gt_dict["joints_3d"][idx, :, 2], color='green')
|
|
|
|
for i in range(21):
|
|
if AUGMENTED_MANO_CHAIN[i] == -1:
|
|
continue
|
|
|
|
parent = pred_dict["joints_3d"][idx][AUGMENTED_MANO_CHAIN[i]]
|
|
child = pred_dict["joints_3d"][idx][i]
|
|
ax.plot([child[0], parent[0]], [child[1], parent[1]], zs=[child[2], parent[2]], color='red')
|
|
|
|
parent_i = gt_dict["joints_3d"][idx][AUGMENTED_MANO_CHAIN[i]]
|
|
child_i = gt_dict["joints_3d"][idx][i]
|
|
|
|
ax.plot([child_i[0], parent_i[0]], [child_i[1], parent_i[1]], [child_i[2], parent_i[2]], color='green')
|
|
|
|
|
|
ax.set_xlabel('X Label')
|
|
ax.set_ylabel('Y Label')
|
|
ax.set_zlabel('Z Label')
|
|
plt.savefig("test.png")
|
|
|
|
return
|
|
|
|
|
|
|
|
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 = f"{os.path.dirname(__file__)}/../../_DATA/BMC"
|
|
|
|
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)]
|
|
|
|
self.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,
|
|
]
|
|
|
|
self.JOINT_ROOT_IDX = 9
|
|
|
|
self.REF_BONE_LINK = (0, 9) # mid mcp
|
|
|
|
# bone indexes in 20 bones setting
|
|
self.ID_ROOT_bone = np.array([0, 4, 8, 12, 16]) # ROOT_bone from wrist to MCP
|
|
self.ID_PIP_bone = np.array([1, 5, 9, 13, 17]) # PIP_bone from MCP to PIP
|
|
self.ID_DIP_bone = np.array([2, 6, 10, 14, 18]) # DIP_bone from PIP to DIP
|
|
self.ID_TIP_bone = np.array([3, 7, 11, 15, 19]) # TIP_bone from DIP to TIP
|
|
|
|
def compute_loss(self, joints):
|
|
'''
|
|
|
|
Args:
|
|
joints: B*21*3
|
|
|
|
Returns:
|
|
|
|
'''
|
|
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[:, self.SNAP_PARENT[i], :]
|
|
) for i in range(21)
|
|
]
|
|
ALL_bones = torch.stack(ALL_bones[1:], dim=1) # (B,20,3)
|
|
ROOT_bones = ALL_bones[:, self.ID_ROOT_bone] # (B,5,3)
|
|
PIP_bones = ALL_bones[:, self.ID_PIP_bone]
|
|
DIP_bones = ALL_bones[:, self.ID_DIP_bone]
|
|
TIP_bones = ALL_bones[:, self.ID_TIP_bone]
|
|
|
|
ALL_Z_axis = normalize(ALL_bones)
|
|
PIP_Z_axis = ALL_Z_axis[:, self.ID_ROOT_bone]
|
|
DIP_Z_axis = ALL_Z_axis[:, self.ID_PIP_bone]
|
|
TIP_Z_axis = ALL_Z_axis[:, self.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 |