| """The face networks in plain torch: SCRFD-10G (InsightFace's det_10g detector) and ArcFace's |
| IResNet (w600k_r50, glintr100), with batch norms folded into the convolutions as InsightFace's |
| exported files have them.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
|
|
| |
|
|
| class _Block(nn.Module): |
| """ResNet-D basic block: two 3x3 convs and a shortcut (average pool then 1x1 conv when it |
| downsamples).""" |
|
|
| def __init__(self, cin: int, cout: int, stride: int): |
| super().__init__() |
| self.conv1 = nn.Conv2d(cin, cout, 3, stride, 1) |
| self.conv2 = nn.Conv2d(cout, cout, 3, 1, 1) |
| self.down = nn.Conv2d(cin, cout, 1) if stride != 1 or cin != cout else None |
| self.stride = stride |
|
|
| def forward(self, x): |
| out = self.conv2(F.relu(self.conv1(x))) |
| if self.down is not None: |
| if self.stride != 1: |
| x = F.avg_pool2d(x, 2, 2, ceil_mode=True, count_include_pad=False) |
| x = self.down(x) |
| return F.relu(out + x) |
|
|
|
|
| class _Head(nn.Module): |
| """One stride's head: three 3x3 conv + relu, then scores (2 anchors), boxes and 5 keypoints.""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.convs = nn.ModuleList([nn.Conv2d(56, 80, 3, 1, 1), nn.Conv2d(80, 80, 3, 1, 1), nn.Conv2d(80, 80, 3, 1, 1)]) |
| self.cls = nn.Conv2d(80, 2, 3, 1, 1) |
| self.reg = nn.Conv2d(80, 8, 3, 1, 1) |
| self.kps = nn.Conv2d(80, 20, 3, 1, 1) |
| self.register_buffer("scale", torch.ones(())) |
|
|
| def forward(self, x): |
| for c in self.convs: |
| x = F.relu(c(x)) |
| |
| flat = lambda t, k: t.permute(2, 3, 0, 1).reshape(-1, k) |
| return flat(self.cls(x), 1).sigmoid(), flat(self.reg(x) * self.scale, 4), flat(self.kps(x), 10) |
|
|
|
|
| class SCRFD(nn.Module): |
| """(1, 3, H, W) RGB in (x - 127.5) / 128 -> per stride 8, 16, 32: scores, box and keypoint |
| distances in units of the stride.""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.stem = nn.ModuleList([nn.Conv2d(3, 28, 3, 2, 1), nn.Conv2d(28, 28, 3, 1, 1), nn.Conv2d(28, 56, 3, 1, 1)]) |
| |
| for i, (cin, cout, n, stride) in enumerate([(56, 56, 3, 1), (56, 88, 4, 2), (88, 88, 2, 2), (88, 224, 3, 2)]): |
| setattr(self, f"layer{i + 1}", nn.Sequential(_Block(cin, cout, stride), *[_Block(cout, cout, 1) for _ in range(n - 1)])) |
| self.lateral = nn.ModuleList([nn.Conv2d(88, 56, 1), nn.Conv2d(88, 56, 1), nn.Conv2d(224, 56, 1)]) |
| self.fpn = nn.ModuleList([nn.Conv2d(56, 56, 3, 1, 1) for _ in range(3)]) |
| self.down = nn.ModuleList([nn.Conv2d(56, 56, 3, 2, 1) for _ in range(2)]) |
| self.pafpn = nn.ModuleList([nn.Conv2d(56, 56, 3, 1, 1) for _ in range(2)]) |
| self.heads = nn.ModuleList([_Head() for _ in range(3)]) |
|
|
| def forward(self, x): |
| for c in self.stem: |
| x = F.relu(c(x)) |
| x = F.max_pool2d(x, 2, 2) |
| c2 = self.layer2(self.layer1(x)) |
| c3 = self.layer3(c2) |
| c4 = self.layer4(c3) |
| p4 = self.lateral[2](c4) |
| p3 = self.lateral[1](c3) + F.interpolate(p4, size=c3.shape[2:], mode="nearest") |
| p2 = self.lateral[0](c2) + F.interpolate(p3, size=c2.shape[2:], mode="nearest") |
| o2, o3, o4 = self.fpn[0](p2), self.fpn[1](p3), self.fpn[2](p4) |
| n3 = o3 + self.down[0](o2) |
| n4 = o4 + self.down[1](n3) |
| return [h(f) for h, f in zip(self.heads, (o2, self.pafpn[0](n3), self.pafpn[1](n4)))] |
|
|
|
|
| |
|
|
| class _IBlock(nn.Module): |
| """IResNet block: BN, 3x3 conv, PReLU, 3x3 conv (the stride), plus the shortcut.""" |
|
|
| def __init__(self, cin: int, cout: int, stride: int): |
| super().__init__() |
| self.bn = nn.BatchNorm2d(cin, eps=1e-5) |
| self.conv1 = nn.Conv2d(cin, cout, 3, 1, 1) |
| self.prelu = nn.PReLU(cout) |
| self.conv2 = nn.Conv2d(cout, cout, 3, stride, 1) |
| self.down = nn.Conv2d(cin, cout, 1, stride) if stride != 1 else None |
|
|
| def forward(self, x): |
| out = self.conv2(self.prelu(self.conv1(self.bn(x)))) |
| return out + (self.down(x) if self.down is not None else x) |
|
|
|
|
| class IResNet(nn.Module): |
| """(N, 3, 112, 112) RGB in (x - 127.5) / 127.5 -> (N, 512) identity embedding (not unit length).""" |
|
|
| LAYERS = {"r50": (3, 4, 14, 3), "r100": (3, 13, 30, 3)} |
|
|
| def __init__(self, layers: tuple[int, ...]): |
| super().__init__() |
| self.conv = nn.Conv2d(3, 64, 3, 1, 1) |
| self.prelu = nn.PReLU(64) |
| blocks, cin = [], 64 |
| for cout, n in zip((64, 128, 256, 512), layers): |
| blocks += [_IBlock(cin, cout, 2)] + [_IBlock(cout, cout, 1) for _ in range(n - 1)] |
| cin = cout |
| self.blocks = nn.Sequential(*blocks) |
| self.bn = nn.BatchNorm2d(512, eps=1e-5) |
| self.fc = nn.Linear(512 * 7 * 7, 512) |
| self.features = nn.BatchNorm1d(512, eps=1e-5) |
|
|
| def forward(self, x): |
| x = self.bn(self.blocks(self.prelu(self.conv(x)))) |
| return self.features(self.fc(x.flatten(1))) |
|
|
