Merge pull request #77 from kaseris/losses/angle-loss

Losses/angle loss

src/skelcast/losses/__init__.py

@@ -2,4 +2,5 @@
 
 LOSSES = Registry()
 
-from .logloss import LogLoss
+from .logloss import LogLoss
+from .euler_angle_loss import EulerAngleLoss

src/skelcast/losses/euler_angle_loss.py

@@ -0,0 +1,29 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+from skelcast.data.human36m.quaternion import qeuler
+from skelcast.losses import LOSSES
+
+
+@LOSSES.register_module()
+class EulerAngleLoss(nn.Module):
+    def __init__(self, order="xyz", reduction="mean"):
+        super(EulerAngleLoss, self).__init__()
+        self._order = order
+        self._reduction = reduction
+
+    def forward(self, predictions: torch.Tensor, targets: torch.Tensor):
+        # Check the shape of predictions and targets
+        assert (
+            predictions.shape == targets.shape
+        ), f"Predictions and targets must have the same shape."
+        assert (
+            predictions.shape[-1] == 3
+        ), f"Predictions and targets must have 3 channels in the last dimension."
+
+        predicted_euler = qeuler(predictions, self._order, epsilon=1e-6)
+        angle_distance = (
+            torch.remainder(predicted_euler - targets + np.pi, 2 * np.pi) - np.pi
+        )
+        return torch.mean(torch.abs(angle_distance))