models/detector.py

"""
SCRFD Full Detector — Backbone + Neck + Head + Loss + Post-processing.

This is the main model class that ties together all components and provides:
1. Training forward: returns losses dict
2. Inference forward: returns detections (boxes, scores, landmarks)
3. ONNX-exportable inference path

Model configurations (WiderFace Hard val / GFLOPs / FPS @VGA on V100):
- SCRFD-34GF:  85.2% / 34 GF / ~80 FPS   (flagship quality)
- SCRFD-10GF:  83.1% / 10 GF / ~140 FPS  (balanced)
- SCRFD-2.5GF: 77.9% / 2.5 GF / ~400 FPS (real-time)
- SCRFD-0.5GF: 68.5% / 0.5 GF / ~1000 FPS (mobile/edge)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple, Dict, Optional
import math

from .backbone import SCRFDBackbone, build_backbone
from .neck import PAFPN, build_neck
from .head import SCRFDHead, build_head
from .anchor import AnchorGenerator, ATSSAssigner
from .losses import GFocalLoss, DIoULoss, FocalLoss, LandmarkLoss


class SCRFD(nn.Module):
    """
    Sample and Computation Redistribution Face Detector.

    Complete pipeline: backbone → PAFPN → shared head → anchors → losses/NMS
    """

    def __init__(self,
                 backbone: SCRFDBackbone,
                 neck: PAFPN,
                 head: SCRFDHead,
                 anchor_generator: AnchorGenerator,
                 assigner: ATSSAssigner,
                 strides: List[int] = [8, 16, 32],
                 score_threshold: float = 0.3,
                 nms_threshold: float = 0.4,
                 max_detections: int = 750,
                 use_gfl: bool = True,
                 cls_weight: float = 1.0,
                 reg_weight: float = 2.0,
                 lmk_weight: float = 0.1):
        super().__init__()
        self.backbone = backbone
        self.neck = neck
        self.head = head
        self.anchor_gen = anchor_generator
        self.assigner = assigner
        self.strides = strides
        self.score_threshold = score_threshold
        self.nms_threshold = nms_threshold
        self.max_detections = max_detections
        self.use_gfl = use_gfl

        # Loss functions
        self.cls_loss_fn = GFocalLoss(beta=2.0) if use_gfl else FocalLoss()
        self.reg_loss_fn = DIoULoss()
        self.lmk_loss_fn = LandmarkLoss() if head.use_landmarks else None

        # Loss weights
        self.cls_weight = cls_weight
        self.reg_weight = reg_weight
        self.lmk_weight = lmk_weight

    def forward(self, images: torch.Tensor,
                targets: Optional[List[Dict]] = None) -> Dict:
        """
        Args:
            images: [B, 3, H, W] batch of images (normalized)
            targets: List of dicts with keys:
                'boxes': [M, 4] face boxes (x1, y1, x2, y2)
                'labels': [M] labels (all 1)
                'landmarks': [M, 10] optional landmarks
                When None, runs inference.

        Returns:
            Training: dict of losses
            Inference: list of dicts with 'boxes', 'scores', 'landmarks'
        """
        # Feature extraction
        features = self.backbone(images)
        features = self.neck(features)
        head_out = self.head(features)

        # Generate anchors
        feat_sizes = [(f.shape[2], f.shape[3]) for f in features]
        anchors_per_level = self.anchor_gen.grid_anchors(feat_sizes, images.device)
        num_anchors_per_level = [a.shape[0] for a in anchors_per_level]

        if targets is not None:
            return self._compute_loss(head_out, anchors_per_level,
                                      num_anchors_per_level, targets, images.shape)
        else:
            return self._inference(head_out, anchors_per_level, images.shape)

    def _compute_loss(self, head_out: Dict, anchors_per_level: List[torch.Tensor],
                      num_per_level: List[int], targets: List[Dict],
                      img_shape: Tuple) -> Dict:
        """Compute training losses."""
        device = anchors_per_level[0].device
        batch_size = len(targets)

        # Flatten predictions across levels
        all_cls = []
        all_reg = []
        all_lmk = []
        for i in range(len(self.strides)):
            B, _, H, W = head_out['cls_scores'][i].shape
            cls = head_out['cls_scores'][i].permute(0, 2, 3, 1).reshape(B, -1, 1)
            reg = head_out['bbox_preds'][i].permute(0, 2, 3, 1).reshape(B, -1, 4)
            all_cls.append(cls)
            all_reg.append(reg)
            if self.head.use_landmarks and 'lmk_preds' in head_out:
                lmk = head_out['lmk_preds'][i].permute(0, 2, 3, 1).reshape(B, -1, 10)
                all_lmk.append(lmk)

        all_cls = torch.cat(all_cls, dim=1)  # [B, N, 1]
        all_reg = torch.cat(all_reg, dim=1)  # [B, N, 4]
        all_anchors = torch.cat(anchors_per_level, dim=0)  # [N, 4]

        has_lmk = len(all_lmk) > 0
        if has_lmk:
            all_lmk = torch.cat(all_lmk, dim=1)

        total_cls_loss = torch.tensor(0.0, device=device)
        total_reg_loss = torch.tensor(0.0, device=device)
        total_lmk_loss = torch.tensor(0.0, device=device)
        num_pos = 0

        for b in range(batch_size):
            gt_boxes = targets[b]['boxes']
            gt_labels = targets[b].get('labels',
                                       torch.ones(gt_boxes.shape[0], dtype=torch.long, device=device))

            # ATSS matching
            assigned_labels, assigned_gt_inds = self.assigner.assign(
                all_anchors, gt_boxes, gt_labels, num_per_level
            )

            pos_mask = assigned_labels > 0
            num_pos += pos_mask.sum().item()

            # Classification loss (all anchors)
            if self.use_gfl:
                # GFL: positive target = IoU, negative target = 0
                cls_targets = torch.zeros(all_anchors.shape[0], device=device)
                if pos_mask.any():
                    pos_anchors = all_anchors[pos_mask]
                    pos_gt = gt_boxes[assigned_gt_inds[pos_mask]]
                    pos_ious = self._compute_iou_single(pos_anchors, pos_gt)
                    cls_targets[pos_mask] = pos_ious
                total_cls_loss += self.cls_loss_fn(
                    all_cls[b].squeeze(-1), cls_targets
                )
            else:
                total_cls_loss += self.cls_loss_fn(
                    all_cls[b].squeeze(-1), (assigned_labels > 0).float()
                )

            # Box regression loss (positive anchors only)
            if pos_mask.any():
                pos_reg = all_reg[b][pos_mask]
                pos_anchors = all_anchors[pos_mask]
                pos_gt = gt_boxes[assigned_gt_inds[pos_mask]]

                # Decode predictions to absolute boxes
                pred_boxes = self._decode_boxes(pos_anchors, pos_reg)
                total_reg_loss += self.reg_loss_fn(pred_boxes, pos_gt)

                # Landmark loss
                if self.head.use_landmarks and 'landmarks' in targets[b] and has_lmk:
                    gt_lmk = targets[b]['landmarks']
                    pos_lmk_pred = all_lmk[b][pos_mask]
                    pos_lmk_gt = gt_lmk[assigned_gt_inds[pos_mask]]
                    # Decode landmarks relative to anchors
                    pred_lmk = self._decode_landmarks(pos_anchors, pos_lmk_pred)
                    total_lmk_loss += self.lmk_loss_fn(pred_lmk, pos_lmk_gt)

        num_pos = max(num_pos, 1)
        losses = {
            'cls_loss': self.cls_weight * total_cls_loss / batch_size,
            'reg_loss': self.reg_weight * total_reg_loss / batch_size,
        }
        if self.head.use_landmarks:
            losses['lmk_loss'] = self.lmk_weight * total_lmk_loss / batch_size

        losses['total_loss'] = sum(losses.values())
        losses['num_pos'] = torch.tensor(num_pos, dtype=torch.float, device=device)
        return losses

    def _inference(self, head_out: Dict, anchors_per_level: List[torch.Tensor],
                   img_shape: Tuple) -> List[Dict]:
        """Run inference with NMS."""
        batch_size = head_out['cls_scores'][0].shape[0]
        device = head_out['cls_scores'][0].device

        results = []
        for b in range(batch_size):
            all_boxes = []
            all_scores = []
            all_lmk = []

            for i in range(len(self.strides)):
                cls = head_out['cls_scores'][i][b].permute(1, 2, 0).reshape(-1, 1).sigmoid()
                reg = head_out['bbox_preds'][i][b].permute(1, 2, 0).reshape(-1, 4)
                anchors = anchors_per_level[i]

                # Filter by score threshold
                scores = cls.squeeze(-1)
                keep = scores > self.score_threshold
                if keep.sum() == 0:
                    continue

                scores = scores[keep]
                reg = reg[keep]
                anc = anchors[keep]

                # Decode boxes
                boxes = self._decode_boxes(anc, reg)

                # Clamp to image boundaries
                boxes[:, 0].clamp_(min=0)
                boxes[:, 1].clamp_(min=0)
                boxes[:, 2].clamp_(max=img_shape[3])
                boxes[:, 3].clamp_(max=img_shape[2])

                all_boxes.append(boxes)
                all_scores.append(scores)

                if self.head.use_landmarks and 'lmk_preds' in head_out:
                    lmk = head_out['lmk_preds'][i][b].permute(1, 2, 0).reshape(-1, 10)[keep]
                    lmk_decoded = self._decode_landmarks(anc, lmk)
                    all_lmk.append(lmk_decoded)

            if not all_boxes:
                results.append({
                    'boxes': torch.empty(0, 4, device=device),
                    'scores': torch.empty(0, device=device),
                })
                continue

            all_boxes = torch.cat(all_boxes, dim=0)
            all_scores = torch.cat(all_scores, dim=0)

            # NMS
            keep = self._nms(all_boxes, all_scores, self.nms_threshold)
            keep = keep[:self.max_detections]

            result = {
                'boxes': all_boxes[keep],
                'scores': all_scores[keep],
            }
            if all_lmk:
                all_lmk = torch.cat(all_lmk, dim=0)
                result['landmarks'] = all_lmk[keep]
            results.append(result)

        return results

    def _decode_boxes(self, anchors: torch.Tensor, pred: torch.Tensor) -> torch.Tensor:
        """Decode box predictions relative to anchors (distance-based)."""
        anchor_cx = (anchors[:, 0] + anchors[:, 2]) / 2
        anchor_cy = (anchors[:, 1] + anchors[:, 3]) / 2
        anchor_w = anchors[:, 2] - anchors[:, 0]
        anchor_h = anchors[:, 3] - anchors[:, 1]

        x1 = anchor_cx - pred[:, 0] * anchor_w
        y1 = anchor_cy - pred[:, 1] * anchor_h
        x2 = anchor_cx + pred[:, 2] * anchor_w
        y2 = anchor_cy + pred[:, 3] * anchor_h

        return torch.stack([x1, y1, x2, y2], dim=1)

    def _decode_landmarks(self, anchors: torch.Tensor, pred: torch.Tensor) -> torch.Tensor:
        """Decode landmark predictions relative to anchors."""
        anchor_cx = (anchors[:, 0] + anchors[:, 2]) / 2
        anchor_cy = (anchors[:, 1] + anchors[:, 3]) / 2
        anchor_w = anchors[:, 2] - anchors[:, 0]
        anchor_h = anchors[:, 3] - anchors[:, 1]

        decoded = pred.clone()
        for i in range(5):
            decoded[:, i*2] = anchor_cx + pred[:, i*2] * anchor_w
            decoded[:, i*2+1] = anchor_cy + pred[:, i*2+1] * anchor_h
        return decoded

    @staticmethod
    def _compute_iou_single(boxes1: torch.Tensor, boxes2: torch.Tensor) -> torch.Tensor:
        """Compute elementwise IoU between paired boxes. [N,4] × [N,4] → [N]"""
        inter_x1 = torch.max(boxes1[:, 0], boxes2[:, 0])
        inter_y1 = torch.max(boxes1[:, 1], boxes2[:, 1])
        inter_x2 = torch.min(boxes1[:, 2], boxes2[:, 2])
        inter_y2 = torch.min(boxes1[:, 3], boxes2[:, 3])
        inter = (inter_x2 - inter_x1).clamp(min=0) * (inter_y2 - inter_y1).clamp(min=0)

        area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
        area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
        union = area1 + area2 - inter
        return inter / (union + 1e-6)

    @staticmethod
    def _nms(boxes: torch.Tensor, scores: torch.Tensor,
             threshold: float) -> torch.Tensor:
        """Non-Maximum Suppression. Returns kept indices."""
        if boxes.shape[0] == 0:
            return torch.empty(0, dtype=torch.long, device=boxes.device)

        # Use torchvision NMS if available, else pure PyTorch
        try:
            from torchvision.ops import nms
            return nms(boxes, scores, threshold)
        except ImportError:
            pass

        # Pure PyTorch NMS fallback
        x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
        areas = (x2 - x1) * (y2 - y1)
        order = scores.argsort(descending=True)
        keep = []

        while order.numel() > 0:
            i = order[0].item()
            keep.append(i)
            if order.numel() == 1:
                break

            xx1 = torch.max(x1[i], x1[order[1:]])
            yy1 = torch.max(y1[i], y1[order[1:]])
            xx2 = torch.min(x2[i], x2[order[1:]])
            yy2 = torch.min(y2[i], y2[order[1:]])
            inter = (xx2 - xx1).clamp(min=0) * (yy2 - yy1).clamp(min=0)
            iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-6)
            mask = iou <= threshold
            order = order[1:][mask]

        return torch.tensor(keep, dtype=torch.long, device=boxes.device)


# ──────────────────────── Model Builder ────────────────────────

MODEL_CONFIGS = {
    'scrfd_34g': {
        'backbone': 'scrfd_34g',
        'neck_out': 64,
        'head_feat': 64,
        'head_convs': 3,
    },
    'scrfd_10g': {
        'backbone': 'scrfd_10g',
        'neck_out': 56,
        'head_feat': 56,
        'head_convs': 2,
    },
    'scrfd_2.5g': {
        'backbone': 'scrfd_2.5g',
        'neck_out': 40,
        'head_feat': 40,
        'head_convs': 2,
    },
    'scrfd_0.5g': {
        'backbone': 'scrfd_0.5g',
        'neck_out': 16,
        'head_feat': 16,
        'head_convs': 2,
    },
}


def build_detector(name: str, use_landmarks: bool = False,
                   score_threshold: float = 0.3,
                   nms_threshold: float = 0.4,
                   **kwargs) -> SCRFD:
    """
    Build a complete SCRFD detector by name.

    Args:
        name: Model name ('scrfd_34g', 'scrfd_10g', 'scrfd_2.5g', 'scrfd_0.5g')
        use_landmarks: Enable 5-point landmark prediction
        score_threshold: Detection confidence threshold
        nms_threshold: NMS IoU threshold

    Returns:
        Complete SCRFD detector ready for training or inference
    """
    if name not in MODEL_CONFIGS:
        raise ValueError(f"Unknown model: {name}. Options: {list(MODEL_CONFIGS.keys())}")

    cfg = MODEL_CONFIGS[name]

    backbone = build_backbone(cfg['backbone'])
    neck = PAFPN(backbone.out_channels, out_channels=cfg['neck_out'])
    head = SCRFDHead(
        in_channels=cfg['neck_out'],
        feat_channels=cfg['head_feat'],
        stacked_convs=cfg['head_convs'],
        use_landmarks=use_landmarks,
    )
    anchor_gen = AnchorGenerator()
    assigner = ATSSAssigner(topk=9)

    model = SCRFD(
        backbone=backbone,
        neck=neck,
        head=head,
        anchor_generator=anchor_gen,
        assigner=assigner,
        score_threshold=score_threshold,
        nms_threshold=nms_threshold,
        **kwargs,
    )

    return model