1457 lines
55 KiB
Python
1457 lines
55 KiB
Python
import os
|
||
import math
|
||
import numpy as np
|
||
import cv2
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torch.distributions import (
|
||
MixtureSameFamily,
|
||
Categorical,
|
||
Normal,
|
||
MultivariateNormal,
|
||
)
|
||
|
||
# from body_model import SMPL_JOINTS, SMPL_PARENTS
|
||
from geometry.rotation import rotation_matrix_to_angle_axis
|
||
from geometry import camera as cam_util
|
||
from optim.bio_loss import BMCLoss
|
||
from util.logger import Logger
|
||
from typing import List, Tuple, NewType
|
||
|
||
Tensor = NewType('Tensor', torch.Tensor)
|
||
|
||
CONTACT_HEIGHT_THRESH = 0.08
|
||
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
|
||
|
||
def solid_angles(
|
||
points: Tensor,
|
||
triangles: Tensor,
|
||
thresh: float = 1e-8
|
||
) -> Tensor:
|
||
''' Compute solid angle between the input points and triangles
|
||
Follows the method described in:
|
||
The Solid Angle of a Plane Triangle
|
||
A. VAN OOSTEROM AND J. STRACKEE
|
||
IEEE TRANSACTIONS ON BIOMEDICAL ENGINEERING,
|
||
VOL. BME-30, NO. 2, FEBRUARY 1983
|
||
Parameters
|
||
-----------
|
||
points: BxQx3
|
||
Tensor of input query points
|
||
triangles: BxFx3x3
|
||
Target triangles
|
||
thresh: float
|
||
float threshold
|
||
Returns
|
||
-------
|
||
solid_angles: BxQxF
|
||
A tensor containing the solid angle between all query points
|
||
and input triangles
|
||
'''
|
||
# Center the triangles on the query points. Size should be BxQxFx3x3
|
||
centered_tris = triangles[:, None] - points[:, :, None, None]
|
||
|
||
# BxQxFx3
|
||
norms = torch.norm(centered_tris, dim=-1)
|
||
|
||
# Should be BxQxFx3
|
||
cross_prod = torch.cross(
|
||
centered_tris[:, :, :, 1], centered_tris[:, :, :, 2], dim=-1)
|
||
# Should be BxQxF
|
||
numerator = (centered_tris[:, :, :, 0] * cross_prod).sum(dim=-1)
|
||
del cross_prod
|
||
|
||
dot01 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 1]).sum(dim=-1)
|
||
dot12 = (centered_tris[:, :, :, 1] * centered_tris[:, :, :, 2]).sum(dim=-1)
|
||
dot02 = (centered_tris[:, :, :, 0] * centered_tris[:, :, :, 2]).sum(dim=-1)
|
||
del centered_tris
|
||
|
||
denominator = (
|
||
norms.prod(dim=-1) +
|
||
dot01 * norms[:, :, :, 2] +
|
||
dot02 * norms[:, :, :, 1] +
|
||
dot12 * norms[:, :, :, 0]
|
||
)
|
||
del dot01, dot12, dot02, norms
|
||
|
||
# Should be BxQ
|
||
solid_angle = torch.atan2(numerator, denominator)
|
||
del numerator, denominator
|
||
|
||
torch.cuda.empty_cache()
|
||
|
||
return 2 * solid_angle
|
||
|
||
def winding_numbers(
|
||
points: Tensor,
|
||
triangles: Tensor,
|
||
thresh: float = 1e-8
|
||
) -> Tensor:
|
||
''' Uses winding_numbers to compute inside/outside
|
||
Robust inside-outside segmentation using generalized winding numbers
|
||
Alec Jacobson,
|
||
Ladislav Kavan,
|
||
Olga Sorkine-Hornung
|
||
Fast Winding Numbers for Soups and Clouds SIGGRAPH 2018
|
||
Gavin Barill
|
||
NEIL G. Dickson
|
||
Ryan Schmidt
|
||
David I.W. Levin
|
||
and Alec Jacobson
|
||
Parameters
|
||
-----------
|
||
points: BxQx3
|
||
Tensor of input query points
|
||
triangles: BxFx3x3
|
||
Target triangles
|
||
thresh: float
|
||
float threshold
|
||
Returns
|
||
-------
|
||
winding_numbers: BxQ
|
||
A tensor containing the Generalized winding numbers
|
||
'''
|
||
# The generalized winding number is the sum of solid angles of the point
|
||
# with respect to all triangles.
|
||
return 1 / (4 * math.pi) * solid_angles(
|
||
points, triangles, thresh=thresh).sum(dim=-1)
|
||
|
||
def pcl_pcl_pairwise_distance(
|
||
x: torch.Tensor,
|
||
y: torch.Tensor,
|
||
use_cuda: bool = True,
|
||
squared: bool = False
|
||
):
|
||
"""
|
||
Calculate the pairse distance between two point clouds.
|
||
"""
|
||
|
||
bs, num_points_x, points_dim = x.size()
|
||
_, num_points_y, _ = y.size()
|
||
|
||
dtype = torch.cuda.LongTensor if \
|
||
use_cuda else torch.LongTensor
|
||
|
||
xx = torch.bmm(x, x.transpose(2, 1))
|
||
yy = torch.bmm(y, y.transpose(2, 1))
|
||
zz = torch.bmm(x, y.transpose(2, 1))
|
||
|
||
diag_ind_x = torch.arange(0, num_points_x).type(dtype)
|
||
diag_ind_y = torch.arange(0, num_points_y).type(dtype)
|
||
rx = (
|
||
xx[:, diag_ind_x, diag_ind_x]
|
||
.unsqueeze(1)
|
||
.expand_as(zz.transpose(2, 1))
|
||
)
|
||
ry = yy[:, diag_ind_y, diag_ind_y].unsqueeze(1).expand_as(zz)
|
||
P = rx.transpose(2, 1) + ry - 2 * zz
|
||
|
||
if not squared:
|
||
P = torch.clamp(P, min=0.0) # make sure we dont get nans
|
||
P = torch.sqrt(P)
|
||
|
||
return P
|
||
|
||
def get_keypoints_rectangle(keypoints: np.array, threshold: float) -> Tuple[float, float, float]:
|
||
"""
|
||
Compute rectangle enclosing keypoints above the threshold.
|
||
Args:
|
||
keypoints (np.array): Keypoint array of shape (N, 3).
|
||
threshold (float): Confidence visualization threshold.
|
||
Returns:
|
||
Tuple[float, float, float]: Rectangle width, height and area.
|
||
"""
|
||
# print(keypoints.shape)
|
||
valid_ind = keypoints[:, -1] > threshold
|
||
if valid_ind.sum() > 0:
|
||
valid_keypoints = keypoints[valid_ind][:, :-1]
|
||
max_x = valid_keypoints[:,0].max()
|
||
max_y = valid_keypoints[:,1].max()
|
||
min_x = valid_keypoints[:,0].min()
|
||
min_y = valid_keypoints[:,1].min()
|
||
width = max_x - min_x
|
||
height = max_y - min_y
|
||
area = width * height
|
||
return width, height, area
|
||
else:
|
||
return 0,0,0
|
||
|
||
def render_keypoints(img: np.array,
|
||
keypoints: np.array,
|
||
pairs: List,
|
||
colors: List,
|
||
thickness_circle_ratio: float,
|
||
thickness_line_ratio_wrt_circle: float,
|
||
pose_scales: List,
|
||
threshold: float = 0.1,
|
||
alpha: float = 1.0) -> np.array:
|
||
"""
|
||
Render keypoints on input image.
|
||
Args:
|
||
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
||
keypoints (np.array): Keypoint array of shape (N, 3).
|
||
pairs (List): List of keypoint pairs per limb.
|
||
colors: (List): List of colors per keypoint.
|
||
thickness_circle_ratio (float): Circle thickness ratio.
|
||
thickness_line_ratio_wrt_circle (float): Line thickness ratio wrt the circle.
|
||
pose_scales (List): List of pose scales.
|
||
threshold (float): Only visualize keypoints with confidence above the threshold.
|
||
Returns:
|
||
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
||
"""
|
||
img_orig = img.copy()
|
||
width, height = img.shape[1], img.shape[2]
|
||
area = width * height
|
||
|
||
lineType = 8
|
||
shift = 0
|
||
numberColors = len(colors)
|
||
thresholdRectangle = 0.1
|
||
|
||
person_width, person_height, person_area = get_keypoints_rectangle(keypoints, thresholdRectangle)
|
||
if person_area > 0:
|
||
ratioAreas = min(1, max(person_width / width, person_height / height))
|
||
thicknessRatio = np.maximum(np.round(math.sqrt(area) * thickness_circle_ratio * ratioAreas), 2)
|
||
thicknessCircle = np.maximum(1, thicknessRatio if ratioAreas > 0.05 else -np.ones_like(thicknessRatio))
|
||
thicknessLine = np.maximum(1, np.round(thicknessRatio * thickness_line_ratio_wrt_circle))
|
||
radius = thicknessRatio / 2
|
||
|
||
img = np.ascontiguousarray(img.copy())
|
||
for i, pair in enumerate(pairs):
|
||
index1, index2 = pair
|
||
if keypoints[index1, -1] > threshold and keypoints[index2, -1] > threshold:
|
||
thicknessLineScaled = int(round(min(thicknessLine[index1], thicknessLine[index2]) * pose_scales[0]))
|
||
colorIndex = index2
|
||
color = colors[colorIndex % numberColors]
|
||
keypoint1 = keypoints[index1, :-1].astype(np.int32)
|
||
keypoint2 = keypoints[index2, :-1].astype(np.int32)
|
||
cv2.line(img, tuple(keypoint1.tolist()), tuple(keypoint2.tolist()), tuple(color.tolist()), thicknessLineScaled, lineType, shift)
|
||
for part in range(len(keypoints)):
|
||
faceIndex = part
|
||
if keypoints[faceIndex, -1] > threshold:
|
||
radiusScaled = int(round(radius[faceIndex] * pose_scales[0]))
|
||
thicknessCircleScaled = int(round(thicknessCircle[faceIndex] * pose_scales[0]))
|
||
colorIndex = part
|
||
color = colors[colorIndex % numberColors]
|
||
center = keypoints[faceIndex, :-1].astype(np.int32)
|
||
cv2.circle(img, tuple(center.tolist()), radiusScaled, tuple(color.tolist()), thicknessCircleScaled, lineType, shift)
|
||
return img
|
||
|
||
def render_hand_keypoints(img, right_hand_keypoints, threshold=0.1, use_confidence=False, map_fn=lambda x: np.ones_like(x), alpha=1.0):
|
||
if use_confidence and map_fn is not None:
|
||
#thicknessCircleRatioLeft = 1./50 * map_fn(left_hand_keypoints[:, -1])
|
||
thicknessCircleRatioRight = 1./50 * map_fn(right_hand_keypoints[:, -1])
|
||
else:
|
||
#thicknessCircleRatioLeft = 1./50 * np.ones(left_hand_keypoints.shape[0])
|
||
thicknessCircleRatioRight = 1./50 * np.ones(right_hand_keypoints.shape[0])
|
||
thicknessLineRatioWRTCircle = 0.75
|
||
pairs = [0,1, 1,2, 2,3, 3,4, 0,5, 5,6, 6,7, 7,8, 0,9, 9,10, 10,11, 11,12, 0,13, 13,14, 14,15, 15,16, 0,17, 17,18, 18,19, 19,20]
|
||
pairs = np.array(pairs).reshape(-1,2)
|
||
|
||
colors = [100., 100., 100.,
|
||
100., 0., 0.,
|
||
150., 0., 0.,
|
||
200., 0., 0.,
|
||
255., 0., 0.,
|
||
100., 100., 0.,
|
||
150., 150., 0.,
|
||
200., 200., 0.,
|
||
255., 255., 0.,
|
||
0., 100., 50.,
|
||
0., 150., 75.,
|
||
0., 200., 100.,
|
||
0., 255., 125.,
|
||
0., 50., 100.,
|
||
0., 75., 150.,
|
||
0., 100., 200.,
|
||
0., 125., 255.,
|
||
100., 0., 100.,
|
||
150., 0., 150.,
|
||
200., 0., 200.,
|
||
255., 0., 255.]
|
||
colors = np.array(colors).reshape(-1,3)
|
||
#colors = np.zeros_like(colors)
|
||
poseScales = [1]
|
||
#img = render_keypoints(img, left_hand_keypoints, pairs, colors, thicknessCircleRatioLeft, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha)
|
||
img = render_keypoints(img, right_hand_keypoints, pairs, colors, thicknessCircleRatioRight, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha)
|
||
#img = render_keypoints(img, right_hand_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1)
|
||
return img
|
||
|
||
def render_body_keypoints(img: np.array,
|
||
body_keypoints: np.array) -> np.array:
|
||
"""
|
||
Render OpenPose body keypoints on input image.
|
||
Args:
|
||
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
||
body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence).
|
||
Returns:
|
||
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
||
"""
|
||
|
||
thickness_circle_ratio = 1./75. * np.ones(body_keypoints.shape[0])
|
||
thickness_line_ratio_wrt_circle = 0.75
|
||
pairs = []
|
||
pairs = [1,8,1,2,1,5,2,3,3,4,5,6,6,7,8,9,9,10,10,11,8,12,12,13,13,14,1,0,0,15,15,17,0,16,16,18,14,19,19,20,14,21,11,22,22,23,11,24]
|
||
pairs = np.array(pairs).reshape(-1,2)
|
||
colors = [255., 0., 85.,
|
||
255., 0., 0.,
|
||
255., 85., 0.,
|
||
255., 170., 0.,
|
||
255., 255., 0.,
|
||
170., 255., 0.,
|
||
85., 255., 0.,
|
||
0., 255., 0.,
|
||
255., 0., 0.,
|
||
0., 255., 85.,
|
||
0., 255., 170.,
|
||
0., 255., 255.,
|
||
0., 170., 255.,
|
||
0., 85., 255.,
|
||
0., 0., 255.,
|
||
255., 0., 170.,
|
||
170., 0., 255.,
|
||
255., 0., 255.,
|
||
85., 0., 255.,
|
||
0., 0., 255.,
|
||
0., 0., 255.,
|
||
0., 0., 255.,
|
||
0., 255., 255.,
|
||
0., 255., 255.,
|
||
0., 255., 255.]
|
||
colors = np.array(colors).reshape(-1,3)
|
||
pose_scales = [1]
|
||
return render_keypoints(img, body_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1)
|
||
|
||
def render_openpose(img: np.array,
|
||
hand_keypoints: np.array) -> np.array:
|
||
"""
|
||
Render keypoints in the OpenPose format on input image.
|
||
Args:
|
||
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
||
body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence).
|
||
Returns:
|
||
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
||
"""
|
||
#img = render_body_keypoints(img, body_keypoints)
|
||
img = render_hand_keypoints(img, hand_keypoints)
|
||
return img
|
||
###################################################
|
||
|
||
class StageLoss(nn.Module):
|
||
def __init__(self, loss_weights, **kwargs):
|
||
super().__init__()
|
||
self.cur_optim_step = 0
|
||
self.set_loss_weights(loss_weights)
|
||
print(kwargs)
|
||
self.setup_losses(loss_weights, **kwargs)
|
||
|
||
def setup_losses(self, *args, **kwargs):
|
||
raise NotImplementedError
|
||
|
||
def set_loss_weights(self, loss_weights):
|
||
self.loss_weights = loss_weights
|
||
Logger.log("Stage loss weights set to:")
|
||
Logger.log(self.loss_weights)
|
||
|
||
|
||
class RootLoss(StageLoss):
|
||
def setup_losses(
|
||
self,
|
||
loss_weights,
|
||
ignore_op_joints=None,
|
||
joints2d_sigma=100,
|
||
use_chamfer=False,
|
||
robust_loss="none",
|
||
robust_tuning_const=4.6851,
|
||
faces=None
|
||
):
|
||
self.joints2d_loss = Joints2DLoss(ignore_op_joints, joints2d_sigma)
|
||
self.points3d_loss = Points3DLoss(use_chamfer, robust_loss, robust_tuning_const)
|
||
self.inter_penetration_loss = GeneralContactLoss(faces)
|
||
self.bio_loss = BMCLoss(lambda_bl=1, lambda_rb=1, lambda_a=1)
|
||
|
||
def forward(self, observed_data, pred_data, valid_mask=None):
|
||
"""
|
||
For fitting just global root trans/orientation.
|
||
Only computes joint/point/vert losses, i.e. no priors.
|
||
"""
|
||
stats_dict = dict()
|
||
loss = 0.0
|
||
|
||
# Joints in 3D space
|
||
if (
|
||
"joints3d" in observed_data
|
||
and "joints3d" in pred_data
|
||
and self.loss_weights["joints3d"] > 0.0
|
||
):
|
||
cur_loss = joints3d_loss(
|
||
observed_data["joints3d"], pred_data["joints3d"], valid_mask
|
||
)
|
||
loss += self.loss_weights["joints3d"] * cur_loss
|
||
stats_dict["joints3d"] = cur_loss
|
||
|
||
if (
|
||
"joints3d" in pred_data
|
||
and self.loss_weights["bio"] > 0.0
|
||
):
|
||
# print('biobiobiobiobiobiobiobiobiobiobiobio')
|
||
cur_loss, _ = self.bio_loss.compute_loss(
|
||
pred_data["joints3d"], valid_mask
|
||
)
|
||
loss += self.loss_weights["bio"] * cur_loss[0]
|
||
stats_dict["bio"] = cur_loss[0]
|
||
|
||
# Select vertices in 3D space
|
||
if (
|
||
"verts3d" in observed_data
|
||
and "verts3d" in pred_data
|
||
and self.loss_weights["verts3d"] > 0.0
|
||
):
|
||
# raise ValueError
|
||
cur_loss = verts3d_loss(
|
||
observed_data["verts3d"], pred_data["verts3d"], valid_mask
|
||
)
|
||
loss += self.loss_weights["verts3d"] * cur_loss
|
||
stats_dict["verts3d"] = cur_loss
|
||
|
||
# All vertices to non-corresponding observed points in 3D space
|
||
if (
|
||
"points3d" in observed_data
|
||
and "points3d" in pred_data
|
||
and self.loss_weights["points3d"] > 0.0
|
||
):
|
||
# raise ValueError
|
||
cur_loss = self.points3d_loss(
|
||
observed_data["points3d"], pred_data["points3d"]
|
||
)
|
||
loss += self.loss_weights["points3d"] * cur_loss
|
||
stats_dict["points3d"] = cur_loss
|
||
|
||
# 2D re-projection loss
|
||
if (
|
||
"joints2d" in observed_data
|
||
and "joints3d_op" in pred_data
|
||
and "cameras" in pred_data
|
||
and self.loss_weights["joints2d"] > 0.0
|
||
):
|
||
# print(4444444444444444444)
|
||
joints2d = cam_util.reproject(
|
||
pred_data["joints3d_op"], *pred_data["cameras"]
|
||
)
|
||
cur_loss = self.joints2d_loss(
|
||
observed_data["joints2d"], joints2d, valid_mask
|
||
)
|
||
|
||
loss += self.loss_weights["joints2d"] * cur_loss
|
||
stats_dict["joints2d"] = cur_loss
|
||
|
||
# smooth 3d joint motion
|
||
if self.loss_weights["joints3d_smooth"] > 0.0:
|
||
# print(55555555555555555555)
|
||
cur_loss = joints3d_smooth_loss(pred_data["joints3d"], valid_mask)
|
||
loss += self.loss_weights["joints3d_smooth"] * cur_loss
|
||
stats_dict["joints3d_smooth"] = cur_loss
|
||
|
||
# If we're optimizing cameras, camera reprojection loss
|
||
if "bg2d_err" in pred_data and self.loss_weights["bg2d"] > 0.0:
|
||
raise ValueError
|
||
cur_loss = pred_data["bg2d_err"]
|
||
loss += self.loss_weights["bg2d"] * cur_loss
|
||
stats_dict["bg2d_err"] = cur_loss
|
||
|
||
# camera smoothness
|
||
if "cam_R" in pred_data and self.loss_weights["cam_R_smooth"] > 0.0:
|
||
raise ValueError
|
||
cam_R = pred_data["cam_R"] # (T, 3, 3)
|
||
cur_loss = rotation_smoothness_loss(cam_R[1:], cam_R[:-1])
|
||
loss += self.loss_weights["cam_R_smooth"] * cur_loss
|
||
stats_dict["cam_R_smooth"] = cur_loss
|
||
|
||
if "cam_t" in pred_data and self.loss_weights["cam_t_smooth"] > 0.0:
|
||
raise ValueError
|
||
cam_t = pred_data["cam_t"] # (T, 3, 3)
|
||
cur_loss = translation_smoothness_loss(cam_t[1:], cam_t[:-1])
|
||
loss += self.loss_weights["cam_t_smooth"] * cur_loss
|
||
stats_dict["cam_t_smooth"] = cur_loss
|
||
|
||
# Depth constraint: keep hand in front of camera
|
||
if (
|
||
"joints3d" in pred_data
|
||
and "cameras" in pred_data
|
||
and self.loss_weights.get("depth_constraint", 0.0) > 0.0
|
||
):
|
||
# Extract cam_R and cam_t from cameras tuple
|
||
# cameras = (cam_R, cam_t, cam_f, cam_center)
|
||
cam_R, cam_t, cam_f, cam_center = pred_data["cameras"]
|
||
cur_loss = depth_constraint_loss(
|
||
pred_data["joints3d"],
|
||
cam_R,
|
||
cam_t,
|
||
min_depth=0.0,
|
||
max_depth=999
|
||
)
|
||
loss += self.loss_weights["depth_constraint"] * cur_loss
|
||
stats_dict["depth_constraint"] = cur_loss
|
||
|
||
print(stats_dict)
|
||
return loss, stats_dict
|
||
|
||
|
||
def rotation_smoothness_loss(R1, R2):
|
||
R12 = torch.einsum("...ij,...jk->...ik", R2, R1.transpose(-1, -2))
|
||
aa12 = rotation_matrix_to_angle_axis(R12)
|
||
return torch.sum(aa12**2)
|
||
|
||
|
||
def translation_smoothness_loss(t1, t2):
|
||
return torch.sum((t2 - t1) ** 2)
|
||
|
||
|
||
def depth_constraint_loss(joints3d, cam_R, cam_t, min_depth=0.0, max_depth=999):
|
||
"""
|
||
Penalize joints that are behind the camera or too far away.
|
||
Ensures hand stays within reasonable depth range in camera space.
|
||
|
||
:param joints3d (B, T, J, 3) or (T, J, 3) joints in world space
|
||
:param cam_R (B, T, 3, 3) or (T, 3, 3) world-to-camera rotation
|
||
:param cam_t (B, T, 3) or (T, 3) world-to-camera translation
|
||
:param min_depth minimum allowed Z in camera space (default 0.1m)
|
||
:param max_depth maximum allowed Z in camera space (default 5.0m)
|
||
"""
|
||
# Handle both (B, T, J, 3) and (T, J, 3) shapes
|
||
if joints3d.ndim == 3:
|
||
# (T, J, 3) -> add batch dimension
|
||
joints3d = joints3d.unsqueeze(0) # (1, T, J, 3)
|
||
if cam_R.ndim == 3:
|
||
# (T, 3, 3) -> add batch dimension
|
||
cam_R = cam_R.unsqueeze(0) # (1, T, 3, 3)
|
||
if cam_t.ndim == 2:
|
||
# (T, 3) -> add batch dimension
|
||
cam_t = cam_t.unsqueeze(0) # (1, T, 3)
|
||
|
||
# Transform joints to camera space
|
||
B, T, J, _ = joints3d.shape
|
||
joints_cam = torch.einsum("btij,btjk->btik", cam_R, joints3d.transpose(-1, -2)).transpose(-1, -2)
|
||
joints_cam = joints_cam + cam_t[..., None, :] # (B, T, J, 3)
|
||
|
||
# Get Z coordinate (depth)
|
||
depth = joints_cam[..., 2] # (B, T, J)
|
||
|
||
# Penalize negative depth (behind camera) heavily
|
||
behind_camera_loss = torch.sum(torch.relu(-depth) ** 2)
|
||
|
||
# Penalize depth outside reasonable range
|
||
too_close_loss = torch.sum(torch.relu(min_depth - depth) ** 2)
|
||
too_far_loss = torch.sum(torch.relu(depth - max_depth) ** 2)
|
||
|
||
return behind_camera_loss * 100.0 + too_close_loss + too_far_loss
|
||
|
||
|
||
def camera_smoothness_loss(R1, t1, R2, t2):
|
||
"""
|
||
:param R1, t1 (N, 3, 3), (N, 3)
|
||
:param R2, t2 (N, 3, 3), (N, 3)
|
||
"""
|
||
R12, t12 = cam_util.compose_cameras(R2, t2, *cam_util.invert_camera(R1, t1))
|
||
aa12 = rotation_matrix_to_angle_axis(R12)
|
||
return torch.sum(aa12**2) + torch.sum(t12**2)
|
||
|
||
"""
|
||
Losses are cumulative
|
||
SMPLLoss setup is same as RootLoss
|
||
"""
|
||
|
||
class SMPLLoss(RootLoss):
|
||
def forward(self, observed_data, pred_data, nsteps, valid_mask=None):
|
||
"""
|
||
For fitting full shape and pose of SMPL.
|
||
nsteps used to scale single-step losses
|
||
"""
|
||
loss, stats_dict = super().forward(
|
||
observed_data, pred_data, valid_mask=valid_mask
|
||
)
|
||
|
||
# prior to keep latent pose likely
|
||
print(f"DEBUG pose_prior check: 'latent_pose' in pred_data = {'latent_pose' in pred_data}, loss_weight = {self.loss_weights.get('pose_prior', 'NOT_FOUND')}")
|
||
if "latent_pose" in pred_data:
|
||
print(f" latent_pose shape: {pred_data['latent_pose'].shape}")
|
||
if "init_latent_pose" in observed_data:
|
||
print(f" init_latent_pose shape: {observed_data['init_latent_pose'].shape}")
|
||
|
||
if "latent_pose" in pred_data and self.loss_weights["pose_prior"] > 0.0:
|
||
# Use initial latent pose from HaMeR as the prior target
|
||
latent_pose_init = observed_data["init_latent_pose"]
|
||
cur_loss = pose_prior_loss(pred_data["latent_pose"], latent_pose_init, valid_mask)
|
||
print("pose_prior: ", cur_loss, pred_data["latent_pose"][0],latent_pose_init[0])
|
||
loss += self.loss_weights["pose_prior"] * cur_loss
|
||
stats_dict["pose_prior"] = cur_loss
|
||
else:
|
||
raise ValueError
|
||
|
||
# prior to keep PCA shape likely
|
||
if "betas" in pred_data and self.loss_weights["shape_prior"] > 0.0:
|
||
cur_loss = shape_prior_loss(pred_data["betas"])
|
||
loss += self.loss_weights["shape_prior"] * nsteps * cur_loss
|
||
stats_dict["shape_prior"] = cur_loss
|
||
|
||
# inter-penetration loss
|
||
if 'verts3d' in pred_data and self.loss_weights["penetration"] > 0.0:
|
||
print('???????????????????????????????? inter-penetration loss')
|
||
# print(pred_data["verts3d"].shape, pred_data['is_right'].shape)
|
||
order = torch.sum(pred_data['is_right'], dim=-1)//pred_data['is_right'].shape[1]
|
||
|
||
if len(order) > 1:
|
||
cur_loss = 0
|
||
if order[0] == 0:
|
||
l_verts = pred_data["verts3d"][0]
|
||
r_verts = pred_data["verts3d"][1]
|
||
else:
|
||
l_verts = pred_data["verts3d"][1]
|
||
r_verts = pred_data["verts3d"][0]
|
||
|
||
for idx in range(len(l_verts)):
|
||
cur_loss += self.inter_penetration_loss(v1=l_verts[idx:idx+1], v2=r_verts[idx:idx+1])
|
||
|
||
loss += self.loss_weights["penetration"] * cur_loss
|
||
stats_dict["penetration"] = cur_loss
|
||
|
||
print(stats_dict)
|
||
# if 'penetration' in stats_dict.keys():
|
||
# if stats_dict['penetration'] == 0:
|
||
# print(pred_data)
|
||
# print('***************')
|
||
# print(observed_data)
|
||
# raise ValueError
|
||
return loss, stats_dict
|
||
|
||
# class MotionLoss(SMPLLoss):
|
||
# def setup_losses(
|
||
# self,
|
||
# loss_weights,
|
||
# init_motion_prior=None,
|
||
# **kwargs,
|
||
# ):
|
||
# super().setup_losses(loss_weights, **kwargs)
|
||
# if loss_weights["init_motion_prior"] > 0.0:
|
||
# self.init_motion_prior_loss = GMMPriorLoss(init_motion_prior)
|
||
|
||
# def forward(
|
||
# self,
|
||
# observed_data,
|
||
# pred_data,
|
||
# cam_pred_data,
|
||
# nsteps,
|
||
# valid_mask=None,
|
||
# init_motion_scale=1.0,
|
||
# ):
|
||
# """
|
||
# For fitting full shape and pose of SMPL with motion prior.
|
||
|
||
# pred_data is data pertinent to the canonical prior coordinate frame
|
||
# cam_pred_data is for the camera coordinate frame
|
||
|
||
# loss rather than standard normal if given.
|
||
# """
|
||
# cam_pred_data["latent_pose"] = pred_data["latent_pose"]
|
||
# loss, stats_dict = super().forward(
|
||
# observed_data, cam_pred_data, nsteps, valid_mask=valid_mask
|
||
# )
|
||
|
||
# # prior to keep latent motion likely
|
||
# if "latent_motion" in pred_data and self.loss_weights["motion_prior"] > 0.0:
|
||
# # NOTE: latent is NOT synchronized in time,
|
||
# # the mask is NOT relevant
|
||
# # Generate the async mask to properly mask motion prior loss
|
||
# # Helps to calibrate the range of 'good' loss values
|
||
# B, T, _ = pred_data["latent_motion"].shape
|
||
# device = pred_data["latent_motion"].device
|
||
# async_mask = (
|
||
# torch.arange(T, device=device)[None].expand(B, -1)
|
||
# < valid_mask.sum(dim=1)[:, None]
|
||
# )
|
||
# cur_loss = motion_prior_loss(
|
||
# pred_data["latent_motion"],
|
||
# cond_prior=pred_data.get("cond_prior", None),
|
||
# mask=async_mask,
|
||
# )
|
||
# loss += self.loss_weights["motion_prior"] * cur_loss
|
||
# stats_dict["motion_prior"] = cur_loss
|
||
|
||
# # prior to keep initial state likely
|
||
# have_init_prior_info = (
|
||
# "joints3d_init" in pred_data
|
||
# and "joints_vel" in pred_data
|
||
# and "trans_vel" in pred_data
|
||
# and "root_orient_vel" in pred_data
|
||
# )
|
||
# if have_init_prior_info and self.loss_weights["init_motion_prior"] > 0.0:
|
||
# cur_loss = self.init_motion_prior_loss(
|
||
# pred_data["joints3d_init"],
|
||
# pred_data["joints_vel"],
|
||
# pred_data["trans_vel"],
|
||
# pred_data["root_orient_vel"],
|
||
# )
|
||
# loss += (
|
||
# self.loss_weights["init_motion_prior"] * init_motion_scale * cur_loss
|
||
# ) # must scale since doesn't scale with more steps
|
||
# stats_dict["init_motion_prior"] = cur_loss
|
||
|
||
# # make sure joints consistent between SMPL and direct motion prior output
|
||
# if (
|
||
# "joints3d_rollout" in pred_data
|
||
# and "joints3d" in pred_data
|
||
# and self.loss_weights["joint_consistency"] > 0.0
|
||
# ):
|
||
# cur_loss = joint_consistency_loss(
|
||
# pred_data["joints3d"],
|
||
# pred_data["joints3d_rollout"],
|
||
# valid_mask,
|
||
# )
|
||
# loss += self.loss_weights["joint_consistency"] * cur_loss
|
||
# stats_dict["joint_consistency"] = cur_loss
|
||
|
||
# # make sure bone lengths between frames of direct motion prior output are consistent
|
||
# if "joints3d_rollout" in pred_data and self.loss_weights["bone_length"] > 0.0:
|
||
# cur_loss = bone_length_loss(pred_data["joints3d_rollout"], valid_mask)
|
||
# loss += self.loss_weights["bone_length"] * cur_loss
|
||
# stats_dict["bone_length"] = cur_loss
|
||
|
||
# # make sure rolled out joints match observations too
|
||
# if (
|
||
# "joints3d" in observed_data
|
||
# and "joints3d_rollout" in pred_data
|
||
# and self.loss_weights["joints3d_rollout"] > 0.0
|
||
# ):
|
||
# raise ValueError
|
||
# # cur_loss = joints3d_loss(
|
||
# # observed_data["joints3d"], pred_data["joints3d_rollout"], valid_mask
|
||
# # )
|
||
# # loss += self.loss_weights["joints3d_rollout"] * cur_loss
|
||
# # stats_dict["joints3d_rollout"] = cur_loss
|
||
|
||
# # velocity 0 during contacts
|
||
# if (
|
||
# self.loss_weights["contact_vel"] > 0.0
|
||
# and "contacts_conf" in pred_data
|
||
# and "joints3d" in pred_data
|
||
# ):
|
||
# raise ValueError
|
||
# # cur_loss = contact_vel_loss(
|
||
# # pred_data["contacts_conf"], pred_data["joints3d"], valid_mask
|
||
# # )
|
||
# # loss += self.loss_weights["contact_vel"] * cur_loss
|
||
# # stats_dict["contact_vel"] = cur_loss
|
||
|
||
# # contacting joints are near the floor
|
||
# if (
|
||
# self.loss_weights["contact_height"] > 0.0
|
||
# and "contacts_conf" in pred_data
|
||
# and "joints3d" in pred_data
|
||
# ):
|
||
# raise ValueError
|
||
# # cur_loss = contact_height_loss(
|
||
# # pred_data["contacts_conf"], pred_data["joints3d"], valid_mask
|
||
# # )
|
||
# # loss += self.loss_weights["contact_height"] * cur_loss
|
||
# # stats_dict["contact_height"] = cur_loss
|
||
|
||
# # floor is close to the initialization
|
||
# if (
|
||
# self.loss_weights["floor_reg"] > 0.0
|
||
# and "floor_plane" in pred_data
|
||
# and "floor_plane" in observed_data
|
||
# ):
|
||
# raise ValueError
|
||
# # cur_loss = floor_reg_loss(
|
||
# # pred_data["floor_plane"], observed_data["floor_plane"]
|
||
# # )
|
||
# # loss += self.loss_weights["floor_reg"] * nsteps * cur_loss
|
||
# # stats_dict["floor_reg"] = cur_loss
|
||
|
||
# return loss, stats_dict
|
||
|
||
|
||
def joints3d_loss(joints3d_obs, joints3d_pred, mask=None):
|
||
"""
|
||
:param joints3d_obs (B, T, J, 3)
|
||
:param joints3d_pred (B, T, J, 3)
|
||
:param mask (optional) (B, T)
|
||
"""
|
||
B, T, *dims = joints3d_obs.shape
|
||
vis_mask = get_visible_mask(joints3d_obs)
|
||
if mask is not None:
|
||
vis_mask = vis_mask & mask.reshape(B, T, *(1,) * len(dims)).bool()
|
||
loss = (joints3d_obs[vis_mask] - joints3d_pred[vis_mask]) ** 2
|
||
loss = 0.5 * torch.sum(loss)
|
||
return loss
|
||
|
||
|
||
def verts3d_loss(verts3d_obs, verts3d_pred, mask=None):
|
||
"""
|
||
:param verts3d_obs (B, T, V, 3)
|
||
:param verts3d_pred (B, T, V, 3)
|
||
:param mask (optional) (B, T)
|
||
"""
|
||
B, T, *dims = verts3d_obs.shape
|
||
vis_mask = get_visible_mask(verts3d_obs)
|
||
if mask is not None:
|
||
assert mask.shape == (B, T)
|
||
vis_mask = vis_mask & mask.reshape(B, T, *(1,) * len(dims)).bool()
|
||
loss = (verts3d_obs[vis_mask] - verts3d_pred[vis_mask]) ** 2
|
||
loss = 0.5 * torch.sum(loss)
|
||
return loss
|
||
|
||
|
||
def get_visible_mask(obs_data):
|
||
"""
|
||
Given observed data gets the mask of visible data (that actually contributes to the loss).
|
||
"""
|
||
return torch.logical_not(torch.isinf(obs_data))
|
||
|
||
|
||
class Joints2DLoss(nn.Module):
|
||
def __init__(self, ignore_op_joints=None, joints2d_sigma=100, normalize_by_scale=True):
|
||
super().__init__()
|
||
self.ignore_op_joints = ignore_op_joints
|
||
self.joints2d_sigma = joints2d_sigma
|
||
self.normalize_by_scale = normalize_by_scale
|
||
|
||
def forward(self, joints2d_obs, joints2d_pred, mask=None):
|
||
"""
|
||
:param joints2d_obs (B, T, 25, 3)
|
||
:param joints2d_pred (B, T, 22, 2)
|
||
:param mask (optional) (B, T)
|
||
"""
|
||
# Safety check: detect NaN in predictions before computing loss
|
||
if torch.isnan(joints2d_pred).any() or torch.isinf(joints2d_pred).any():
|
||
print(f"ERROR: joints2d_pred contains NaN or Inf!")
|
||
print(f" NaN count: {torch.isnan(joints2d_pred).sum().item()}")
|
||
print(f" This means MANO model produced invalid output")
|
||
return torch.tensor(1e8, device=joints2d_pred.device, requires_grad=True)
|
||
|
||
if mask is not None:
|
||
mask = mask.bool()
|
||
joints2d_obs = joints2d_obs[mask] # (N, 25, 3)
|
||
joints2d_pred = joints2d_pred[mask] # (N, 22, 2)
|
||
|
||
joints2d_obs_conf = joints2d_obs[..., 2:3]
|
||
if self.ignore_op_joints is not None:
|
||
# set confidence to 0 so not weighted
|
||
joints2d_obs_conf[..., self.ignore_op_joints, :] = 0.0
|
||
|
||
# Compute error
|
||
error = joints2d_pred - joints2d_obs[..., :2] # (N, 22, 2)
|
||
|
||
# Normalize by hand scale to make loss resolution-independent
|
||
if self.normalize_by_scale:
|
||
# Estimate hand scale from observed keypoints spread
|
||
valid_obs = joints2d_obs[:, :, :2] # (N, 25, 2)
|
||
# Compute bbox size of observed keypoints for each hand
|
||
kp_min = valid_obs.min(dim=1, keepdim=True)[0] # (N, 1, 2)
|
||
kp_max = valid_obs.max(dim=1, keepdim=True)[0] # (N, 1, 2)
|
||
hand_scale = torch.sqrt(((kp_max - kp_min) ** 2).sum(dim=-1, keepdim=True)) # (N, 1, 1)
|
||
|
||
# Safety check: detect invalid/degenerate cases
|
||
if torch.isnan(hand_scale).any() or torch.isinf(hand_scale).any() or (hand_scale < 5.0).any():
|
||
print(f"ERROR: Invalid hand_scale! min={hand_scale.min():.2f}, max={hand_scale.max():.2f}")
|
||
print(f" This indicates MANO produced degenerate output (collapsed joints)")
|
||
print(f" kp_min: {kp_min[0, 0]}, kp_max: {kp_max[0, 0]}")
|
||
# Return a large constant loss to signal the optimizer this is a bad configuration
|
||
return torch.tensor(1e6, device=hand_scale.device, requires_grad=True)
|
||
|
||
# Clamp to reasonable range (5-1000 pixels)
|
||
hand_scale = torch.clamp(hand_scale, min=5.0, max=1000.0)
|
||
|
||
# Normalize error by hand scale AND scale up to reasonable magnitude
|
||
error = error / hand_scale * 100.0 # Now error is in normalized units (0-100 scale)
|
||
|
||
# Use FIXED sigma in normalized space (not divided by hand_scale!)
|
||
normalized_sigma = self.joints2d_sigma # Already in normalized units
|
||
else:
|
||
normalized_sigma = self.joints2d_sigma
|
||
|
||
# weight errors by detection confidence
|
||
robust_sqr_dist = gmof(error, normalized_sigma)
|
||
reproj_err = (joints2d_obs_conf**2) * robust_sqr_dist
|
||
loss = torch.mean(reproj_err)
|
||
return loss
|
||
|
||
|
||
# class Joints2DLoss(nn.Module):
|
||
# def __init__(self, ignore_op_joints=None, joints2d_sigma=100):
|
||
# super().__init__()
|
||
# self.ignore_op_joints = ignore_op_joints
|
||
# self.joints2d_sigma = joints2d_sigma
|
||
|
||
# def forward(self, joints2d_obs, joints2d_pred, mask=None):
|
||
# """
|
||
# :param joints2d_obs (B, T, 25, 3)
|
||
# :param joints2d_pred (B, T, 22, 2)
|
||
# :param mask (optional) (B, T)
|
||
# """
|
||
# if mask is not None:
|
||
# mask = mask.bool()
|
||
# joints2d_obs = joints2d_obs[mask] # (N, 25, 3)
|
||
# joints2d_pred = joints2d_pred[mask] # (N, 22, 2)
|
||
|
||
# joints2d_obs_conf = joints2d_obs[..., 2:3]
|
||
# if self.ignore_op_joints is not None:
|
||
# # set confidence to 0 so not weighted
|
||
# joints2d_obs_conf[..., self.ignore_op_joints, :] = 0.0
|
||
|
||
# # Compute per-joint pixel errors
|
||
# error = joints2d_pred - joints2d_obs[..., :2] # (N, J, 2)
|
||
# pixel_error = torch.sqrt((error**2).sum(dim=-1, keepdim=True)) # (N, J, 1)
|
||
|
||
# # Use Smooth L1 (Huber) loss: L2 for small errors, L1 for large errors
|
||
# # This prevents gradient explosion while keeping strong gradients for normal errors
|
||
# # Beta=1000 means: errors <1000px get L2 (strong gradients), errors >1000px get L1 (bounded)
|
||
# beta = 1000.0 # Transition point: L2 below, L1 above
|
||
|
||
# smooth_l1 = torch.nn.functional.smooth_l1_loss(
|
||
# joints2d_pred,
|
||
# joints2d_obs[..., :2],
|
||
# reduction='none',
|
||
# beta=beta
|
||
# ) # (N, J, 2)
|
||
|
||
# sqr_dist = smooth_l1
|
||
|
||
# # DEBUG
|
||
# print(f'Pixel errors: mean={pixel_error.mean():.1f}, max={pixel_error.max():.1f}, median={pixel_error.median():.1f}')
|
||
# num_large = (pixel_error > beta).sum().item()
|
||
# print(f'L2 loss for {pixel_error.numel() - num_large}/{pixel_error.numel()} joints, L1 for {num_large} large >{beta}px')
|
||
# print(f'Loss stats: mean={sqr_dist.mean():.2f}, max={sqr_dist.max():.2f}')
|
||
|
||
# # Weight by confidence and compute loss
|
||
# reproj_err = (joints2d_obs_conf**2) * sqr_dist
|
||
# reproj_err = torch.clamp(reproj_err, max=1000)
|
||
# loss = torch.mean(reproj_err)
|
||
# return loss
|
||
|
||
|
||
class Points3DLoss(nn.Module):
|
||
def __init__(
|
||
self,
|
||
use_chamfer=False,
|
||
robust_loss="bisquare",
|
||
robust_tuning_const=4.6851,
|
||
):
|
||
super().__init__()
|
||
|
||
if not use_chamfer:
|
||
self.active = False
|
||
return
|
||
|
||
self.active = True
|
||
|
||
robust_choices = ["none", "bisquare", "gm"]
|
||
if robust_loss not in robust_choices:
|
||
Logger.log(
|
||
"Not a valid robust loss: %s. Please use %s"
|
||
% (robust_loss, str(robust_choices))
|
||
)
|
||
exit()
|
||
|
||
from utils.chamfer_distance import ChamferDistance
|
||
|
||
self.chamfer_dist = ChamferDistance()
|
||
|
||
self.robust_loss = robust_loss
|
||
self.robust_tuning_const = robust_tuning_const
|
||
|
||
def forward(self, points3d_obs, points3d_pred):
|
||
if not self.active:
|
||
return torch.tensor(0.0, dtype=torch.float32, device=points3d_obs.device)
|
||
|
||
# one-way chamfer
|
||
B, T, N_obs, _ = points3d_obs.size()
|
||
N_pred = points3d_pred.size(2)
|
||
points3d_obs = points3d_obs.reshape((B * T, -1, 3))
|
||
points3d_pred = points3d_pred.reshape((B * T, -1, 3))
|
||
|
||
obs2pred_sqr_dist, pred2obs_sqr_dist = self.chamfer_dist(
|
||
points3d_obs, points3d_pred
|
||
)
|
||
obs2pred_sqr_dist = obs2pred_sqr_dist.reshape((B, T * N_obs))
|
||
pred2obs_sqr_dist = pred2obs_sqr_dist.reshape((B, T * N_pred))
|
||
|
||
weighted_obs2pred_sqr_dist, w = apply_robust_weighting(
|
||
obs2pred_sqr_dist.sqrt(),
|
||
robust_loss_type=self.robust_loss,
|
||
robust_tuning_const=self.robust_tuning_const,
|
||
)
|
||
|
||
loss = torch.sum(weighted_obs2pred_sqr_dist)
|
||
loss = 0.5 * loss
|
||
return loss
|
||
|
||
|
||
######################
|
||
# losses
|
||
######################
|
||
class GeneralContactLoss(nn.Module):
|
||
def __init__(
|
||
self,
|
||
faces,
|
||
region_aggregation_type: str = 'sum',
|
||
squared_dist: bool = False,
|
||
model_type: str = 'mano',
|
||
body_model_utils_folder: str = 'body_model_utils',
|
||
**kwargs
|
||
):
|
||
super().__init__()
|
||
"""
|
||
Compute intersection and contact between two meshes and resolves.
|
||
"""
|
||
|
||
self.region_aggregation_type = region_aggregation_type
|
||
self.squared = squared_dist
|
||
|
||
self.criterion = self.init_loss()
|
||
|
||
# create extra vertex and faces to close back of the mouth to maske
|
||
# the mano mesh watertight.
|
||
self.model_type = model_type
|
||
|
||
if self.model_type == 'mano':
|
||
# add faces that make the hand mesh watertight
|
||
faces_new = torch.tensor([[92, 38, 234],
|
||
[234, 38, 239],
|
||
[38, 122, 239],
|
||
[239, 122, 279],
|
||
[122, 118, 279],
|
||
[279, 118, 215],
|
||
[118, 117, 215],
|
||
[215, 117, 214],
|
||
[117, 119, 214],
|
||
[214, 119, 121],
|
||
[119, 120, 121],
|
||
[121, 120, 78],
|
||
[120, 108, 78],
|
||
[78, 108, 79]])
|
||
|
||
r_faces = torch.cat([faces, faces_new.to(faces.device)], dim=0)
|
||
l_faces = r_faces[:,[0,2,1]].clone()
|
||
|
||
self.register_buffer('l_faces', l_faces)
|
||
self.register_buffer('r_faces', r_faces)
|
||
|
||
# low resolution mesh
|
||
# inner_mouth_verts_path = f'{body_model_utils_folder}/lowres_{model_type}.pkl'
|
||
# self.low_res_mesh = pickle.load(open(inner_mouth_verts_path, 'rb'))
|
||
|
||
# def close_mouth(self, v):
|
||
# mv = torch.mean(v[:,self.vert_ids_wt,:], 1, keepdim=True)
|
||
# v = torch.cat((v, mv), 1)
|
||
# return v
|
||
|
||
def to_lowres(self, v, is_right):
|
||
v = v
|
||
if is_right == 0:
|
||
t = v[:,self.l_faces,:]
|
||
elif is_right == 1:
|
||
t = v[:,self.r_faces,:]
|
||
else:
|
||
raise ValueError
|
||
|
||
return v, t
|
||
|
||
def init_loss(self):
|
||
def loss_func(v1, v2, factor=100, wn_batch=True):
|
||
"""
|
||
Compute loss between region r1 on meshes v1 and
|
||
region r2 on mesh v2.
|
||
"""
|
||
|
||
nn = 1000
|
||
|
||
loss = torch.tensor(0.0, device=v1.device)
|
||
|
||
if wn_batch:
|
||
# close mouth for self-intersection test
|
||
v1l, t1l = self.to_lowres(v1, 0)
|
||
v2l, t2l = self.to_lowres(v2, 1)
|
||
|
||
# compute intersection between v1 and v2
|
||
# print(v1.shape, t1l.shape)
|
||
# exit()
|
||
interior_v1 = winding_numbers(v1, t2l).ge(0.99)
|
||
interior_v2 = winding_numbers(v2, t1l).ge(0.99)
|
||
|
||
# visu results
|
||
# if True:
|
||
# import trimesh
|
||
# mesh = trimesh.Trimesh(
|
||
# vertices=v2l[0].detach().cpu().numpy(),
|
||
# faces=self.l_faces.cpu().numpy().astype(np.int32),
|
||
# )
|
||
# col = 255 * np.ones((778, 4))
|
||
# inside_idx = torch.where(interior_v2[0])[0].detach().cpu().numpy()
|
||
# col[inside_idx] = [0, 255, 0, 255]
|
||
# mesh.visual.vertex_colors = col
|
||
# _ = mesh.export('/vol/bitbucket/zy3023/code/hand/slam-hand//interior_v2_lowres.obj')
|
||
|
||
# vmc = v1 # self.close_mouth(v1)
|
||
# mesh = trimesh.Trimesh(vertices=vmc[0].detach().cpu().numpy(), faces=self.l_faces.cpu().numpy())
|
||
# col = 255 * np.ones((778, 4))
|
||
# inside_idx = torch.where(interior_v1[0])[0].detach().cpu().numpy()
|
||
# col[inside_idx] = [255, 0, 0, 255]
|
||
# mesh.visual.vertex_colors = col
|
||
# _ = mesh.export('/vol/bitbucket/zy3023/code/hand/slam-hand//interior_v1_highres.obj')
|
||
|
||
# # Other direction
|
||
# mesh = trimesh.Trimesh(
|
||
# vertices=v1l[0].detach().cpu().numpy(),
|
||
# faces=self.r_faces.cpu().numpy().astype(np.int32),
|
||
# )
|
||
# col = 255 * np.ones((778, 4))
|
||
# inside_idx = torch.where(interior_v1[0])[0].detach().cpu().numpy()
|
||
# col[inside_idx] = [0, 255, 0, 255]
|
||
# mesh.visual.vertex_colors = col
|
||
# _ = mesh.export('/vol/bitbucket/zy3023/code/hand/slam-hand//interior_v1_lowres.obj')
|
||
|
||
# vmc = v2 # self.close_mouth(v2)
|
||
# mesh = trimesh.Trimesh(vertices=vmc[0].detach().cpu().numpy(), faces=self.r_faces.cpu().numpy())
|
||
# col = 255 * np.ones((778, 4))
|
||
# inside_idx = torch.where(interior_v2[0])[0].detach().cpu().numpy()
|
||
# col[inside_idx] = [255, 0, 0, 255]
|
||
# mesh.visual.vertex_colors = col
|
||
# _ = mesh.export('/vol/bitbucket/zy3023/code/hand/slam-hand//interior_v2_highres.obj')
|
||
# exit()
|
||
|
||
batch_losses = []
|
||
for bidx in range(v1.shape[0]):
|
||
if not wn_batch:
|
||
# close mouth for self-intersection test
|
||
v1l, t1l = self.to_lowres(v1[[bidx]], nn)
|
||
v2l, t2l = self.to_lowres(v2[[bidx]], nn)
|
||
|
||
# compute intersection between v1 and v2
|
||
curr_interior_v1 = winding_numbers(v1[[bidx]], t2l.detach()).ge(0.99)[0]
|
||
curr_interior_v2 = winding_numbers(v2[[bidx]], t1l.detach()).ge(0.99)[0]
|
||
crit_v1, crit_v2 = torch.any(curr_interior_v1), torch.any(curr_interior_v2)
|
||
else:
|
||
curr_interior_v1 = interior_v1[bidx]
|
||
curr_interior_v2 = interior_v2[bidx]
|
||
crit_v1, crit_v2 = torch.any(interior_v1[bidx]), torch.any(interior_v2[bidx])
|
||
|
||
bloss = torch.tensor(0.0, device=v1.device)
|
||
if crit_v1 and crit_v2:
|
||
# find vertices that are close to each other between v1 and v2
|
||
#squared_dist = pcl_pcl_pairwise_distance(
|
||
# v1[:,interior_v1[bidx],:], v2[:, interior_v2[bidx], :], squared=self.squared
|
||
#)
|
||
squared_dist_v1v2 = pcl_pcl_pairwise_distance(
|
||
v1[[[bidx]],curr_interior_v1,:], v2[[bidx]], squared=self.squared)
|
||
squared_dist_v2v1 = pcl_pcl_pairwise_distance(
|
||
v2[[[bidx]], curr_interior_v2, :], v1[[bidx]], squared=self.squared)
|
||
|
||
v1_to_v2 = (squared_dist_v1v2[0].min(1)[0] * factor)**2
|
||
#v1_to_v2 = 10.0 * (torch.tanh(v1_to_v2 / 10.0)**2)
|
||
bloss += v1_to_v2.sum()
|
||
|
||
v2_to_v1 = (squared_dist_v2v1[0].min(1)[0] * factor)**2
|
||
#v2_to_v1 = 10.0 * (torch.tanh(v2_to_v1 / 10.0)**2)
|
||
bloss += v2_to_v1.sum()
|
||
|
||
batch_losses.append(bloss)
|
||
|
||
# compute loss
|
||
if len(batch_losses) > 0:
|
||
loss = sum(batch_losses) / len(batch_losses)
|
||
|
||
#if loss > 1.0:
|
||
# import ipdb; ipdb.set_trace()
|
||
|
||
return loss
|
||
|
||
return loss_func
|
||
|
||
def forward(self, **args):
|
||
losses = self.criterion(**args)
|
||
return losses
|
||
|
||
def pose_prior_loss(latent_pose_pred, latent_pose_init=None, mask=None):
|
||
"""
|
||
:param latent_pose_pred (B, T, D) - optimized latent pose
|
||
:param latent_pose_init (optional) (B, T, D) - initial HaMeR prediction
|
||
:param mask (optional) (B, T)
|
||
"""
|
||
if latent_pose_init is not None:
|
||
# Penalize deviation from original HaMeR prediction
|
||
loss = (latent_pose_pred - latent_pose_init)**2
|
||
else:
|
||
# Fallback: 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 shape_prior_loss(betas_pred):
|
||
# prior is isotropic gaussian so take L2 distance from 0
|
||
loss = betas_pred**2
|
||
loss = torch.sum(loss)
|
||
return loss
|
||
|
||
|
||
def joints3d_smooth_loss(joints3d_pred, mask=None, normalize_by_scale=True):
|
||
"""
|
||
:param joints3d_pred (B, T, J, 3)
|
||
:param mask (optional) (B, T)
|
||
:param normalize_by_scale: If True, normalize by hand scale for size-invariance
|
||
"""
|
||
# Safety check: detect NaN in predictions before computing loss
|
||
if torch.isnan(joints3d_pred).any() or torch.isinf(joints3d_pred).any():
|
||
print(f"ERROR: joints3d_pred contains NaN or Inf in smooth loss!")
|
||
print(f" NaN count: {torch.isnan(joints3d_pred).sum().item()}")
|
||
return torch.tensor(1e8, device=joints3d_pred.device, requires_grad=True)
|
||
|
||
B, T, *dims = joints3d_pred.shape
|
||
|
||
# Normalize by hand scale FIRST to make loss size-invariant
|
||
if normalize_by_scale:
|
||
# Estimate hand scale from 3D joint positions at each frame
|
||
# Use bbox of joints as scale reference
|
||
joints_min = joints3d_pred.min(dim=2, keepdim=True)[0] # (B, T, 1, 3)
|
||
joints_max = joints3d_pred.max(dim=2, keepdim=True)[0] # (B, T, 1, 3)
|
||
hand_scale = torch.sqrt(((joints_max - joints_min) ** 2).sum(dim=-1, keepdim=True)) # (B, T, 1, 1)
|
||
|
||
# Safety check: detect invalid/degenerate cases
|
||
if torch.isnan(hand_scale).any() or torch.isinf(hand_scale).any() or (hand_scale < 0.001).any():
|
||
print(f"ERROR: Invalid hand_scale in joints3d_smooth! min={hand_scale.min():.4f}, max={hand_scale.max():.4f}")
|
||
print(f" This indicates MANO produced degenerate 3D output")
|
||
# Return a large constant loss
|
||
return torch.tensor(1e6, device=hand_scale.device, requires_grad=True)
|
||
|
||
# Clamp to reasonable range (0.001m - 1m for hand size)
|
||
hand_scale = torch.clamp(hand_scale, min=0.001, max=1.0)
|
||
|
||
# Normalize joints FIRST by their own scale at each frame (same as joints2d)
|
||
joints3d_pred_normalized = joints3d_pred / hand_scale # (B, T, J, 3) - same scale factor as joints2d
|
||
else:
|
||
joints3d_pred_normalized = joints3d_pred
|
||
|
||
# Now compute delta on normalized joints
|
||
delta = joints3d_pred_normalized[:, 1:, :, :] - joints3d_pred_normalized[:, :-1, :, :] # (B, T-1, J, 3)
|
||
|
||
loss = delta ** 2
|
||
if mask is not None:
|
||
mask = mask.bool()
|
||
mask = mask[:, 1:] & mask[:, :-1]
|
||
loss = loss[mask]
|
||
|
||
loss = torch.mean(loss, dim=0)
|
||
loss = 0.5 * torch.sum(loss)
|
||
return loss
|
||
|
||
# def motion_prior_loss(latent_motion_pred, cond_prior=None, mask=None):
|
||
# """
|
||
# :param latent_motion_pred (B, T, D)
|
||
# :param cond_prior (optional) (B, T, D, 2) stacked mean and var of post distribution
|
||
# :param mask (optional) (B, T)
|
||
# """
|
||
# if mask is not None:
|
||
# latent_motion_pred = latent_motion_pred[mask.bool()]
|
||
|
||
# if cond_prior is None:
|
||
# # assume standard normal
|
||
# loss = latent_motion_pred**2
|
||
# else:
|
||
# pm, pv = cond_prior[..., 0], cond_prior[..., 1]
|
||
# if mask is not None:
|
||
# mask = mask.bool()
|
||
# pm, pv = pm[mask], pv[mask]
|
||
# loss = -log_normal(latent_motion_pred, pm, pv)
|
||
|
||
# return torch.sum(loss)
|
||
|
||
# class GMMPriorLoss(nn.Module):
|
||
# def __init__(self, init_motion_prior=None):
|
||
# super().__init__()
|
||
# if init_motion_prior is None:
|
||
# self.active = False
|
||
# return
|
||
|
||
# # build pytorch GMM
|
||
# self.active = True
|
||
# self.init_motion_prior = dict()
|
||
# gmm = build_gmm(*init_motion_prior["gmm"])
|
||
# self.init_motion_prior["gmm"] = gmm
|
||
|
||
# def forward(self, joints, joints_vel, trans_vel, root_orient_vel):
|
||
# if not self.active:
|
||
# return torch.tensor(0.0, dtype=torch.float32, device=joints.device)
|
||
|
||
# # create input
|
||
# B = joints.size(0)
|
||
|
||
# joints = joints.reshape((B, -1))
|
||
# joints_vel = joints_vel.reshape((B, -1))
|
||
# trans_vel = trans_vel.reshape((B, -1))
|
||
# root_orient_vel = root_orient_vel.reshape((B, -1))
|
||
# init_state = torch.cat([joints, joints_vel, trans_vel, root_orient_vel], dim=-1)
|
||
|
||
# loss = -self.init_motion_prior["gmm"].log_prob(init_state)
|
||
# loss = torch.sum(loss)
|
||
|
||
# return loss
|
||
|
||
|
||
# def build_gmm(gmm_weights, gmm_means, gmm_covs):
|
||
# mix = Categorical(gmm_weights)
|
||
# comp = MultivariateNormal(gmm_means, covariance_matrix=gmm_covs)
|
||
# gmm_distrib = MixtureSameFamily(mix, comp)
|
||
# return gmm_distrib
|
||
|
||
|
||
# def joint_consistency_loss(smpl_joints3d, rollout_joints3d, mask=None):
|
||
# """
|
||
# :param smpl_joints3d (B, T, J, 3)
|
||
# :param rollout_joints3d (B, T, J, 3)
|
||
# :param mask (optional) (B, T)
|
||
# """
|
||
# loss = (smpl_joints3d - rollout_joints3d) ** 2
|
||
# if mask is not None:
|
||
# loss = loss[mask.bool()]
|
||
# loss = 0.5 * torch.sum(loss)
|
||
# return loss
|
||
|
||
|
||
# def bone_length_loss(rollout_joints3d, mask=None):
|
||
# bones = rollout_joints3d[:, :, 1:]
|
||
# parents = rollout_joints3d[:, :, SMPL_PARENTS[1:]]
|
||
# bone_lengths = torch.norm(bones - parents, dim=-1)
|
||
# loss = bone_lengths[:, 1:] - bone_lengths[:, :-1]
|
||
# if mask is not None:
|
||
# mask = mask.bool()
|
||
# mask = mask[:, 1:] & mask[:, :-1]
|
||
# loss = loss[mask]
|
||
# loss = 0.5 * torch.sum(loss**2)
|
||
# return loss
|
||
|
||
|
||
def kl_normal(qm, qv, pm, pv):
|
||
"""
|
||
Computes the elem-wise KL divergence between two normal distributions KL(q || p) and
|
||
sum over the last dimension
|
||
|
||
Args:
|
||
qm: tensor: (batch, dim): q mean
|
||
qv: tensor: (batch, dim): q variance
|
||
pm: tensor: (batch, dim): p mean
|
||
pv: tensor: (batch, dim): p variance
|
||
|
||
Return:
|
||
kl: tensor: (batch,): kl between each sample
|
||
"""
|
||
element_wise = 0.5 * (
|
||
torch.log(pv) - torch.log(qv) + qv / pv + (qm - pm).pow(2) / pv - 1
|
||
)
|
||
kl = element_wise.sum(-1)
|
||
return kl
|
||
|
||
|
||
def log_normal(x, m, v):
|
||
"""
|
||
Computes the elem-wise log probability of a Gaussian and then sum over the
|
||
last dim. Basically we're assuming all dims are batch dims except for the
|
||
last dim. Args:
|
||
x: tensor: (batch_1, batch_2, ..., batch_k, dim): Observation
|
||
m: tensor: (batch_1, batch_2, ..., batch_k, dim): Mean
|
||
v: tensor: (batch_1, batch_2, ..., batch_k, dim): Variance Return:
|
||
log_prob: tensor: (batch_1, batch_2, ..., batch_k): log probability of
|
||
each sample. Note that the summation dimension is not kept
|
||
"""
|
||
log_prob = (
|
||
-torch.log(torch.sqrt(v))
|
||
- math.log(math.sqrt(2 * math.pi))
|
||
- ((x - m) ** 2 / (2 * v))
|
||
)
|
||
log_prob = torch.sum(log_prob, dim=-1)
|
||
return log_prob
|
||
|
||
|
||
def apply_robust_weighting(
|
||
res, robust_loss_type="bisquare", robust_tuning_const=4.6851
|
||
):
|
||
"""
|
||
Returns robustly weighted squared residuals.
|
||
- res : torch.Tensor (B x N), take the MAD over each batch dimension independently.
|
||
"""
|
||
robust_choices = ["none", "bisquare"]
|
||
if robust_loss_type not in robust_choices:
|
||
print(
|
||
"Not a valid robust loss: %s. Please use %s"
|
||
% (robust_loss_type, str(robust_choices))
|
||
)
|
||
|
||
w = None
|
||
detach_res = (
|
||
res.clone().detach()
|
||
) # don't want gradients flowing through the weights to avoid degeneracy
|
||
if robust_loss_type == "none":
|
||
w = torch.ones_like(detach_res)
|
||
elif robust_loss_type == "bisquare":
|
||
w = bisquare_robust_weights(detach_res, tune_const=robust_tuning_const)
|
||
|
||
# apply weights to squared residuals
|
||
weighted_sqr_res = w * (res**2)
|
||
return weighted_sqr_res, w
|
||
|
||
|
||
def robust_std(res):
|
||
"""
|
||
Compute robust estimate of standarad deviation using median absolute deviation (MAD)
|
||
of the given residuals independently over each batch dimension.
|
||
|
||
- res : (B x N)
|
||
|
||
Returns:
|
||
- std : B x 1
|
||
"""
|
||
B = res.size(0)
|
||
med = torch.median(res, dim=-1)[0].reshape((B, 1))
|
||
abs_dev = torch.abs(res - med)
|
||
MAD = torch.median(abs_dev, dim=-1)[0].reshape((B, 1))
|
||
std = MAD / 0.67449
|
||
return std
|
||
|
||
|
||
def bisquare_robust_weights(res, tune_const=4.6851):
|
||
"""
|
||
Bisquare (Tukey) loss.
|
||
See https://www.mathworks.com/help/curvefit/least-squares-fitting.html
|
||
|
||
- residuals
|
||
"""
|
||
# print(res.size())
|
||
norm_res = res / (robust_std(res) * tune_const)
|
||
# NOTE: this should use absolute value, it's ok right now since only used for 3d point cloud residuals
|
||
# which are guaranteed positive, but generally this won't work)
|
||
outlier_mask = norm_res >= 1.0
|
||
|
||
# print(torch.sum(outlier_mask))
|
||
# print('Outlier frac: %f' % (float(torch.sum(outlier_mask)) / res.size(1)))
|
||
|
||
w = (1.0 - norm_res**2) ** 2
|
||
w[outlier_mask] = 0.0
|
||
|
||
return w
|
||
|
||
|
||
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)
|