| """Backbone (IR-50) split into F_c / F_s, and the U-Net feature-inversion network. |
| |
| The IR-50 definition follows face.evoLVe / AdaFace exactly so the pretrained |
| AdaFace MS1Mv2 checkpoint loads without surgery. The client encoder F_c is the |
| stem plus the first residual stage, whose output is 64x56x56 -- the intermediate |
| feature used throughout DecoyFace (paper Sec. V-B). F_s is the remainder. |
| """ |
| from collections import namedtuple |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| |
| class Flatten(nn.Module): |
| def forward(self, x): |
| return x.view(x.size(0), -1) |
|
|
|
|
| class BasicBlockIR(nn.Module): |
| def __init__(self, in_channel, depth, stride): |
| super().__init__() |
| if in_channel == depth: |
| self.shortcut_layer = nn.MaxPool2d(1, stride) |
| else: |
| self.shortcut_layer = nn.Sequential( |
| nn.Conv2d(in_channel, depth, (1, 1), stride, bias=False), |
| nn.BatchNorm2d(depth)) |
| self.res_layer = nn.Sequential( |
| nn.BatchNorm2d(in_channel), |
| nn.Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), |
| nn.BatchNorm2d(depth), |
| nn.PReLU(depth), |
| nn.Conv2d(depth, depth, (3, 3), stride, 1, bias=False), |
| nn.BatchNorm2d(depth)) |
|
|
| def forward(self, x): |
| return self.res_layer(x) + self.shortcut_layer(x) |
|
|
|
|
| Bottleneck = namedtuple('Block', ['in_channel', 'depth', 'stride']) |
|
|
|
|
| def get_block(in_channel, depth, num_units, stride=2): |
| return [Bottleneck(in_channel, depth, stride)] + \ |
| [Bottleneck(depth, depth, 1) for _ in range(num_units - 1)] |
|
|
|
|
| def get_blocks_ir50(): |
| return [ |
| get_block(64, 64, 3), |
| get_block(64, 128, 4), |
| get_block(128, 256, 14), |
| get_block(256, 512, 3), |
| ] |
|
|
|
|
| class IR50(nn.Module): |
| """Full IR-50 recognizer (input 112x112 -> 512-d embedding).""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.input_layer = nn.Sequential( |
| nn.Conv2d(3, 64, (3, 3), 1, 1, bias=False), |
| nn.BatchNorm2d(64), nn.PReLU(64)) |
| modules = [] |
| for block in get_blocks_ir50(): |
| for b in block: |
| modules.append(BasicBlockIR(b.in_channel, b.depth, b.stride)) |
| self.body = nn.Sequential(*modules) |
| self.output_layer = nn.Sequential( |
| nn.BatchNorm2d(512), nn.Dropout(0.4), Flatten(), |
| nn.Linear(512 * 7 * 7, 512), nn.BatchNorm1d(512, affine=False)) |
|
|
| def forward(self, x): |
| x = self.input_layer(x) |
| x = self.body(x) |
| return self.output_layer(x) |
|
|
|
|
| |
| |
| SPLIT_INDEX = 3 |
|
|
|
|
| def adaface_preprocess(x_rgb01): |
| """RGB image in [0,1] -> AdaFace input (BGR, normalized to [-1,1]).""" |
| x_bgr = x_rgb01.flip(1) |
| return (x_bgr - 0.5) / 0.5 |
|
|
|
|
| class ClientEncoder(nn.Module): |
| """F_c: AdaFace stem + stage 1. Frozen. Consumes RGB [0,1], emits 64x56x56.""" |
|
|
| def __init__(self, ir50: IR50): |
| super().__init__() |
| self.input_layer = ir50.input_layer |
| self.stage1 = ir50.body[:SPLIT_INDEX] |
|
|
| def forward(self, x_rgb01): |
| x = adaface_preprocess(x_rgb01) |
| x = self.input_layer(x) |
| return self.stage1(x) |
|
|
|
|
| class ServerBackbone(nn.Module): |
| """F_s: AdaFace stage 2-4 + output layer. 64x56x56 -> 512-d embedding.""" |
|
|
| def __init__(self, ir50: IR50): |
| super().__init__() |
| self.rest = ir50.body[SPLIT_INDEX:] |
| self.output_layer = ir50.output_layer |
|
|
| def forward(self, feat): |
| return self.output_layer(self.rest(feat)) |
|
|
|
|
| def load_ir50(ckpt_path, device): |
| net = IR50() |
| sd = torch.load(ckpt_path, map_location="cpu") |
| if isinstance(sd, dict) and "state_dict" in sd: |
| sd = sd["state_dict"] |
| missing, unexpected = net.load_state_dict(sd, strict=True) |
| net.eval().to(device) |
| return net |
|
|
|
|
| |
| class ConvBlock(nn.Module): |
| def __init__(self, cin, cout, norm="batch"): |
| super().__init__() |
| if norm == "batch": |
| norm1 = nn.BatchNorm2d(cout) |
| norm2 = nn.BatchNorm2d(cout) |
| elif norm == "instance": |
| norm1 = nn.InstanceNorm2d(cout) |
| norm2 = nn.InstanceNorm2d(cout) |
| else: |
| raise ValueError(norm) |
| self.net = nn.Sequential( |
| nn.Conv2d(cin, cout, 3, 1, 1, bias=False), |
| norm1, nn.LeakyReLU(0.2, inplace=True), |
| nn.Conv2d(cout, cout, 3, 1, 1, bias=False), |
| norm2, nn.LeakyReLU(0.2, inplace=True)) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class UNetInverter(nn.Module): |
| """Feature-inversion U-Net: 64x56x56 feature -> 3x112x112 RGB image in [0,1]. |
| |
| Encoder-decoder with skip connections over the feature map, plus a final |
| upsampling stage to reach the 112x112 image resolution. |
| """ |
|
|
| def __init__(self, in_ch=64, norm="batch"): |
| super().__init__() |
| self.enc0 = ConvBlock(in_ch, 64, norm) |
| self.down1 = nn.MaxPool2d(2) |
| self.enc1 = ConvBlock(64, 128, norm) |
| self.down2 = nn.MaxPool2d(2) |
| self.enc2 = ConvBlock(128, 256, norm) |
| self.bott = ConvBlock(256, 256, norm) |
|
|
| self.up2 = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) |
| self.dec2 = ConvBlock(256 + 128, 128, norm) |
| self.up1 = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) |
| self.dec1 = ConvBlock(128 + 64, 64, norm) |
| self.up0 = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) |
| self.dec0 = ConvBlock(64, 32, norm) |
| self.head = nn.Conv2d(32, 3, 3, 1, 1) |
|
|
| def forward(self, feat): |
| e0 = self.enc0(feat) |
| e1 = self.enc1(self.down1(e0)) |
| e2 = self.enc2(self.down2(e1)) |
| b = self.bott(e2) |
| d2 = self.dec2(torch.cat([self.up2(b), e1], 1)) |
| d1 = self.dec1(torch.cat([self.up1(d2), e0], 1)) |
| d0 = self.dec0(self.up0(d1)) |
| return torch.sigmoid(self.head(d0)) |
|
|
