Spaces:
Running
Running
File size: 6,492 Bytes
17f1f54 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | # coding: utf-8
import torch
import torch.nn as nn
from typing import Optional
from helpers import getSkeletalModelStructure, getSkeletalParentsDict
from pinn_losses import PINNLoss, PINNConfig
class Loss(nn.Module):
def __init__(self, cfg, target_pad=0.0):
super(Loss, self).__init__()
self.loss = cfg["training"]["loss"].lower()
self.bone_loss = cfg["training"]["bone_loss"].lower()
if self.loss == "l1":
self.criterion = nn.L1Loss()
elif self.loss == "mse":
self.criterion = nn.MSELoss()
else:
print("Loss not found - revert to default L1 loss")
self.criterion = nn.L1Loss()
if self.bone_loss == "l1":
self.criterion_bone = nn.L1Loss()
elif self.bone_loss == "mse":
self.criterion_bone = nn.MSELoss()
else:
print("Loss not found - revert to default MSE loss")
self.criterion_bone = nn.MSELoss()
model_cfg = cfg["model"]
training_cfg = cfg["training"]
self.target_pad = target_pad
self.loss_scale = model_cfg.get("loss_scale", 1.0)
# PINN Configuration
self.use_pinn = training_cfg.get("use_pinn", False)
self.lambda_pinn = float(training_cfg.get("lambda_pinn", 0.5))
if self.use_pinn:
# Create PINN configuration from training config with proper type conversion
pinn_cfg = PINNConfig(
lambda_bone=float(training_cfg.get("pinn_lambda_bone", 1.0)),
lambda_vel=float(training_cfg.get("pinn_lambda_vel", 0.1)),
lambda_acc=float(training_cfg.get("pinn_lambda_acc", 0.05)),
lambda_fk=float(training_cfg.get("pinn_lambda_fk", 0.5)),
eps=float(training_cfg.get("pinn_eps", 1e-8)),
dt=float(training_cfg.get("pinn_dt", 1.0)),
rest_from=str(training_cfg.get("pinn_rest_from", "first_valid")),
detach_rest=bool(training_cfg.get("pinn_detach_rest", True)),
use_huber=bool(training_cfg.get("pinn_use_huber", True)),
huber_delta=float(training_cfg.get("pinn_huber_delta", 1.0))
)
# Get parents dictionary from skeletal structure
parents_dict = getSkeletalParentsDict()
# Initialize PINN loss module
self.pinn_loss = PINNLoss(
parents=parents_dict,
num_joints=50,
cfg=pinn_cfg
)
print(f"PINN Loss initialized with {len(parents_dict)} parent-child pairs")
print(f" - lambda_pinn: {self.lambda_pinn}")
print(f" - lambda_bone: {pinn_cfg.lambda_bone}")
print(f" - lambda_vel: {pinn_cfg.lambda_vel}")
print(f" - lambda_acc: {pinn_cfg.lambda_acc}")
print(f" - lambda_fk: {pinn_cfg.lambda_fk}")
else:
self.pinn_loss = None
print("PINN Loss is disabled")
def forward(self, preds, targets, mask: Optional[torch.Tensor] = None):
"""
Compute loss with optional PINN regularization
Args:
preds: predicted skeleton (B, T, 150)
targets: target skeleton (B, T, 150)
mask: optional mask (B, T) for valid frames
"""
# Create loss mask from target padding
loss_mask = (targets != self.target_pad)
# Find the masked predictions and targets using loss mask
preds_masked = preds * loss_mask
targets_masked = targets * loss_mask
# Compute bone length and direction features
preds_masked_length, preds_masked_direct = get_length_direct(preds_masked)
targets_masked_length, targets_masked_direct = get_length_direct(targets_masked)
preds_masked_length = preds_masked_length * loss_mask[:, :, :50]
targets_masked_length = targets_masked_length * loss_mask[:, :, :50]
preds_masked_direct = preds_masked_direct * loss_mask[:, :, :150]
targets_masked_direct = targets_masked_direct * loss_mask[:, :, :150]
# Calculate base reconstruction loss
recon_loss = self.criterion(preds_masked, targets_masked) + \
0.1 * self.criterion_bone(preds_masked_direct, targets_masked_direct)
# Add PINN loss if enabled
if self.use_pinn and self.pinn_loss is not None:
# Create mask for PINN (B, T) - True for valid frames
if mask is None:
# Infer mask from targets: a frame is valid if it's not all padding
# Check if any coordinate in the frame is non-pad
pinn_mask = (targets[:, :, 0] != self.target_pad) # (B, T)
else:
pinn_mask = mask
# Compute PINN losses on predictions
pinn_out = self.pinn_loss(preds, pinn_mask)
# Total PINN loss
pinn_total = pinn_out["total"]
# Combined loss
loss = recon_loss + self.lambda_pinn * pinn_total
else:
loss = recon_loss
# Multiply loss by the loss scale
if self.loss_scale != 1.0:
loss = loss * self.loss_scale
return loss
def get_length_direct(trg):
"""
Compute bone lengths and unit directions from skeleton coordinates
Args:
trg: skeleton tensor (B, T, 150) - 50 joints x 3 coordinates
Returns:
lengths: (B, T, num_bones) - length of each bone
directs: (B, T, 3*num_bones) - unit direction vectors (3D) for each bone
"""
trg_reshaped = trg.view(trg.shape[0], trg.shape[1], 50, 3)
trg_list = trg_reshaped.split(1, dim=2)
trg_list_squeeze = [t.squeeze(dim=2) for t in trg_list]
skeletons = getSkeletalModelStructure()
length = []
direct = []
for skeleton in skeletons:
Skeleton_length = torch.norm(trg_list_squeeze[skeleton[0]]-trg_list_squeeze[skeleton[1]], p=2, dim=2, keepdim=True)
result_length = Skeleton_length
result_direct = (trg_list_squeeze[skeleton[0]]-trg_list_squeeze[skeleton[1]]) / (Skeleton_length+torch.finfo(Skeleton_length.dtype).tiny)
direct.append(result_direct)
length.append(result_length)
lengths = torch.stack(length, dim=-1).squeeze()
directs = torch.stack(direct, dim=2).view(trg.shape[0], trg.shape[1], -1)
return lengths, directs
|