| """DecoyFace client-side protection (DGCM) and server-side recovery (ACM). |
| |
| All projections act on the channel dimension of a C x H x W feature at every |
| spatial location. C = 64 (F_c output), r = 32, G = 16 groups of size 4. |
| |
| The RSSD statistics (P_rec, P_N, mu_rec), the fixed channel-shuffle permutation |
| and the auxiliary complementary bank D_N are built once (see build_rssd.py / |
| build_bank.py) and loaded into DecoyProtector. |
| """ |
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def proj(P, z): |
| """Apply channel-space projection P (C,C) to feature z (B,C,H,W).""" |
| b, c, h, w = z.shape |
| out = P @ z.reshape(b, c, h * w) |
| return out.reshape(b, c, h, w) |
|
|
|
|
| def group_shuffle(z, perm, G): |
| """CSG: permute channels by `perm`, then split into G contiguous groups. |
| z (B,C,H,W) -> (B,G,c_g,H,W).""" |
| b, c, h, w = z.shape |
| z = z[:, perm] |
| return z.reshape(b, G, c // G, h, w) |
|
|
|
|
| def group_unshuffle(zg, inv_perm): |
| """Inverse CSG: (B,G,c_g,H,W) -> (B,C,H,W) in the original channel order.""" |
| b, G, cg, h, w = zg.shape |
| z = zg.reshape(b, G * cg, h, w) |
| return z[:, inv_perm] |
|
|
|
|
| def group_instance_norm(zg, eps=1e-5): |
| """InstanceNorm applied to every group's channels over the spatial dims.""" |
| b, G, cg, h, w = zg.shape |
| x = zg.reshape(b * G, cg, h, w) |
| x = F.instance_norm(x, eps=eps) |
| return x.reshape(b, G, cg, h, w) |
|
|
|
|
| class DecoyProtector: |
| def __init__(self, P_rec, P_N, mu_rec, aux_bank, perm, cfg, device): |
| self.device = device |
| self.P_rec = P_rec.to(device) |
| self.P_N = P_N.to(device) |
| self.mu_rec = mu_rec.to(device) |
| self.perm = perm.to(device) |
| self.inv_perm = torch.argsort(perm).to(device) |
| self.cfg = cfg |
| self.G = cfg["G"] |
| self.K = cfg["K"] |
| self.M = aux_bank.shape[0] |
| |
| |
| ag = group_shuffle(aux_bank.to(device), self.perm, self.G) |
| self.aux_grouped_norm = group_instance_norm(ag) |
|
|
| def _sample_routing(self, B): |
| """Per-sample mask m (B,G) and decoy source indices src (B,G).""" |
| G, K, M, cfg = self.G, self.K, self.M, self.cfg |
| mask = torch.zeros(B, G, dtype=torch.bool) |
| src = torch.zeros(B, G, dtype=torch.long) |
| for b in range(B): |
| keep = torch.randperm(G)[:K] |
| mask[b, keep] = True |
| g0 = torch.nonzero(~mask[b], as_tuple=False).flatten() |
| sc = int(torch.randint(cfg["sc_lo"], cfg["sc_hi"] + 1, (1,))) |
| sc = min(sc, g0.numel()) |
| L = int(torch.randint(cfg["L_lo"], cfg["L_hi"] + 1, (1,))) |
| L = min(L, sc) if sc > 0 else L |
| perm_g0 = g0[torch.randperm(g0.numel())] |
| coherent = perm_g0[:sc] |
| independent = perm_g0[sc:] |
| shared_pool = torch.randperm(M)[:L] |
| for g in coherent.tolist(): |
| src[b, g] = shared_pool[torch.randint(L, (1,))] |
| for g in independent.tolist(): |
| src[b, g] = torch.randint(M, (1,)) |
| return mask.to(self.device), src.to(self.device) |
|
|
| def protect(self, z_A, z_B): |
| """Client-side DGCM. z_A, z_B: (B,C,H,W). Returns transmitted z_t.""" |
| B = z_A.shape[0] |
| r_B = proj(self.P_rec, z_B) |
| n_A = proj(self.P_N, z_A) |
|
|
| uA = group_instance_norm(group_shuffle(n_A, self.perm, self.G)) |
| mask, src = self._sample_routing(B) |
| ar = torch.arange(self.G, device=self.device) |
| decoy = self.aux_grouped_norm[src, ar] |
| m = mask[:, :, None, None, None] |
| mixed = torch.where(m, uA, decoy) |
| n_tilde = group_unshuffle(mixed, self.inv_perm) |
|
|
| z_tilde = n_tilde + r_B |
| k = torch.randint(-self.cfg["t"], self.cfg["t"] + 1, (B, 1, 1, 1), |
| device=self.device) |
| z_t = (2.0 ** k.float()) * z_tilde |
| return z_t |
|
|
| def server_recover(self, z_t): |
| """Server-side ACM. z_t: (B,C,H,W). Returns canonicalized z_can.""" |
| z_s = proj(self.P_N, z_t) |
| zg = group_shuffle(z_s, self.perm, self.G) |
| b, G, cg, h, w = zg.shape |
| flat = zg.reshape(b, G, cg * h * w).abs() |
| q = torch.quantile(flat, self.cfg["q"], dim=2) |
| scale = self.cfg["tau_s"] / (q + self.cfg["eps"]) |
| zg = zg * scale[:, :, None, None, None] |
| n_cal = group_unshuffle(zg, self.inv_perm) |
| return self.mu_rec[None] + n_cal |
|
|
