File size: 2,751 Bytes
8e5456b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Reconstruction losses for DWPose sign-language motion.
Differs from T2M-GPT's ReConsLoss in two ways that matter for this data:
* masked -- undetected keypoints (low DWPose score / NaN) carry no gradient;
* keypoint-weighted -- hands dominate meaning in sign language, so they are
weighted above body, and face is down-weighted.
The velocity term is computed on the same weighted/masked coordinates, which for
a raw-coordinate representation is just the first temporal difference.
"""
import torch
import torch.nn as nn
from dataset.dataset_vsl import kp_weights
def _base_loss(name):
if name == "l1":
return lambda a, b: (a - b).abs()
if name == "l2":
return lambda a, b: (a - b) ** 2
if name == "l1_smooth":
return lambda a, b: nn.functional.smooth_l1_loss(a, b, reduction="none", beta=0.1)
raise ValueError(f"unknown recons loss {name}")
class VSLReConsLoss(nn.Module):
def __init__(self, recons_loss="l1_smooth", w_body=1.0, w_face=0.5, w_hand=3.0,
layout=None, w_finger=None, w_fingertip=None):
super().__init__()
self.fn = _base_loss(recons_loss)
w = torch.from_numpy(kp_weights(w_body, w_face, w_hand, layout=layout,
finger=w_finger, fingertip=w_fingertip))
self.register_buffer("w", w.view(1, 1, -1)) # [1,1,dim]
def _reduce(self, err, mask):
wm = mask * self.w
return (err * wm).sum() / wm.sum().clamp(min=1.0)
def forward(self, pred, gt, mask):
"""pred/gt/mask: [B,T,256]. mask is 1 where the keypoint is usable."""
return self._reduce(self.fn(pred, gt), mask)
def forward_vel(self, pred, gt, mask):
dp = pred[:, 1:] - pred[:, :-1]
dg = gt[:, 1:] - gt[:, :-1]
m = mask[:, 1:] * mask[:, :-1] # a velocity is valid only if both frames are
return self._reduce(self.fn(dp, dg), m)
@torch.no_grad()
def mpjpe_groups(pred_xy, gt_xy, valid, groups=None):
"""Mean per-joint position error in raw frame-normalized units, per group.
pred_xy / gt_xy: [B,T,2*n_kpts] un-normalized coordinates. valid: [B,T,n_kpts].
`groups` comes from Layout.metric_groups(); defaults to the full 128 layout.
Returns a dict of scalars (all / body / face / hands).
"""
B, T, D = pred_xy.shape
nk = D // 2
p = pred_xy.view(B, T, nk, 2)
g = gt_xy.view(B, T, nk, 2)
d = torch.linalg.norm(p - g, dim=-1) # [B,T,nk]
out = {}
if groups is None:
groups = {"all": (0, 128), "body": (0, 18), "face": (18, 86), "hands": (86, 128)}
for name, (a, b) in groups.items():
m = valid[..., a:b]
out[name] = float((d[..., a:b] * m).sum() / m.sum().clamp(min=1.0))
return out
|