"""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 # ----------------------------- IR-50 backbone ------------------------------ # 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), # stage 1 -> 64 x 56 x 56 get_block(64, 128, 4), # stage 2 -> 128 x 28 x 28 get_block(128, 256, 14), # stage 3 -> 256 x 14 x 14 get_block(256, 512, 3), # stage 4 -> 512 x 7 x 7 ] 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) # The first residual stage (stage 1) spans body[0:3]; splitting there yields the # 64x56x56 intermediate feature the paper transmits. 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) # RGB -> BGR on the channel dim 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 # ------------------------------- U-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) # 56 self.down1 = nn.MaxPool2d(2) # 28 self.enc1 = ConvBlock(64, 128, norm) self.down2 = nn.MaxPool2d(2) # 14 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) # 28 self.up1 = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) self.dec1 = ConvBlock(128 + 64, 64, norm) # 56 self.up0 = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) self.dec0 = ConvBlock(64, 32, norm) # 112 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))