import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from ..registry import LOSSES
[docs]@LOSSES.register_module()
class SmoothL1Loss(nn.Module):
"""SmoothL1Loss loss .
Args:
use_target_weight (bool): Option to use weighted MSE loss.
Different joint types may have different target weights.
loss_weight (float): Weight of the loss. Default: 1.0.
"""
def __init__(self, use_target_weight=False, loss_weight=1.):
super().__init__()
self.criterion = F.smooth_l1_loss
self.use_target_weight = use_target_weight
self.loss_weight = loss_weight
[docs] def forward(self, output, target, target_weight):
"""Forward function.
Note:
batch_size: N
num_keypoints: K
dimension of keypoints: D (D=2 or D=3)
Args:
output (torch.Tensor[N, K, D]): Output regression.
target (torch.Tensor[N, K, D]): Target regression.
target_weight (torch.Tensor[N, K, D]):
Weights across different joint types.
"""
if self.use_target_weight:
loss = self.criterion(output * target_weight,
target * target_weight)
else:
loss = self.criterion(output, target)
return loss * self.loss_weight
[docs]@LOSSES.register_module()
class WingLoss(nn.Module):
"""Wing Loss 'Wing Loss for Robust Facial Landmark Localisation with
Convolutional Neural Networks' Feng et al. CVPR'2018.
Args:
omega (float), epsilon (float) are hyper-parameters.
use_target_weight (bool): Option to use weighted MSE loss.
Different joint types may have different target weights.
loss_weight (float): Weight of the loss. Default: 1.0.
"""
def __init__(self,
omega=10.0,
epsilon=2.0,
use_target_weight=False,
loss_weight=1.):
super().__init__()
self.omega = omega
self.epsilon = epsilon
self.use_target_weight = use_target_weight
self.loss_weight = loss_weight
# constant that smoothly links the piecewise-defined linear
# and nonlinear parts
self.C = self.omega * (1.0 - math.log(1.0 + self.omega / self.epsilon))
[docs] def criterion(self, pred, target):
"""Criterion of wingloss.
Note:
batch_size: N
num_keypoints: K
dimension of keypoints: D (D=2 or D=3)
Args:
pred (torch.Tensor[N, K, D]): Output regression.
target (torch.Tensor[N, K, D]): Target regression.
"""
delta = (target - pred).abs()
losses = torch.where(
delta < self.omega,
self.omega * torch.log(1.0 + delta / self.epsilon), delta - self.C)
return torch.mean(torch.sum(losses, dim=[1, 2]), dim=0)
[docs] def forward(self, output, target, target_weight):
"""Forward function.
Note:
batch_size: N
num_keypoints: K
dimension of keypoints: D (D=2 or D=3)
Args:
output (torch.Tensor[N, K, D]): Output regression.
target (torch.Tensor[N, K, D]): Target regression.
target_weight (torch.Tensor[N, K, D]):
Weights across different joint types.
"""
if self.use_target_weight:
loss = self.criterion(output * target_weight,
target * target_weight)
else:
loss = self.criterion(output, target)
return loss * self.loss_weight
[docs]@LOSSES.register_module()
class MPJPELoss(nn.Module):
"""MPJPE (Mean Per Joint Position Error) loss.
Args:
use_target_weight (bool): Option to use weighted MSE loss.
Different joint types may have different target weights.
loss_weight (float): Weight of the loss. Default: 1.0.
"""
def __init__(self, use_target_weight=False, loss_weight=1.):
super().__init__()
self.use_target_weight = use_target_weight
self.loss_weight = loss_weight
[docs] def forward(self, output, target, target_weight):
"""Forward function.
Note:
batch_size: N
num_keypoints: K
dimension of keypoints: D (D=2 or D=3)
Args:
output (torch.Tensor[N, K, D]): Output regression.
target (torch.Tensor[N, K, D]): Target regression.
target_weight (torch.Tensor[N, K, D]):
Weights across different joint types.
"""
if self.use_target_weight:
loss = torch.mean(
torch.norm((output - target) * target_weight, dim=-1))
else:
loss = torch.mean(torch.norm(output - target, dim=-1))
return loss * self.loss_weight
[docs]@LOSSES.register_module()
class L1Loss(nn.Module):
"""L1Loss loss ."""
def __init__(self, use_target_weight=False, loss_weight=1.):
super().__init__()
self.criterion = F.l1_loss
self.use_target_weight = use_target_weight
self.loss_weight = loss_weight
[docs] def forward(self, output, target, target_weight):
"""Forward function.
Note:
batch_size: N
num_keypoints: K
Args:
output (torch.Tensor[N, K, 2]): Output regression.
target (torch.Tensor[N, K, 2]): Target regression.
target_weight (torch.Tensor[N, K, 2]):
Weights across different joint types.
"""
if self.use_target_weight:
loss = self.criterion(output * target_weight,
target * target_weight)
else:
loss = self.criterion(output, target)
return loss * self.loss_weight
[docs]@LOSSES.register_module()
class MSELoss(nn.Module):
"""MSE loss for coordinate regression."""
def __init__(self, use_target_weight=False, loss_weight=1.):
super().__init__()
self.criterion = F.mse_loss
self.use_target_weight = use_target_weight
self.loss_weight = loss_weight
[docs] def forward(self, output, target, target_weight):
"""Forward function.
Note:
batch_size: N
num_keypoints: K
Args:
output (torch.Tensor[N, K, 2]): Output regression.
target (torch.Tensor[N, K, 2]): Target regression.
target_weight (torch.Tensor[N, K, 2]):
Weights across different joint types.
"""
if self.use_target_weight:
loss = self.criterion(output * target_weight,
target * target_weight)
else:
loss = self.criterion(output, target)
return loss * self.loss_weight