Upload models/losses.py with huggingface_hub

models/losses.py

@@ -0,0 +1,197 @@
+"""
+Loss functions for SCRFD face detection.
+
+SCRFD uses:
+1. Generalized Focal Loss (GFL/QFL) for classification — jointly represents
+   classification score and localization quality in a single prediction.
+2. DIoU Loss for bounding box regression — better gradient signal for
+   non-overlapping boxes and directly minimizes distance between box centers.
+
+References:
+- GFL: "Generalized Focal Loss" (Li et al., 2020)
+- DIoU: "Distance-IoU Loss" (Zheng et al., 2020)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+
+
+class GFocalLoss(nn.Module):
+    """
+    Quality Focal Loss (QFL) — Generalized Focal Loss for classification.
+
+    Instead of binary {0,1} targets, QFL uses continuous quality scores
+    [0, 1] where the target is the IoU between predicted and GT boxes.
+    This jointly trains classification confidence and localization quality.
+
+    Loss = -|y - σ|^β * ((1-y)log(1-σ) + y*log(σ))
+
+    where y ∈ [0,1] is quality target, σ is predicted score, β is focusing param.
+    """
+
+    def __init__(self, beta: float = 2.0, reduction: str = 'mean'):
+        super().__init__()
+        self.beta = beta
+        self.reduction = reduction
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor,
+                weight: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            pred: [N] predicted scores (logits)
+            target: [N] quality targets in [0, 1]
+            weight: [N] optional sample weights
+        """
+        pred_sigmoid = pred.sigmoid()
+        scale_factor = (pred_sigmoid - target).abs().pow(self.beta)
+
+        # Binary cross-entropy with continuous targets
+        bce = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
+        loss = scale_factor * bce
+
+        if weight is not None:
+            loss = loss * weight
+
+        if self.reduction == 'mean':
+            return loss.sum() / max(weight.sum() if weight is not None else target.gt(0).sum(), 1)
+        elif self.reduction == 'sum':
+            return loss.sum()
+        return loss
+
+
+class FocalLoss(nn.Module):
+    """
+    Standard Focal Loss for binary classification.
+
+    FL(p) = -α * (1-p)^γ * log(p)    for positive
+          = -(1-α) * p^γ * log(1-p)  for negative
+
+    Used as fallback when QFL is not appropriate.
+    """
+
+    def __init__(self, alpha: float = 0.25, gamma: float = 2.0,
+                 reduction: str = 'mean'):
+        super().__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.reduction = reduction
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        pred_sigmoid = pred.sigmoid()
+        target = target.float()
+
+        # Focal weights
+        pt = pred_sigmoid * target + (1 - pred_sigmoid) * (1 - target)
+        focal_weight = (1 - pt).pow(self.gamma)
+        alpha_weight = self.alpha * target + (1 - self.alpha) * (1 - target)
+
+        bce = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
+        loss = alpha_weight * focal_weight * bce
+
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        return loss
+
+
+class DIoULoss(nn.Module):
+    """
+    Distance-IoU Loss for bounding box regression.
+
+    DIoU = IoU - (ρ²(b, b_gt) / c²)
+
+    where ρ is Euclidean distance between box centers and c is diagonal
+    length of the smallest enclosing box. This provides better gradients
+    for non-overlapping boxes (common with tiny faces) and directly
+    optimizes center alignment.
+
+    Loss = 1 - DIoU ∈ [0, 2]
+    """
+
+    def __init__(self, reduction: str = 'mean'):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor,
+                weight: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            pred: [N, 4] predicted boxes (x1, y1, x2, y2)
+            target: [N, 4] target boxes (x1, y1, x2, y2)
+            weight: [N] optional per-box weights
+        """
+        # Intersection
+        inter_x1 = torch.max(pred[:, 0], target[:, 0])
+        inter_y1 = torch.max(pred[:, 1], target[:, 1])
+        inter_x2 = torch.min(pred[:, 2], target[:, 2])
+        inter_y2 = torch.min(pred[:, 3], target[:, 3])
+        inter = (inter_x2 - inter_x1).clamp(min=0) * (inter_y2 - inter_y1).clamp(min=0)
+
+        # Union
+        area_pred = (pred[:, 2] - pred[:, 0]) * (pred[:, 3] - pred[:, 1])
+        area_target = (target[:, 2] - target[:, 0]) * (target[:, 3] - target[:, 1])
+        union = area_pred + area_target - inter
+
+        iou = inter / (union + 1e-6)
+
+        # Center distance
+        pred_cx = (pred[:, 0] + pred[:, 2]) / 2
+        pred_cy = (pred[:, 1] + pred[:, 3]) / 2
+        target_cx = (target[:, 0] + target[:, 2]) / 2
+        target_cy = (target[:, 1] + target[:, 3]) / 2
+        center_dist_sq = (pred_cx - target_cx).pow(2) + (pred_cy - target_cy).pow(2)
+
+        # Smallest enclosing box diagonal
+        enclose_x1 = torch.min(pred[:, 0], target[:, 0])
+        enclose_y1 = torch.min(pred[:, 1], target[:, 1])
+        enclose_x2 = torch.max(pred[:, 2], target[:, 2])
+        enclose_y2 = torch.max(pred[:, 3], target[:, 3])
+        enclose_diag_sq = (enclose_x2 - enclose_x1).pow(2) + (enclose_y2 - enclose_y1).pow(2)
+
+        diou = iou - center_dist_sq / (enclose_diag_sq + 1e-6)
+        loss = 1 - diou
+
+        if weight is not None:
+            loss = loss * weight
+
+        if self.reduction == 'mean':
+            return loss.sum() / max(weight.sum() if weight is not None else loss.shape[0], 1)
+        elif self.reduction == 'sum':
+            return loss.sum()
+        return loss
+
+
+class LandmarkLoss(nn.Module):
+    """
+    Smooth L1 loss for facial landmark regression (optional multi-task head).
+
+    Used when landmark annotations are available (e.g., RetinaFace 5-point
+    landmarks on WIDER FACE). Auxiliary landmark supervision improves
+    detection AP by ~1% (RetinaFace paper finding).
+    """
+
+    def __init__(self, beta: float = 1.0, reduction: str = 'mean'):
+        super().__init__()
+        self.beta = beta
+        self.reduction = reduction
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor,
+                weight: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            pred: [N, 10] predicted landmarks (5 points × 2 coords)
+            target: [N, 10] target landmarks
+            weight: [N] optional mask for visible landmarks
+        """
+        loss = F.smooth_l1_loss(pred, target, beta=self.beta, reduction='none')
+        loss = loss.sum(dim=1)  # Sum over 10 coords per face
+
+        if weight is not None:
+            loss = loss * weight
+
+        if self.reduction == 'mean':
+            return loss.sum() / max(weight.sum() if weight is not None else loss.shape[0], 1)
+        return loss.sum()