models/backbone.py

"""
SCRFD Backbones — NAS-searched ResNet-style with computation redistribution.

Key insight from paper: Standard classification backbones over-invest compute in
C5 features (stride 32), which are useless for tiny face detection. SCRFD
redistributes compute toward earlier stages (C2/C3) for stride-8 feature quality.

Configurations (from paper Table 3):
- SCRFD-34GF: stages=[3,12,28,3], widths=[56,88,248,304], groups=[1,1,1,1]
- SCRFD-10GF: stages=[3,10,16,3], widths=[36,64,144,224], groups=[1,1,1,1]
- SCRFD-2.5GF: stages=[2,4,4,3],  widths=[24,48,96,160],  groups=[1,1,1,1]
- SCRFD-0.5GF: stages=[2,2,4,2],  widths=[16,32,64,128],  groups=[1,1,1,1]
"""

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


class ConvBNReLU(nn.Module):
    """Conv + BatchNorm + ReLU building block."""

    def __init__(self, in_ch: int, out_ch: int, kernel: int = 3,
                 stride: int = 1, padding: int = 1, groups: int = 1,
                 use_relu: bool = True):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, kernel, stride, padding,
                              groups=groups, bias=False)
        self.bn = nn.BatchNorm2d(out_ch)
        self.relu = nn.ReLU(inplace=True) if use_relu else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.relu(self.bn(self.conv(x)))


class BasicBlock(nn.Module):
    """ResNet BasicBlock with optional group convolution."""
    expansion = 1

    def __init__(self, in_ch: int, out_ch: int, stride: int = 1,
                 groups: int = 1, downsample: Optional[nn.Module] = None):
        super().__init__()
        self.conv1 = ConvBNReLU(in_ch, out_ch, 3, stride, 1, groups)
        self.conv2 = ConvBNReLU(out_ch, out_ch, 3, 1, 1, groups, use_relu=False)
        self.downsample = downsample
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = x
        out = self.conv1(x)
        out = self.conv2(out)
        if self.downsample is not None:
            identity = self.downsample(x)
        out += identity
        return self.relu(out)


class BottleneckBlock(nn.Module):
    """ResNet Bottleneck with optional group convolution."""
    expansion = 4

    def __init__(self, in_ch: int, out_ch: int, stride: int = 1,
                 groups: int = 1, downsample: Optional[nn.Module] = None):
        super().__init__()
        mid_ch = out_ch  # bottleneck width
        self.conv1 = ConvBNReLU(in_ch, mid_ch, 1, 1, 0)
        self.conv2 = ConvBNReLU(mid_ch, mid_ch, 3, stride, 1, groups)
        self.conv3 = ConvBNReLU(mid_ch, out_ch * self.expansion, 1, 1, 0, use_relu=False)
        self.downsample = downsample
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = x
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        if self.downsample is not None:
            identity = self.downsample(x)
        out += identity
        return self.relu(out)


class SCRFDBackbone(nn.Module):
    """
    SCRFD backbone with NAS-searched stage depths and widths.

    For SCRFD, we use BasicBlock (expansion=1) since the searched widths
    already account for channel capacity — no need for bottleneck expansion.

    Returns feature maps at strides [8, 16, 32] (C3, C4, C5).
    """

    def __init__(self, stages: List[int], widths: List[int],
                 groups: List[int] = None, in_channels: int = 3,
                 block_type: str = 'basic'):
        super().__init__()
        assert len(stages) == 4 and len(widths) == 4

        if groups is None:
            groups = [1, 1, 1, 1]

        Block = BasicBlock if block_type == 'basic' else BottleneckBlock

        # Stem: stride 2 conv + stride 2 maxpool → effective stride 4
        self.stem = nn.Sequential(
            ConvBNReLU(in_channels, widths[0], 3, 2, 1),
            ConvBNReLU(widths[0], widths[0], 3, 1, 1),
            nn.MaxPool2d(3, 2, 1),
        )

        # Stage 1: stride 1 (output stride = 4)
        self.layer1 = self._make_layer(Block, widths[0], widths[0], stages[0],
                                       stride=1, groups=groups[0])
        # Stage 2: stride 2 (output stride = 8) → C3
        self.layer2 = self._make_layer(Block, widths[0], widths[1], stages[1],
                                       stride=2, groups=groups[1])
        # Stage 3: stride 2 (output stride = 16) → C4
        self.layer3 = self._make_layer(Block, widths[1], widths[2], stages[2],
                                       stride=2, groups=groups[2])
        # Stage 4: stride 2 (output stride = 32) → C5
        self.layer4 = self._make_layer(Block, widths[2], widths[3], stages[3],
                                       stride=2, groups=groups[3])

        self.out_channels = [widths[1], widths[2], widths[3]]
        self._init_weights()

    def _make_layer(self, block, in_ch: int, out_ch: int, num_blocks: int,
                    stride: int = 1, groups: int = 1) -> nn.Sequential:
        downsample = None
        if stride != 1 or in_ch != out_ch * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(in_ch, out_ch * block.expansion, 1, stride, bias=False),
                nn.BatchNorm2d(out_ch * block.expansion),
            )

        layers = [block(in_ch, out_ch, stride, groups, downsample)]
        in_ch = out_ch * block.expansion
        for _ in range(1, num_blocks):
            layers.append(block(in_ch, out_ch, 1, groups))
        return nn.Sequential(*layers)

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        x = self.stem(x)
        c2 = self.layer1(x)     # stride 4
        c3 = self.layer2(c2)    # stride 8
        c4 = self.layer3(c3)    # stride 16
        c5 = self.layer4(c4)    # stride 32
        return c3, c4, c5


# ──────────────────────── Configuration presets ────────────────────────

BACKBONE_CONFIGS = {
    'scrfd_34g': dict(stages=[3, 12, 28, 3], widths=[56, 88, 248, 304]),
    'scrfd_10g': dict(stages=[3, 10, 16, 3], widths=[36, 64, 144, 224]),
    'scrfd_2.5g': dict(stages=[2, 4, 4, 3], widths=[24, 48, 96, 160]),
    'scrfd_0.5g': dict(stages=[2, 2, 4, 2], widths=[16, 32, 64, 128]),
    # ResNet variants for comparison
    'resnet50': dict(stages=[3, 4, 6, 3], widths=[64, 128, 256, 512], block_type='bottleneck'),
    'resnet18': dict(stages=[2, 2, 2, 2], widths=[64, 128, 256, 512]),
}


def build_backbone(name: str, **kwargs) -> SCRFDBackbone:
    """Build a backbone by name."""
    if name not in BACKBONE_CONFIGS:
        raise ValueError(f"Unknown backbone: {name}. Options: {list(BACKBONE_CONFIGS.keys())}")
    cfg = {**BACKBONE_CONFIGS[name], **kwargs}
    return SCRFDBackbone(**cfg)